Add test for retrieval model with transformer block

merlin/models/tf/core/encoder.py

@@ -71,7 +71,7 @@ def encode(
         elif isinstance(index, Tags):
             output_schema = self.schema.select_by_tag(index)
         else:
-            raise ValueError(f"Invalid index: {index}")
+            output_schema = None
 
         return self.batch_predict(
             dataset,

merlin/models/tf/outputs/contrastive.py

@@ -314,7 +314,9 @@ def sample_negatives(
         return negatives
 
     def embedding_lookup(self, ids: tf.Tensor):
-        return self.to_call.embedding_lookup(tf.squeeze(ids))
+        if ids.shape.rank == 2:
+            ids = tf.squeeze(ids, axis=1)
+        return self.to_call.embedding_lookup(ids)
 
     def to_dataset(self, gpu=None) -> merlin.io.Dataset:
         return merlin.io.Dataset(tf_utils.tensor_to_df(self.to_call.embeddings, gpu=gpu))

merlin/models/tf/transforms/sequence.py

@@ -219,7 +219,7 @@ def call(
         self._check_seq_inputs_targets(inputs)
 
         # Shifts the target column to be the next item of corresponding input column
-        new_target = inputs[self.target_name][:, 1:]
+        new_target = inputs[self.target_name][:, -1:]
         if targets is None:
             targets = dict({self.target_name: new_target})
         elif isinstance(targets, dict):

merlin/models/tf/utils/tf_utils.py

@@ -114,8 +114,12 @@ def rescore_false_negatives(
     Zeroes the logits of accidental negatives.
     """
     # Removing dimensions of size 1 from the shape of the item ids, if applicable
-    positive_item_ids = tf.cast(tf.squeeze(positive_item_ids), neg_samples_item_ids.dtype)
-    neg_samples_item_ids = tf.squeeze(neg_samples_item_ids)
+    if positive_item_ids.shape.rank == 2:
+        positive_item_ids = tf.squeeze(positive_item_ids, axis=1)
+    positive_item_ids = tf.cast(positive_item_ids, neg_samples_item_ids.dtype)
+
+    if neg_samples_item_ids.shape.rank == 2:
+        neg_samples_item_ids = tf.squeeze(neg_samples_item_ids, axis=1)
 
     # Reshapes positive and negative ids so that false_negatives_mask matches the scores shape
     false_negatives_mask = tf.equal(

tests/unit/tf/transformers/test_block.py

@@ -25,6 +25,59 @@ 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_next = mm.SequencePredictNext(schema=seq_schema, target=target)
+    loader = Loader(sequence_testing_data, batch_size=8, shuffle=False, transform=predict_next)
+
+    query_schema = seq_schema
+    output_schema = seq_schema.select_by_name(target)
+
+    query_encoder = mm.Encoder(
+        mm.InputBlockV2(
+            query_schema,
+            embeddings=mm.Embeddings(
+                query_schema.select_by_tag(Tags.CATEGORICAL), sequence_combiner=None
+            ),
+        ),
+        GPT2Block(d_model=48, n_head=4, n_layer=2, pre=mm.ReplaceMaskedEmbeddings()),
+        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={},
+    )
+
+    predictions = model.predict(loader)
+    assert list(predictions.shape) == [100, 51997]
+
+    query_embeddings = query_encoder.predict(loader)
+    assert list(query_embeddings.shape) == [100, 48]
+
+    # query_embeddings = model.query_embeddings(sequence_testing_data, batch_size=10).compute()
+    item_embeddings = model.candidate_embeddings().compute().to_numpy()
+
+    assert list(item_embeddings.shape) == [51997, 48]
+    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