Extend ContrastiveOutput to support sequential encoders

merlin/models/tf/outputs/contrastive.py

@@ -223,11 +223,17 @@ def call(self, inputs, features=None, targets=None, training=False, testing=Fals
     def call_contrastive(self, inputs, features, targets, training=False, testing=False):
         if isinstance(inputs, dict) and self.query_name in inputs:
             query_embedding = inputs[self.query_name]
-        elif isinstance(inputs, tf.Tensor):
+        elif isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
             query_embedding = inputs
         else:
             raise ValueError("Couldn't infer query embedding")
 
+        is_ragged = isinstance(query_embedding, tf.RaggedTensor)
+        if is_ragged:
+            # Get flat values of the ragged tensor
+            original_query_embedding = query_embedding
+            query_embedding = query_embedding.flat_values
+
         if self.has_candidate_weights:
             positive_id = targets
             if isinstance(targets, dict):
@@ -237,6 +243,17 @@ def call_contrastive(self, inputs, features, targets, training=False, testing=Fa
             positive_id = features[self.col_schema.name]
             positive_embedding = inputs[self.candidate_name]
 
+        if isinstance(positive_id, tf.RaggedTensor):
+            # Select positive candidates at masked positions
+            target_mask = positive_id._keras_mask.with_row_splits_dtype(
+                positive_id.row_splits.dtype
+            )
+            # Flatten target tensor to have the same shape as the query tensor
+            positive_id = tf.ragged.boolean_mask(positive_id, target_mask)
+            original_target = positive_id
+            positive_id = positive_id.flat_values
+            positive_embedding = tf.ragged.boolean_mask(positive_embedding, target_mask).flat_values
+
         positive = Candidate(id=positive_id, metadata={**features}).with_embedding(
             positive_embedding
         )
@@ -248,7 +265,13 @@ def call_contrastive(self, inputs, features, targets, training=False, testing=Fa
         ):
             negative = negative.with_embedding(self.embedding_lookup(negative.id))
 
-        return self.outputs(query_embedding, positive, negative)
+        logits = self.outputs(query_embedding, positive, negative)
+        if is_ragged:
+            logits.copy_with_updates(
+                outputs=original_query_embedding.with_flat_values(logits.outputs),
+                targets=original_target.with_flat_values(logits.targets),
+            )
+        return logits
 
     def outputs(
         self, query_embedding: tf.Tensor, positive: Candidate, negative: Candidate
@@ -382,7 +405,7 @@ def sample_negatives(
         return negatives, positive
 
     def embedding_lookup(self, ids: tf.Tensor):
-        return self.to_call.embedding_lookup(tf.squeeze(ids))
+        return self.to_call.embedding_lookup(tf.squeeze(ids, axis=-1))
 
     def to_dataset(self, gpu=None) -> merlin.io.Dataset:
         return merlin.io.Dataset(tf_utils.tensor_to_df(self.to_call.embeddings, gpu=gpu))

tests/unit/tf/transformers/test_block.py

@@ -507,3 +507,46 @@ def test_transformer_model_with_masking_and_broadcast_to_sequence(
         reload_model=False,
         fit_kwargs={"pre": fit_pre},
     )
+
+
+@pytest.mark.parametrize("pre", [mm.SequenceMaskRandom, mm.SequencePredictNext])
+def test_transformer_encoder_with_contrastive_output(sequence_testing_data: Dataset, pre):
+    dmodel = 32
+    seq_schema = sequence_testing_data.schema.select_by_tag(Tags.SEQUENCE)
+    target_schema = sequence_testing_data.schema.select_by_tag(Tags.ITEM_ID)
+    model_schema = seq_schema + target_schema
+    target = target_schema.column_names[0]
+
+    input_block = mm.InputBlockV2(
+        model_schema,
+        categorical=mm.Embeddings(
+            model_schema.select_by_tag(Tags.CATEGORICAL), dim=dmodel, sequence_combiner=None
+        ),
+    )
+    transformer_block = XLNetBlock(d_model=dmodel, n_head=4, n_layer=1)
+    session_encoder = mm.Encoder(
+        input_block,
+        mm.MLPBlock([dmodel]),
+        transformer_block,
+    )
+    output_block = mm.ContrastiveOutput(
+        to_call=target_schema,
+        negative_samplers=mm.PopularityBasedSamplerV2(
+            max_num_samples=10,
+            max_id=1000,
+            min_id=1,
+        ),
+        logq_sampling_correction=True,
+    )
+    model = mm.RetrievalModelV2(query=session_encoder, output=output_block)
+
+    sequence_testing_data.schema = model_schema
+    loader = Loader(sequence_testing_data, batch_size=64, shuffle=False)
+
+    fit_pre = pre(schema=seq_schema, target=target, transformer=transformer_block)
+    model.compile()
+    _ = model.fit(loader, pre=fit_pre)
+
+    inputs, _ = loader.peek()
+    predictions = model(inputs)
+    assert list(predictions.shape) == [64, 51997]