Fixing merge-conflicts

Author: marcromeyn

merlin/models/torch/__init__.py

@@ -16,6 +16,7 @@
 
 from merlin.models.torch.batch import Batch, Sequence
 from merlin.models.torch.block import Block, ParallelBlock
+from merlin.models.torch.blocks.mlp import MLPBlock
 from merlin.models.torch.inputs.embedding import EmbeddingTable, EmbeddingTables
 from merlin.models.torch.inputs.tabular import TabularInputBlock
 from merlin.models.torch.outputs.base import ModelOutput
@@ -31,6 +32,7 @@
     "EmbeddingTable",
     "EmbeddingTables",
     "ParallelBlock",
+    "MLPBlock",
     "ModelOutput",
     "Sequence",
     "RegressionOutput",

merlin/models/torch/router.py

@@ -182,7 +182,9 @@ def select(self, selection: Selection) -> "RouterBlock":
 
         selected_branches = {}
         for key, val in self.branches.items():
-            selected_branches[key] = select_container(val, selection)
+            selected = select_container(val, selection)
+            if selected:
+                selected_branches[key] = selected
 
         selectable = self.__class__(self.selectable.select(selection))
         for key, val in selected_branches.items():
@@ -222,6 +224,7 @@ class SelectKeys(nn.Module, Selectable):
 
     def __init__(self, schema: Optional[Schema] = None):
         super().__init__()
+        self.col_names: List[str] = []
         if schema:
             self.setup_schema(schema)
 
@@ -231,7 +234,7 @@ def setup_schema(self, schema: Schema):
 
         super().setup_schema(schema)
 
-        self.col_names: List[str] = schema.column_names
+        self.col_names = schema.column_names
 
     def select(self, selection: Selection) -> "SelectKeys":
         """Select a subset of the schema based on the provided selection.
@@ -277,14 +280,28 @@ def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
 
         return outputs
 
+    def extra_repr(self) -> str:
+        return f"select: {', '.join(self.col_names)}"
+
+    def __bool__(self) -> bool:
+        return bool(self.col_names)
+
 
 def select_container(container: BlockContainer, selection: Selection) -> BlockContainer:
     outputs = []
 
-    for module in container:
-        if module.accepts_dict:
-            outputs.append(module.select(selection))
+    if hasattr(container.values[0], "select"):
+        first = container.values[0].select(selection)
+        if first:
+            outputs.append(first)
         else:
-            outputs.append(module)
+            return BlockContainer()
+
+    if len(container.values) > 1:
+        for module in container.values[1:]:
+            if hasattr(module, "select"):
+                outputs.append(module.select(selection))
+            else:
+                outputs.append(module)
 
     return BlockContainer(*outputs, name=container._name)

merlin/models/torch/utils/selection_utils.py

@@ -33,15 +33,7 @@ def select_schema(schema: Schema, selection: Selection) -> Schema:
     if isinstance(selection, Schema):
         selected = selection
     elif isinstance(selection, ColumnSchema):
-<<<<<<< HEAD
-<<<<<<< HEAD
         selected = Schema([schema[selection.name]])
-=======
-        selected = schema[selection.name]
->>>>>>> a2644079 (Add selection_utils)
-=======
-        selected = Schema([schema[selection.name]])
->>>>>>> 89a6f043 (Increase test-coverage)
     elif callable(selection):
         selected = selection(schema)
     elif isinstance(selection, Tags):

tests/conftest.py

@@ -20,14 +20,12 @@
 from pathlib import Path
 
 import distributed
-import numpy as np
 import pytest
 
 from merlin.core.utils import Distributed
 from merlin.datasets.synthetic import generate_data
 from merlin.io import Dataset
 from merlin.models.utils import ci_utils
-from merlin.schema import ColumnSchema, Tags
 
 REPO_ROOT = Path(__file__).parent.parent
 
@@ -86,26 +84,6 @@ def dask_client() -> distributed.Client:
     pass
 
 
-@pytest.fixture
-def item_id_col_schema() -> ColumnSchema:
-    return ColumnSchema(
-        "item_id",
-        dtype=np.int32,
-        properties={"domain": {"min": 0, "max": 10, "name": "item_id"}},
-        tags=[Tags.CATEGORICAL, Tags.ITEM_ID],
-    )
-
-
-@pytest.fixture
-def user_id_col_schema() -> ColumnSchema:
-    return ColumnSchema(
-        "user_id",
-        dtype=np.int32,
-        properties={"domain": {"min": 0, "max": 20, "name": "user_id"}},
-        tags=[Tags.CATEGORICAL, Tags.USER_ID],
-    )
-
-
 def pytest_collection_modifyitems(items):
     changed_backends = ci_utils.get_changed_backends()
     full_name_to_alias = {v: k for k, v in ci_utils.BACKEND_ALIASES.items()}

tests/unit/torch/_conftest.py

@@ -0,0 +1,24 @@
+import numpy as np
+import pytest
+
+from merlin.schema import ColumnSchema, Tags
+
+
+@pytest.fixture
+def item_id_col_schema() -> ColumnSchema:
+    return ColumnSchema(
+        "item_id",
+        dtype=np.int32,
+        properties={"domain": {"min": 0, "max": 10, "name": "item_id"}},
+        tags=[Tags.CATEGORICAL, Tags.ITEM_ID],
+    )
+
+
+@pytest.fixture
+def user_id_col_schema() -> ColumnSchema:
+    return ColumnSchema(
+        "user_id",
+        dtype=np.int32,
+        properties={"domain": {"min": 0, "max": 20, "name": "user_id"}},
+        tags=[Tags.CATEGORICAL, Tags.USER_ID],
+    )

tests/unit/torch/test_router.py

@@ -1,5 +1,3 @@
-<<<<<<< HEAD
-<<<<<<< HEAD
 from typing import Dict
 
 import pytest
@@ -99,142 +97,19 @@ def test_select(self):
         plus_one = PlusOneDict()
 
         self.router.add_route(Tags.CONTINUOUS)
-        self.router.add_route(Tags.CATEGORICAL)
+        self.router.add_route(Tags.CATEGORICAL, mm.MLPBlock([10]))
+        self.router.add_route(Tags.USER, mm.MLPBlock([10]).prepend(mm.SelectKeys(self.schema)))
+        self.router.add_route(Tags.ITEM, mm.MLPBlock([10]))
         self.router.prepend(plus_one)
 
         router = self.router.select(Tags.CATEGORICAL)
         assert router.selectable.schema == self.schema.select_by_tag(Tags.CATEGORICAL)
         assert router[0][0] == plus_one
 
-    def test_double_add(self):
-        self.router.add_route(Tags.CONTINUOUS)
-        with pytest.raises(ValueError):
-            self.router.add_route(Tags.CONTINUOUS)
-
-    def test_nested(self):
-        self.router.add_route(Tags.CONTINUOUS)
-
-        nested = self.router.nested_router()
-        nested.add_route(Tags.USER)
-        assert "user" in nested
-
-        outputs = module_utils.module_test(nested, self.batch.features)
-        assert list(outputs.keys()) == ["user_age"]
-        assert "user_age" in nested.output_schema().column_names
-<<<<<<< HEAD
-=======
-import pytest
-=======
-from typing import Dict
->>>>>>> 89a6f043 (Increase test-coverage)
-
-import pytest
-import torch
-from torch import nn
-
-import merlin.models.torch as mm
-from merlin.models.torch.batch import Batch, sample_batch
-from merlin.models.torch.utils import module_utils
-from merlin.schema import ColumnSchema, Schema, Tags
-
-
-class ToFloat(nn.Module):
-    def forward(self, x):
-        return x.float()
-
-
-class PlusOneDict(nn.Module):
-    def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
-        return {k: v + 1 for k, v in inputs.items()}
-
-
-class TestRouterBlock:
-<<<<<<< HEAD
-    ...
->>>>>>> a2644079 (Add selection_utils)
-=======
-    @pytest.fixture(autouse=True)
-    def setup_method(self, music_streaming_data):
-        self.schema = music_streaming_data.schema
-        self.router: mm.RouterBlock = mm.RouterBlock(self.schema)
-        self.batch: Batch = sample_batch(music_streaming_data, batch_size=10)
-
-    def test_add_route(self):
-        self.router.add_route(Tags.CONTINUOUS)
-
-        outputs = module_utils.module_test(self.router, self.batch.features)
-        assert set(outputs.keys()) == set(self.schema.select_by_tag(Tags.CONTINUOUS).column_names)
-        assert "continuous" in self.router
-        assert len(self.router["continuous"]) == 1
-        assert isinstance(self.router["continuous"][0], mm.SelectKeys)
-
-    def test_add_route_module(self):
-        class CustomSelect(mm.SelectKeys):
-            ...
-
-        self.router.add_route(Tags.CONTINUOUS, CustomSelect())
-
-        outputs = self.router(self.batch.features)
-        assert set(outputs.keys()) == set(self.schema.select_by_tag(Tags.CONTINUOUS).column_names)
-        assert len(self.router["continuous"]) == 2
-        assert isinstance(self.router["continuous"][0], mm.SelectKeys)
-        assert isinstance(self.router["continuous"][1], CustomSelect)
-
-    def test_module_with_setup(self):
-        class Dummy(nn.Module):
-            def setup_schema(self, schema: Schema):
-                self.schema = schema
-
-            def forward(self, x):
-                return x
-
-        dummy = Dummy()
-        self.router.add_route(Tags.CONTINUOUS, dummy)
-        assert dummy.schema == mm.select_schema(self.schema, Tags.CONTINUOUS)
-
-        dummy_2 = Dummy()
-        self.router.add_route_for_each(ColumnSchema("user_id"), dummy_2, shared=True)
-        assert dummy_2.schema == mm.select_schema(self.schema, ColumnSchema("user_id"))
-
-    def test_add_route_parallel_block(self):
-        class FakeEmbeddings(mm.ParallelBlock):
-            ...
-
-        self.router.add_route(Tags.CATEGORICAL, FakeEmbeddings())
-        assert isinstance(self.router["categorical"], FakeEmbeddings)
-
-    @pytest.mark.parametrize("shared", [True, False])
-    def test_add_route_for_each(self, shared):
-        block = mm.Block(mm.Concat(), ToFloat(), nn.LazyLinear(10)).to(self.batch.device())
-        self.router.add_route_for_each(Tags.CONTINUOUS, block, shared=shared)
-
-        dense_pos = self.router.branches["position"][1][-1]
-        dense_age = self.router.branches["user_age"][1][-1]
-        if shared:
-            assert dense_pos == dense_age
-        else:
-            assert dense_pos != dense_age
-
-        outputs = self.router(self.batch.features)
-        assert set(outputs.keys()) == set(self.schema.select_by_tag(Tags.CONTINUOUS).column_names)
-
-        for value in outputs.values():
-            assert value.shape[-1] == 10
-
-    def test_add_route_for_each_list(self):
-        self.router.add_route_for_each([ColumnSchema("user_id")], ToFloat())
-        assert isinstance(self.router.branches["user_id"][1], ToFloat)
-
-    def test_select(self):
-        plus_one = PlusOneDict()
-
-        self.router.add_route(Tags.CONTINUOUS)
-        self.router.add_route(Tags.CATEGORICAL)
-        self.router.prepend(plus_one)
-
-        router = self.router.select(Tags.CATEGORICAL)
-        assert router.selectable.schema == self.schema.select_by_tag(Tags.CATEGORICAL)
-        assert router[0][0] == plus_one
+        user = self.router.select(Tags.USER)
+        assert "item_recency" not in user.branches["continuous"][0].col_names
+        assert "item_id" not in user.branches["categorical"][0].col_names
+        assert "item" not in user.branches
 
     def test_double_add(self):
         self.router.add_route(Tags.CONTINUOUS)
@@ -250,52 +125,12 @@ def test_nested(self):
 
         outputs = module_utils.module_test(nested, self.batch.features)
         assert list(outputs.keys()) == ["user_age"]
->>>>>>> 89a6f043 (Increase test-coverage)
-=======
->>>>>>> 78386932 (Fix failined nested router test)
+        assert "user_age" in nested.output_schema().column_names
 
 
 class TestSelectKeys:
     @pytest.fixture(autouse=True)
     def setup_method(self, music_streaming_data):
-<<<<<<< HEAD
-<<<<<<< HEAD
-        self.batch: Batch = sample_batch(music_streaming_data, batch_size=10)
-        self.schema: Schema = music_streaming_data.schema
-        self.user_schema: Schema = mm.select_schema(self.schema, Tags.USER)
-
-    def test_forward(self):
-        select_user = mm.SelectKeys(self.user_schema)
-        outputs = select_user(self.batch.features)
-
-        assert select_user.schema == self.user_schema
-
-        for col in {"user_id", "country", "user_age"}:
-            assert col in outputs
-
-        assert "user_genres__values" in outputs
-        assert "user_genres__offsets" in outputs
-
-    def test_select(self):
-        select_user = mm.SelectKeys(self.user_schema)
-
-        user_id = Schema([self.user_schema["user_id"]])
-        assert select_user.select(ColumnSchema("user_id")).schema == user_id
-        assert select_user.select(Tags.USER).schema == self.user_schema
-
-    def test_setup_schema(self):
-        select_user = mm.SelectKeys()
-        select_user.setup_schema(self.user_schema["user_id"])
-        assert select_user.schema == Schema([self.user_schema["user_id"]])
-=======
-        self.data = music_streaming_data
-        self.schema = music_streaming_data.schema
-        self.select_keys = SelectKeys(music_streaming_data.schema)
-
-    def test_forward(self):
-        ...
->>>>>>> a2644079 (Add selection_utils)
-=======
         self.batch: Batch = sample_batch(music_streaming_data, batch_size=10)
         self.schema: Schema = music_streaming_data.schema
         self.user_schema: Schema = mm.select_schema(self.schema, Tags.USER)
@@ -323,4 +158,3 @@ def test_setup_schema(self):
         select_user = mm.SelectKeys()
         select_user.setup_schema(self.user_schema["user_id"])
         assert select_user.schema == Schema([self.user_schema["user_id"]])
->>>>>>> 89a6f043 (Increase test-coverage)