Added import torch action to arch editor. Fixes #447 (#449)

deepforge-dev · Jul 7, 2016 · 14c0af8 · 14c0af8
1 parent 56a0b78
commit 14c0af8
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 7 deletions.
diff --git a/src/plugins/ImportTorch/ImportTorch.js b/src/plugins/ImportTorch/ImportTorch.js
@@ -57,12 +57,17 @@ define([
 
         this.blobClient.getMetadata(srcHash)
             .then(mdata => {  // Create the new model
-                var name = mdata.name.replace('.lua', '');
-                this.tgtNode = this.core.createNode({
-                    base: this.META.Architecture,
-                    parent: this.activeNode
-                });
-                this.core.setAttribute(this.tgtNode, 'name', name);
+                // If the current node is an architecture, assume we are just extending it
+                this.importedName = mdata.name.replace('.lua', '');
+                if (this.isMetaTypeOf(this.activeNode, this.META.Architecture)) {
+                    this.tgtNode = this.activeNode;
+                } else {  // Create a new architecture
+                    this.tgtNode = this.core.createNode({
+                        base: this.META.Architecture,
+                        parent: this.activeNode
+                    });
+                    this.core.setAttribute(this.tgtNode, 'name', this.importedName);
+                }
                 return this.blobClient.getObjectAsString(srcHash);
             })
             .then(src => {  // Retrieved the source code
@@ -81,7 +86,7 @@ define([
                 return this.save('ImportTorch updated model.');
             })
             .then(() => {  // changes saved successfully
-                var name = this.core.getAttribute(this.tgtNode, 'name');
+                var name = this.importedName;
                 this.result.setSuccess(true);
                 this.createMessage(this.tgtNode,
                     `Successfully imported ${name} architecture`);

diff --git a/src/visualizers/panels/ForgeActionButton/Actions.js b/src/visualizers/panels/ForgeActionButton/Actions.js
@@ -210,6 +210,13 @@ define([
                     this.addOperation();
                 }
             }
+        ],
+        Architecture: [
+            {
+                name: 'Import Torch Architecture',
+                icon: 'swap_vert',
+                action: importTorch
+            }
         ]
     };
 });