Add test for retrieval model with transformer block

merlin/models/tf/models/base.py

@@ -707,6 +707,11 @@ def train_step(self, data):
                 out = call_layer(self.train_pre, x, targets=y, features=x, training=True)
                 if isinstance(out, Prediction):
                     x, y = out.outputs, out.targets
+                elif isinstance(out, tuple):
+                    assert (
+                        len(out) == 2
+                    ), "output of `pre` must be a 2-tuple of x, y or `Prediction` tuple"
+                    x, y = out
                 else:
                     x = out
 

tests/unit/tf/transformers/test_block.py

@@ -25,6 +25,61 @@ def test_import():
     assert transformers is not None
 
 
+@pytest.mark.parametrize("run_eagerly", [True])
+def test_retrieval_transformer(sequence_testing_data: Dataset, run_eagerly):
+
+    seq_schema = sequence_testing_data.schema.select_by_tag(Tags.SEQUENCE).select_by_tag(
+        Tags.CATEGORICAL
+    )
+
+    target = sequence_testing_data.schema.select_by_tag(Tags.ITEM_ID).column_names[0]
+    predict_last = mm.SequencePredictLast(schema=seq_schema, target=target)
+    loader = Loader(sequence_testing_data, batch_size=8, shuffle=False)
+
+    query_schema = seq_schema
+    output_schema = seq_schema.select_by_name(target)
+
+    d_model = 48
+    query_encoder = mm.Encoder(
+        mm.InputBlockV2(
+            query_schema,
+            embeddings=mm.Embeddings(
+                query_schema.select_by_tag(Tags.CATEGORICAL), sequence_combiner=None
+            ),
+        ),
+        mm.MLPBlock([d_model]),
+        GPT2Block(d_model=d_model, n_head=2, n_layer=2),
+        tf.keras.layers.Lambda(lambda x: tf.reduce_mean(x, axis=1)),
+    )
+
+    model = mm.RetrievalModelV2(
+        query=query_encoder,
+        output=mm.ContrastiveOutput(output_schema, negative_samplers="in-batch"),
+    )
+
+    testing_utils.model_test(
+        model,
+        loader,
+        run_eagerly=run_eagerly,
+        reload_model=False,
+        metrics={},
+        fit_kwargs={"pre": predict_last},
+    )
+
+    predictions = model.predict(loader)
+    assert list(predictions.shape) == [100, 51997]
+
+    query_embeddings = query_encoder.predict(loader)
+    assert list(query_embeddings.shape) == [100, d_model]
+
+    item_embeddings = model.candidate_embeddings().compute().to_numpy()
+
+    assert list(item_embeddings.shape) == [51997, d_model]
+    predicitons_2 = np.dot(query_embeddings, item_embeddings.T)
+
+    np.testing.assert_allclose(predictions, predicitons_2, atol=1e-7)
+
+
 def test_transformer_encoder():
     NUM_ROWS = 100
     SEQ_LENGTH = 10