From 972b879f666a48a17c1481fd22e4ec60740c21ce Mon Sep 17 00:00:00 2001
From: Lutz Roeder <lutzroeder@users.noreply.github.com>
Date: Thu, 19 Dec 2019 23:19:14 -1000
Subject: [PATCH] CoreML itemSimilarityRecommender support (#193)

---
 src/coreml-metadata.json | 15 +++++++++++++++
 src/coreml.js            | 11 +++++++++++
 test/models.json         |  7 +++++++
 3 files changed, 33 insertions(+)

diff --git a/src/coreml-metadata.json b/src/coreml-metadata.json
index 5fbe896f38..2a197d672d 100644
--- 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
index 4820341925..4fecbb0d32 100644
--- 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
index bb3768d898..5a7094ff1d 100644
--- 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",