Make Model creation explicit (NVIDIA-Merlin#432)

* Remove the mypy plugin for model creation return type determination

No longer needed if we make model creation explicit

* Make `Model` creation explicit

Currently, we're assuming that if you give us something that computes a loss, we should create a `Model`, which is a bit of black magic that has some weird edge cases. This change brings the API closer to standard Keras.

We're also looking forward to future changes that make the `Model` responsible for computing losses and metrics, which will break this somewhat magical functionality.

Co-authored-by: Marc Romeyn <[email protected]>

* Remove unnecessary `ml.ModelContext` params

* Switch dataset tests back to `connect`

* Add tasks to masking tests

* Set the block schema in all blocks inside `SequentialBlock`

Co-authored-by: Julio Perez <[email protected]>

* Fix the return type annotation on `connect` method

Co-authored-by: Marc Romeyn <[email protected]>
Co-authored-by: Julio Perez <[email protected]>

Author: 3 people

docs/source/models_overview.md

@@ -41,7 +41,7 @@ High-level API:
 import merlin.models.tf as ml
 
 block = ml.TwoTowerBlock(schema, ml.MLPBlock([512, 256]))
-model = block.connect(ml.ItemRetrievalTask())
+model = ml.Model(block, ml.ItemRetrievalTask())
 ```
 
 Low-level API:
@@ -53,7 +53,7 @@ from merlin.schema import Tags
 user_tower = ml.InputBlock(schema.select_by_tag(Tags.USER), ml.MLPBlock([512, 256]))
 item_tower = ml.InputBlock(schema.select_by_tag(Tags.ITEM), ml.MLPBlock([512, 256]))
 two_tower = ml.ParallelBlock({"user": user_tower, "item": item_tower})
-model = two_tower.connect(ml.ItemRetrievalTask())
+model = ml.Model(two_tower, ml.ItemRetrievalTask())
 ```
 
 ## Ranking
@@ -78,7 +78,7 @@ dlrm = ml.DLRMBlock(
     bottom_block=ml.MLPBlock([512, 128]),
     top_block=ml.MLPBlock([512, 128])
 )
-model = dlrm.connect(ml.BinaryClassificationTask(schema))
+model = ml.Model(dlrm, ml.BinaryClassificationTask(schema))
 ```
 
 Low-level API:
@@ -140,7 +140,7 @@ inputs = ml.InputBlock(schema)
 prediction_tasks = ml.PredictionTasks(schema)
 block = ml.MLPBlock([64])
 mmoe = ml.MMOEBlock(prediction_tasks, expert_block=ml.MLPBlock([64]), num_experts=4)
-model = inputs.connect(block, mmoe, prediction_tasks)
+model = ml.Model(inputs, block, mmoe, prediction_tasks)
 ```
 
 ### Progressive Layered Extraction
@@ -163,5 +163,5 @@ block = ml.MLPBlock([64])
 cgc = ml.CGCBlock(
     prediction_tasks, expert_block=ml.MLPBlock([64]), num_task_experts=2, num_shared_experts=2
 )
-model = inputs.connect(ml.MLPBlock([64]), cgc, prediction_tasks)
+model = ml.Model(inputs, ml.MLPBlock([64]), cgc, prediction_tasks)
 ```

examples/06-Define-your-own-architecture-with-Merlin-Models.ipynb

@@ -886,7 +886,7 @@
     }
    ],
    "source": [
-    "model = deep_dlrm_interaction.connect(binary_task)\n",
+    "model = mm.Model(deep_dlrm_interaction, binary_task)\n",
     "type(model)"
    ]
   },

merlin/models/mypy.py

@@ -385,7 +385,7 @@ def connect(
         *block: Union[tf.keras.layers.Layer, str],
         block_name: Optional[str] = None,
         context: Optional[ModelContext] = None,
-    ) -> Union["SequentialBlock", "Model", "RetrievalModel"]:
+    ) -> "SequentialBlock":
         """Connect the block to other blocks sequentially.
 
         Parameters
@@ -399,8 +399,6 @@ def connect(
 
         """
         from merlin.models.tf.blocks.core.combinators import SequentialBlock
-        from merlin.models.tf.models.base import Model, RetrievalBlock, RetrievalModel
-        from merlin.models.tf.prediction_tasks.retrieval import ItemRetrievalTask
 
         blocks = [self.parse(b) for b in block]
 
@@ -413,16 +411,6 @@ def connect(
             [self, *blocks], copy_layers=False, block_name=block_name, context=context
         )
 
-        if isinstance(blocks[-1], ModelLikeBlock):
-            if (
-                any(isinstance(b, RetrievalBlock) for b in blocks)
-                or isinstance(self, RetrievalBlock)
-                and any(isinstance(b, ItemRetrievalTask) for b in blocks)
-            ):
-                return RetrievalModel(output)
-
-            return Model(output)
-
         return output
 
     def connect_with_residual(

merlin/models/tf/blocks/core/base.py

@@ -81,6 +81,10 @@ def __init__(
         if getattr(layers[0], "has_schema", None):
             super().set_schema(layers[0].schema)
 
+            for layer in layers[1:]:
+                if hasattr(layer, "set_schema"):
+                    layer.set_schema(layers[0].schema)
+
         layers = copy.copy(layers) if copy_layers else layers
         if filter:
             if not isinstance(filter, Filter):

merlin/models/tf/blocks/core/combinators.py

@@ -91,6 +91,6 @@ def NCFModel(
     ncf = ParallelBlock({"mf": mf_branch, "mlp": mlp_branch}, aggregation="concat")
 
     prediction_tasks = parse_prediction_tasks(schema, prediction_tasks)
-    model = ncf.connect(prediction_tasks)
+    model = Model(ncf, prediction_tasks)
 
     return model

merlin/models/tf/models/benchmark.py

@@ -68,7 +68,7 @@ def DLRMModel(
         bottom_block=bottom_block,
         top_block=top_block,
     )
-    model = dlrm_body.connect(prediction_tasks)
+    model = Model(dlrm_body, prediction_tasks)
 
     return model
 
@@ -152,7 +152,7 @@ def DCNModel(
     else:
         dcn_body = input_block.connect_branch(CrossBlock(depth), deep_block, aggregation="concat")
 
-    model = dcn_body.connect(prediction_tasks)
+    model = Model(dcn_body, prediction_tasks)
 
     return model
 
@@ -230,6 +230,6 @@ def DeepFMModel(
     )
 
     prediction_tasks = parse_prediction_tasks(schema, prediction_tasks)
-    model = deep_fm.connect(prediction_tasks)
+    model = Model(deep_fm, prediction_tasks)
 
     return model

merlin/models/tf/models/ranking.py

@@ -36,7 +36,7 @@ def MatrixFactorizationModel(
     metrics: MetricOrMetrics = ItemRetrievalTask.DEFAULT_METRICS,
     samplers: Sequence[ItemSampler] = (),
     **kwargs,
-) -> Union[Model, RetrievalModel]:
+) -> RetrievalModel:
     """Builds a matrix factorization model.
 
     Example Usage::
@@ -74,7 +74,7 @@ def MatrixFactorizationModel(
 
     Returns
     -------
-    Union[Model, RetrievalModel]
+    RetrievalModel
     """
 
     if not prediction_tasks:
@@ -99,7 +99,7 @@ def MatrixFactorizationModel(
         **kwargs,
     )
 
-    model = two_tower.connect(prediction_tasks)
+    model = RetrievalModel(two_tower, prediction_tasks)
 
     return model
 
@@ -125,7 +125,7 @@ def TwoTowerModel(
     metrics: MetricOrMetrics = ItemRetrievalTask.DEFAULT_METRICS,
     samplers: Sequence[ItemSampler] = (),
     **kwargs,
-) -> Union[Model, RetrievalModel]:
+) -> RetrievalModel:
     """Builds the Two-tower architecture, as proposed in [1].
 
     Example Usage::
@@ -178,7 +178,7 @@ def TwoTowerModel(
 
     Returns
     -------
-    Union[Model, RetrievalModel]
+    RetrievalModel
     """
 
     if not prediction_tasks:
@@ -203,7 +203,7 @@ def TwoTowerModel(
         **kwargs,
     )
 
-    model = two_tower.connect(prediction_tasks)
+    model = RetrievalModel(two_tower, prediction_tasks)
 
     return model
 
@@ -292,4 +292,6 @@ def YoutubeDNNRetrievalModel(
         num_sampled=num_sampled,
     )
 
-    return inputs.connect(top_block, task)
+    # TODO: Figure out how to make this fit as
+    # a RetrievalModel (which must have a RetrievalBlock)
+    return Model(inputs, top_block, task)

merlin/models/tf/models/retrieval.py

@@ -35,7 +35,7 @@ def mark_run_eagerly_modes(*args, **kwargs):
 
 
 def assert_body_works_in_model(dataset, body, run_eagerly, num_epochs=5):
-    model = body.connect(BinaryClassificationTask("click"))
+    model = Model(body, BinaryClassificationTask("click"))
     model.compile(optimizer="adam", run_eagerly=run_eagerly)
 
     losses = model.fit(dataset, batch_size=50, epochs=num_epochs)

merlin/models/tf/utils/testing_utils.py

@@ -43,7 +43,6 @@ parentdir_prefix = merlin-models-
 [mypy]
 ignore_missing_imports = True
 no_implicit_optional = True
-plugins = merlin.models.mypy
 
 [codespell]
 skip = .*pb2.py,./.git,./.github,./bench,./dist,./docs/build,.*egg-info.*,versioneer.py,*.csv,*.parquet,./.mypy_cache