CoreML itemSimilarityRecommender support (#193)

lutzroeder · Dec 20, 2019 · 972b879 · 972b879
1 parent 45136c3
commit 972b879
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 0 deletions.
diff --git a/src/coreml-metadata.json b/src/coreml-metadata.json
@@ -491,5 +491,20 @@
       "category": "Normalization",
       "description": "A layer that performs mean variance normalization, along axis = -3."
     }
+  },
+  {
+    "name": "itemSimilarityRecommender",
+    "schema": {
+      "inputs": [
+        { "name": "item" },
+        { "name": "numRecommendations" },
+        { "name": "itemRestriction" },
+        { "name": "itemExclusion" }
+      ],
+      "outputs": [
+        { "name": "recommendedItemList" },
+        { "name": "recommendedItemScore" }
+      ]
+    }
   }
 ]
diff --git a/src/coreml.js b/src/coreml.js
@@ -405,6 +405,17 @@ coreml.Graph = class {
             this._updateClassifierOutput(group, model.kNearestNeighborsClassifier);
             return 'kNearestNeighborsClassifier';
         }
+        else if (model.itemSimilarityRecommender) {
+            const itemSimilarityRecommenderParams = {
+                itemStringIds: model.itemSimilarityRecommender.itemStringIds.vector,
+                itemItemSimilarities: model.itemSimilarityRecommender.itemItemSimilarities
+            }
+            this._createNode(scope, group, 'itemSimilarityRecommender', null,
+                itemSimilarityRecommenderParams,
+                model.description.input.map((feature) => feature.name),
+                model.description.output.map((feature) => feature.name));
+            return 'itemSimilarityRecommender'
+        }
         else if (model.customModel) {
             this._createNode(scope, group, 'customModel', null,
                 { className: model.customModel.className, parameters: model.customModel.parameters },

diff --git a/test/models.json b/test/models.json
@@ -1610,6 +1610,13 @@
     "format": "Core ML v1",
     "link":   "https://github.com/tf-coreml/tf-coreml/issues/252"
   },
+  {
+    "type":   "coreml",
+    "target": "MovieRecommender.mlmodel",
+    "source": "https://github.com/anupamchugh/iowncode/blob/master/CoreMLRecommender/CoreMLRecommender/MovieRecommender.mlmodel?raw=true",
+    "format": "Core ML v4",
+    "link":   "https://github.com/anupamchugh/iowncode/tree/master/CoreMLRecommender"
+  },
   {
     "type":   "coreml",
     "target": "NamesDT.mlmodel",