Update pickle.js

lutzroeder · Aug 6, 2023 · 13844a0 · 13844a0
1 parent f068867
commit 13844a0
Showing 1 changed file with 61 additions and 65 deletions.
diff --git a/source/pickle.js b/source/pickle.js
@@ -54,117 +54,113 @@ pickle.ModelFactory = class {
 pickle.Model = class {
 
     constructor(value, format) {
-        this._format = format;
-        this._graphs = [ new pickle.Graph(value) ];
-    }
-
-    get format() {
-        return this._format;
-    }
-
-    get graphs() {
-        return this._graphs;
+        this.format = format;
+        this.graphs = [ new pickle.Graph(value) ];
     }
 };
 
 pickle.Graph = class {
 
     constructor(obj) {
-        this._inputs = [];
-        this._outputs = [];
-        this._nodes = [];
-
+        this.inputs = [];
+        this.outputs = [];
+        this.nodes = [];
         if (Array.isArray(obj) && (obj.every((item) => item.__class__) || (obj.every((item) => Array.isArray(item))))) {
             for (const item of obj) {
-                this._nodes.push(new pickle.Node(item));
+                this.nodes.push(new pickle.Node(item));
             }
         } else if (obj && obj instanceof Map && !Array.from(obj.values()).some((value) => typeof value === 'string' || typeof value === 'number')) {
             for (const entry of obj) {
-                this._nodes.push(new pickle.Node(entry[1], entry[0]));
+                this.nodes.push(new pickle.Node(entry[1], entry[0]));
             }
         } else if (obj && obj.__class__) {
-            this._nodes.push(new pickle.Node(obj));
+            this.nodes.push(new pickle.Node(obj));
         } else if (obj && Object(obj) === obj) {
-            this._nodes.push(new pickle.Node(obj));
+            this.nodes.push(new pickle.Node(obj));
         }
     }
-
-    get inputs() {
-        return this._inputs;
-    }
-
-    get outputs() {
-        return this._outputs;
-    }
-
-    get nodes() {
-        return this._nodes;
-    }
 };
 
 pickle.Node = class {
 
     constructor(obj, name) {
-        this._name = name || '';
-        this._inputs = [];
-        this._outputs = [];
-        this._attributes = [];
+        this.name = name || '';
+        this.inputs = [];
+        this.outputs = [];
+        this.attributes = [];
+        const isArray = (obj) => {
+            return obj && obj.__class__ && obj.__class__.__module__ === 'numpy' && obj.__class__.__name__ === 'ndarray';
+        };
         if (Array.isArray(obj)) {
-            this._type = { name: 'List' };
-            this._attributes.push(new pickle.Attribute('value', obj));
+            this.type = { name: 'List' };
+            const attribute = new pickle.Attribute('value', obj);
+            this.attributes.push(attribute);
         } else {
             const type = obj.__class__ ? obj.__class__.__module__ + '.' + obj.__class__.__name__ : 'Object';
-            this._type = { name: type };
+            this.type = { name: type };
             const entries = obj instanceof Map ? Array.from(obj.entries()) : Object.entries(obj);
             for (const entry of entries) {
                 const name = entry[0];
                 const value = entry[1];
-                this._attributes.push(new pickle.Attribute(name, value));
+                if (value && isArray(value)) {
+                    const tensor = new pickle.Tensor(value);
+                    const attribute = new pickle.Attribute(name, 'tensor', tensor);
+                    this.attributes.push(attribute);
+                } else if (Array.isArray(value) && value.length > 0 && value.every((obj) => isArray(obj))) {
+                    const values = value.map((value) => new pickle.Tensor(value));
+                    const attribute = new pickle.Attribute(name, 'tensor[]', values);
+                    this.attributes.push(attribute);
+                } else {
+                    const attribute = new pickle.Attribute(name, null, value);
+                    this.attributes.push(attribute);
+                }
+
             }
         }
     }
+};
 
-    get type() {
-        return this._type;
-    }
-
-    get name() {
-        return this._name;
-    }
+pickle.Attribute = class {
 
-    get inputs() {
-        return this._inputs;
+    constructor(name, type, value) {
+        this.name = name;
+        this.type = type;
+        this.value = value;
+        if (!type && value && value.__class__) {
+            this.type = value.__class__.__module__ + '.' + value.__class__.__name__;
+        }
     }
+};
 
-    get outputs() {
-        return this._outputs;
-    }
+pickle.Tensor = class {
 
-    get attributes() {
-        return this._attributes;
+    constructor(array) {
+        this.type = new pickle.TensorType(array.dtype.__name__, new pickle.TensorShape(array.shape));
+        this.layout = this.type.dataType == 'string' || this.type.dataType == 'object' ? '|' : array.dtype.byteorder;
+        this.values = this.type.dataType == 'string' || this.type.dataType == 'object' ? array.tolist() : array.tobytes();
     }
 };
 
-pickle.Attribute = class {
+pickle.TensorType = class {
 
-    constructor(name, value) {
-        this._name = name;
-        this._value = value;
-        if (value && value.__class__) {
-            this._type = value.__class__.__module__ + '.' + value.__class__.__name__;
-        }
+    constructor(dataType, shape) {
+        this.dataType = dataType;
+        this.shape = shape;
     }
 
-    get name() {
-        return this._name;
+    toString() {
+        return this.dataType + this.shape.toString();
     }
+};
+
+pickle.TensorShape = class {
 
-    get value() {
-        return this._value;
+    constructor(dimensions) {
+        this.dimensions = dimensions;
     }
 
-    get type() {
-        return this._type;
+    toString() {
+        return this.dimensions ? ('[' + this.dimensions.map((dimension) => dimension.toString()).join(',') + ']') : '';
     }
 };