Add Operator for Implicit Models

ci/ignore_codespell_words.txt

@@ -5,3 +5,4 @@ fo
 ot
 lik
 usera
+als

merlin/systems/dag/ops/implicit.py

@@ -0,0 +1,137 @@
+#
+# Copyright (c) 2022, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import importlib
+import json
+import pathlib
+
+from merlin.dag import ColumnSelector  # noqa
+from merlin.schema import ColumnSchema, Schema
+from merlin.systems.dag.ops.operator import InferenceDataFrame, PipelineableInferenceOperator
+
+try:
+    import implicit
+    from packaging.version import Version
+
+    if Version(implicit.__version__) < Version("0.6.0"):
+        raise RuntimeError(
+            "Implicit version 0.6.0 or higher required. (for model save/load methods)."
+        )
+except ImportError:
+    implicit = None
+
+
+class PredictImplicit(PipelineableInferenceOperator):
+    """Operator for running inference on Implicit models.."""
+
+    def __init__(self, model, num_to_recommend: int = 10, **kwargs):
+        """Instantiate an Implicit prediction operator.
+
+        Parameters
+        ----------
+        model : An Implicit Model instance
+        num_to_recommend : int
+           the number of items to return
+        """
+        self.model = model
+        self.num_to_recommend = num_to_recommend
+        super().__init__(**kwargs)
+
+    def compute_output_schema(
+        self,
+        input_schema: Schema,
+        col_selector: ColumnSelector,
+        prev_output_schema: Schema = None,
+    ) -> Schema:
+        """Return the output schema representing the columns this operator returns."""
+        return Schema([ColumnSchema("ids", dtype="int64"), ColumnSchema("scores", dtype="float64")])
+
+    def compute_input_schema(
+        self,
+        root_schema: Schema,
+        parents_schema: Schema,
+        deps_schema: Schema,
+        selector: ColumnSelector,
+    ) -> Schema:
+        """Return the input schema representing the input columns this operator expects to use."""
+        return Schema([ColumnSchema("user_id", dtype="int64")])
+
+    def export(self, path, input_schema, output_schema, params=None, node_id=None, version=1):
+        """Export the class and related files to the path specified."""
+        node_name = f"{node_id}_{self.export_name}" if node_id is not None else self.export_name
+        version_path = pathlib.Path(path) / node_name / str(version)
+        version_path.mkdir(parents=True, exist_ok=True)
+        model_path = version_path / "model.npz"
+        self.model.save(str(model_path))
+        params = params or {}
+        params["model_module_name"] = self.model.__module__
+        params["model_class_name"] = self.model.__class__.__name__
+        params["num_to_recommend"] = self.num_to_recommend
+        return super().export(
+            path,
+            input_schema,
+            output_schema,
+            params=params,
+            node_id=node_id,
+            version=version,
+        )
+
+    @classmethod
+    def from_config(cls, config: dict, **kwargs) -> "PredictImplicit":
+        """Instantiate the class from a dictionary representation.
+
+        Expected config structure:
+        {
+            "input_dict": str  # JSON dict with input names and schemas
+            "params": str  # JSON dict with params saved at export
+        }
+
+        """
+        params = json.loads(config["params"])
+
+        model_repository = kwargs["model_repository"]
+        model_name = kwargs["model_name"]
+        model_version = kwargs["model_version"]
+
+        # load implicit model
+        model_module_name = params["model_module_name"]
+        model_class_name = params["model_class_name"]
+        model_module = importlib.import_module(model_module_name)
+        model_cls = getattr(model_module, model_class_name)
+        model_file = pathlib.Path(model_repository) / model_name / str(model_version) / "model.npz"
+        model = model_cls.load(str(model_file))
+
+        num_to_recommend = params["num_to_recommend"]
+
+        return cls(model, num_to_recommend=num_to_recommend)
+
+    def transform(self, df: InferenceDataFrame) -> InferenceDataFrame:
+        """Transform the dataframe by applying this operator to the set of input columns.
+
+        Parameters
+        -----------
+        df: InferenceDataFrame
+            A pandas or cudf dataframe that this operator will work on
+
+        Returns
+        -------
+        InferenceDataFrame
+            Returns a transformed dataframe for this operator"""
+        user_id = df["user_id"].ravel()
+        user_items = None
+        ids, scores = self.model.recommend(
+            user_id, user_items, N=self.num_to_recommend, filter_already_liked_items=False
+        )
+        return InferenceDataFrame({"ids": ids, "scores": scores})

merlin/systems/triton/oprunner_model.py

@@ -129,14 +129,13 @@ def execute(self, requests):
 
 def _parse_model_repository(model_repository: str) -> str:
     """
-    Extract the model repository path from the value passed to the TritonPythonModel
-    initialize method.
+    Extract the model repository path from the model_repository value
+    passed to the TritonPythonModel initialize method.
     """
-    model_repository_path = pathlib.Path(model_repository).parent
-
     # Handle bug in Tritonserver 22.06
     # model_repository argument became path to model.py
-    if str(model_repository).endswith(".py"):
-        model_repository_path = model_repository_path.parent
-
-    return str(model_repository_path)
+    # instead of path to model directory within the model repository
+    if model_repository.endswith(".py"):
+        return str(pathlib.Path(model_repository).parent.parent.parent)
+    else:
+        return str(pathlib.Path(model_repository).parent)

requirements-test.txt

@@ -15,6 +15,7 @@ testbook==0.4.2
 # packages necessary to run tests and push PRs
 feast==0.19.4
 xgboost==1.6.1
+implicit==0.6.0
 
 # TODO: do we need more of these?
 # https://github.com/NVIDIA-Merlin/Merlin/blob/a1cc48fe23c4dfc627423168436f26ef7e028204/ci/dockerfile.ci#L13-L18

tests/unit/systems/implicit/__init__.py

@@ -0,0 +1,19 @@
+#
+# Copyright (c) 2022, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import pytest
+
+pytest.importorskip("implicit")

tests/unit/systems/implicit/test_implicit.py

@@ -0,0 +1,127 @@
+#
+# Copyright (c) 2022, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import json
+from distutils.spawn import find_executable
+
+import implicit
+import numpy as np
+import pytest
+from scipy.sparse import csr_matrix
+
+from merlin.schema import ColumnSchema, Schema
+from merlin.systems.dag.ensemble import Ensemble
+from merlin.systems.dag.ops.implicit import PredictImplicit
+from merlin.systems.triton.utils import run_triton_server
+
+TRITON_SERVER_PATH = find_executable("tritonserver")
+
+triton = pytest.importorskip("merlin.systems.triton")
+grpcclient = pytest.importorskip("tritonclient.grpc")
+
+
+@pytest.mark.parametrize(
+    "model_cls",
+    [
+        implicit.bpr.BayesianPersonalizedRanking,
+        implicit.als.AlternatingLeastSquares,
+        implicit.lmf.LogisticMatrixFactorization,
+    ],
+)
+def test_reload_from_config(model_cls, tmpdir):
+    model = model_cls()
+    n = 10
+    user_items = csr_matrix(np.random.choice([0, 1], size=n * n).reshape(n, n))
+    model.fit(user_items)
+
+    op = PredictImplicit(model)
+
+    config = op.export(tmpdir, Schema(), Schema())
+
+    node_config = json.loads(config.parameters[config.name].string_value)
+
+    cls = PredictImplicit.from_config(
+        node_config,
+        model_repository=tmpdir,
+        model_name=config.name,
+        model_version=1,
+    )
+    reloaded_model = cls.model
+
+    num_to_recommend = np.random.randint(1, n)
+    user_items = None
+    ids, scores = model.recommend(
+        1, user_items, N=num_to_recommend, filter_already_liked_items=False
+    )
+
+    reloaded_ids, reloaded_scores = reloaded_model.recommend(
+        1, user_items, N=num_to_recommend, filter_already_liked_items=False
+    )
+
+    np.testing.assert_array_equal(ids, reloaded_ids)
+    np.testing.assert_array_equal(scores, reloaded_scores)
+
+
+@pytest.mark.skipif(not TRITON_SERVER_PATH, reason="triton server not found")
+@pytest.mark.parametrize(
+    "model_cls",
+    [
+        implicit.bpr.BayesianPersonalizedRanking,
+        implicit.als.AlternatingLeastSquares,
+        implicit.lmf.LogisticMatrixFactorization,
+    ],
+)
+def test_ensemble(model_cls, tmpdir):
+    model = model_cls()
+    n = 100
+    user_items = csr_matrix(np.random.choice([0, 1], size=n * n, p=[0.9, 0.1]).reshape(n, n))
+    model.fit(user_items)
+
+    num_to_recommend = np.random.randint(1, n)
+
+    user_items = None
+    ids, scores = model.recommend(
+        [0, 1], user_items, N=num_to_recommend, filter_already_liked_items=False
+    )
+
+    implicit_op = PredictImplicit(model, num_to_recommend=num_to_recommend)
+
+    input_schema = Schema([ColumnSchema("user_id", dtype="int64")])
+
+    triton_chain = input_schema.column_names >> implicit_op
+
+    triton_ens = Ensemble(triton_chain, input_schema)
+    triton_ens.export(tmpdir)
+
+    model_name = triton_ens.name
+    input_user_id = np.array([[0], [1]], dtype=np.int64)
+    inputs = [
+        grpcclient.InferInput(
+            "user_id", input_user_id.shape, triton.np_to_triton_dtype(input_user_id.dtype)
+        ),
+    ]
+    inputs[0].set_data_from_numpy(input_user_id)
+    outputs = [grpcclient.InferRequestedOutput("scores"), grpcclient.InferRequestedOutput("ids")]
+
+    response = None
+
+    with run_triton_server(tmpdir) as client:
+        response = client.infer(model_name, inputs, outputs=outputs)
+
+    response_ids = response.as_numpy("ids")
+    response_scores = response.as_numpy("scores")
+
+    np.testing.assert_array_equal(ids, response_ids)
+    np.testing.assert_array_equal(scores, response_scores)