Update SentencePiece test files (#748)

lutzroeder · Jan 15, 2024 · 67d0c76 · 67d0c76
1 parent 099c43b
commit 67d0c76
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 5 deletions.
diff --git a/source/sentencepiece.js b/source/sentencepiece.js
@@ -40,10 +40,53 @@ sentencepiece.ModelFactory = class {
 
 sentencepiece.Model = class {
 
-    constructor() {
+    constructor(model) {
         this.format = 'SentencePiece';
-        this.graphs = [];
-        throw new sentencepiece.Error("Invalid file content. File contains sentencepiece.ModelProto data.");
+        this.graphs = [ new sentencepiece.Graph(model) ];
+    }
+};
+
+sentencepiece.Graph = class {
+
+    constructor(model) {
+        this.inputs = [];
+        this.outputs = [];
+        this.nodes = [];
+        for (const [name, value] of Object.entries(model)) {
+            const node = new sentencepiece.Node(name, value);
+            this.nodes.push(node);
+        }
+    }
+};
+
+sentencepiece.Argument = class {
+
+    constructor(name, value) {
+        this.name = name;
+        this.value = value;
+    }
+};
+
+sentencepiece.Node = class {
+
+    constructor(name, obj) {
+        this.name = name;
+        this.inputs = [];
+        this.outputs = [];
+        this.attributes = [];
+        if (Array.isArray(obj)) {
+            const type = new Set(obj.map((value) => value.constructor.name));
+            this.type = { name: `${Array.from(type)[0]}[]` };
+            const attribute = new sentencepiece.Argument(name, obj);
+            this.attributes.push(attribute);
+        } else {
+            this.type = { name: obj.constructor.name };
+            for (const [name, value] of Object.entries(obj)) {
+                const data = ArrayBuffer.isView(value) ? Array.from(value) : value;
+                const attribute = new sentencepiece.Argument(name, data);
+                this.attributes.push(attribute);
+            }
+        }
     }
 };
 

diff --git a/test/models.json b/test/models.json
@@ -6004,14 +6004,21 @@
     "type":     "sentencepiece",
     "target":   "spm_char.model",
     "source":   "https://github.com/lutzroeder/netron/files/6611189/covost2_en_asr_vocab_char.zip[spm_char.model]",
-    "error":    "Invalid file content. File contains sentencepiece.ModelProto data.",
+    "format":   "SentencePiece",
+    "link":     "https://github.com/lutzroeder/netron/issues/748"
+  },
+  {
+    "type":     "sentencepiece",
+    "target":   "tokenizer.model",
+    "source":   "https://github.com/lutzroeder/netron/files/13935506/tokenizer.model.zip[tokenizer.model]",
+    "format":   "SentencePiece",
     "link":     "https://github.com/lutzroeder/netron/issues/748"
   },
   {
     "type":     "sentencepiece",
     "target":   "universal_encoder_8k_spm.model.model",
     "source":   "https://github.com/lutzroeder/netron/files/7232403/universal_encoder_8k_spm.model.zip[universal_encoder_8k_spm.model]",
-    "error":    "Invalid file content. File contains sentencepiece.ModelProto data.",
+    "format":   "SentencePiece",
     "link":     "https://github.com/lutzroeder/netron/issues/748"
   },
   {