Add pre-trained embeddings support to T4Rec input blocks (#690)

* add support of pretrained embeddings to t4rec input blocks -- v0

* add sequence_combiner

* fix merge conflict

* add pre-trained op to t4rec trainer

* remove custom padding from T4rec dataloader

* revert back custom padding

* Add tabular-normalization

* fix linting

* replace value_count with is_list to check if a feature is a list column

* extend support of 2-D pretrained embeddings

* update docstrings

* fix support of 2d/3d pretrained embeddings

* try to fix failing test in gpu-ci

* adding ops-bot yaml file (#708)

* try to fix torchscript error raised by gpu-ci

---------

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

Author: 3 people

tests/unit/torch/features/test_sequential.py

@@ -14,12 +14,17 @@
 # limitations under the License.
 #
 
+import numpy as np
 import pytest
+from merlin.dataloader.ops.embeddings import EmbeddingOperator
+from merlin.io import Dataset
+from merlin.schema import ColumnSchema
 from merlin.schema import Schema as CoreSchema
 from merlin.schema import Tags
 
 import transformers4rec.torch as tr
 from tests.conftest import parametrize_schemas
+from transformers4rec.torch.utils.data_utils import MerlinDataLoader
 
 
 @parametrize_schemas("yoochoose")
@@ -217,3 +222,173 @@ def test_sequential_and_non_sequential_tabular_features(schema, torch_yoochoose_
     outputs = tab_module(torch_yoochoose_like)
 
     assert list(outputs.shape) == [100, 20, 203]
+
+
+@pytest.mark.parametrize(
+    "pretrained_dim",
+    [None, 128, {"pretrained_item_id_embeddings": 128, "pretrained_user_id_embeddings": 128}],
+)
+def test_sequential_input_block_with_pretrained_embeddings(pretrained_dim):
+    data = tr.data.music_streaming_testing_data
+    seq_schema = data.merlin_schema.select_by_name(["item_id"])
+    # Set the property `dims` for the non-sequential feature: "user_id"
+    user_cardinality = data.merlin_schema["user_id"].int_domain.max + 1
+    seq_schema = seq_schema + CoreSchema(
+        [
+            ColumnSchema(
+                "user_id",
+                dtype=np.int32,
+                tags=[Tags.USER, Tags.CATEGORICAL],
+                properties={
+                    "domain": {"name": "user_id", "min": 0, "max": user_cardinality},
+                },
+                dims=(None,),
+            )
+        ]
+    )
+    batch_size, max_length = 128, 20
+    embedding_dim_default, item_dim, user_dim = 8, 32, 16
+
+    # generate pre-trained embeddings tables
+    item_cardinality = seq_schema["item_id"].int_domain.max + 1
+    np_emb_item_id = np.random.rand(item_cardinality, item_dim)
+    np_emb_user_id = np.random.rand(user_cardinality, user_dim)
+    embeddings_op_item = EmbeddingOperator(
+        np_emb_item_id, lookup_key="item_id", embedding_name="pretrained_item_id_embeddings"
+    )
+    embeddings_op_user = EmbeddingOperator(
+        np_emb_user_id, lookup_key="user_id", embedding_name="pretrained_user_id_embeddings"
+    )
+
+    # set dataloader with pre-trained embeddings
+    data_loader = MerlinDataLoader.from_schema(
+        seq_schema,
+        data.path,
+        batch_size=batch_size,
+        max_sequence_length=max_length,
+        transforms=[embeddings_op_item, embeddings_op_user],
+        shuffle=False,
+    )
+
+    batch, _ = next(iter(data_loader))
+
+    # Sequential input block with pre-trained features
+    inputs = tr.TabularSequenceFeatures.from_schema(
+        data_loader.output_schema,
+        max_sequence_length=20,
+        pretrained_output_dims=pretrained_dim,
+        aggregation=None,
+    )
+
+    # Sequential input + concat aggregation, which inherently performs broadcasting of 2-D features.
+    inputs_with_concat = tr.TabularSequenceFeatures.from_schema(
+        data_loader.output_schema,
+        embedding_dim_default=embedding_dim_default,
+        max_sequence_length=20,
+        aggregation="concat",
+    )
+
+    output = inputs.to(batch["item_id"].device).double()(batch)
+    concat_output = inputs_with_concat.to(batch["item_id"].device).double()(batch)
+
+    assert concat_output.shape[-1] == embedding_dim_default * 2 + item_dim + user_dim
+
+    assert "pretrained_item_id_embeddings" in output
+    if pretrained_dim is not None:
+        assert list(output["pretrained_item_id_embeddings"].shape) == [
+            batch_size,
+            max_length,
+            128,
+        ]
+        assert list(output["pretrained_user_id_embeddings"].shape) == [
+            batch_size,
+            128,
+        ]
+    else:
+        assert list(output["pretrained_item_id_embeddings"].shape) == [
+            batch_size,
+            max_length,
+            item_dim,
+        ]
+        assert list(output["pretrained_user_id_embeddings"].shape) == [
+            batch_size,
+            user_dim,
+        ]
+
+
+@pytest.mark.parametrize(
+    "pretrained_dim",
+    [None, 128, {"pretrained_item_id_embeddings": 128, "pretrained_user_id_embeddings": 128}],
+)
+def test_non_sequential_input_block_with_pretrained_embeddings(pretrained_dim):
+    data = tr.data.music_streaming_testing_data
+    seq_schema = data.merlin_schema.select_by_name(["item_id"])
+    # Set the property `dims` for the non-sequential feature: "user_id"
+    user_cardinality = data.merlin_schema["user_id"].int_domain.max + 1
+    seq_schema = seq_schema + CoreSchema(
+        [
+            ColumnSchema(
+                "user_id",
+                dtype=np.int32,
+                tags=[Tags.USER, Tags.CATEGORICAL],
+                properties={
+                    "domain": {"name": "user_id", "min": 0, "max": user_cardinality},
+                },
+                dims=(None,),
+            )
+        ]
+    )
+    batch_size, max_length = 128, 20
+    item_dim, user_dim = 32, 16
+
+    # generate pre-trained embeddings tables
+    item_cardinality = seq_schema["item_id"].int_domain.max + 1
+    np_emb_item_id = np.random.rand(item_cardinality, item_dim)
+    np_emb_user_id = np.random.rand(user_cardinality, user_dim)
+    embeddings_op_item = EmbeddingOperator(
+        np_emb_item_id, lookup_key="item_id", embedding_name="pretrained_item_id_embeddings"
+    )
+    embeddings_op_user = EmbeddingOperator(
+        np_emb_user_id, lookup_key="user_id", embedding_name="pretrained_user_id_embeddings"
+    )
+
+    # set dataloader with pre-trained embeddings
+    data_loader = MerlinDataLoader.from_schema(
+        seq_schema,
+        Dataset(data.path, schema=seq_schema),
+        batch_size=batch_size,
+        max_sequence_length=max_length,
+        transforms=[embeddings_op_item, embeddings_op_user],
+        shuffle=False,
+    )
+
+    batch, _ = next(iter(data_loader))
+
+    # Non-Sequential input block with a 3-D pre-trained feature
+    inputs = tr.TabularFeatures.from_schema(
+        data_loader.output_schema,
+        pretrained_output_dims=pretrained_dim,
+        sequence_combiner="mean",
+        aggregation=None,
+    )
+    output = inputs.to(batch["item_id"].device).double()(batch)
+
+    assert "pretrained_item_id_embeddings" in output
+    if pretrained_dim is not None:
+        assert list(output["pretrained_item_id_embeddings"].shape) == [
+            batch_size,
+            128,
+        ]
+        assert list(output["pretrained_user_id_embeddings"].shape) == [
+            batch_size,
+            128,
+        ]
+    else:
+        assert list(output["pretrained_item_id_embeddings"].shape) == [
+            batch_size,
+            item_dim,
+        ]
+        assert list(output["pretrained_user_id_embeddings"].shape) == [
+            batch_size,
+            user_dim,
+        ]

tests/unit/torch/test_trainer.py

@@ -598,3 +598,75 @@ def test_trainer_trop_k_with_wrong_task():
         recsys_trainer.predict(data.path)
 
     assert "Top-k prediction is specific to NextItemPredictionTask" in str(excinfo.value)
+
+
+def test_trainer_with_pretrained_embeddings():
+    import numpy as np
+    from merlin.dataloader.ops.embeddings import EmbeddingOperator
+    from merlin.io import Dataset
+
+    from transformers4rec.torch.utils.data_utils import MerlinDataLoader
+
+    data = tr.data.music_streaming_testing_data
+    schema = data.merlin_schema.select_by_name(
+        ["item_id", "item_category", "item_recency", "item_genres", "user_id"]
+    )
+    batch_size, max_length, pretrained_dim = 128, 20, 16
+    item_cardinality = schema["item_id"].int_domain.max + 1
+    np_emb_item_id = np.random.rand(item_cardinality, pretrained_dim)
+
+    embeddings_op = EmbeddingOperator(
+        np_emb_item_id, lookup_key="item_id", embedding_name="pretrained_item_id_embeddings"
+    )
+    # set dataloader with pre-trained embeddings
+    data_loader = MerlinDataLoader.from_schema(
+        schema,
+        Dataset(data.path, schema=schema),
+        max_sequence_length=max_length,
+        batch_size=batch_size,
+        transforms=[embeddings_op],
+        shuffle=False,
+    )
+
+    # set the model schema from data-loader
+    model_schema = data_loader.output_schema
+    inputs = tr.TabularSequenceFeatures.from_schema(
+        model_schema,
+        max_sequence_length=max_length,
+        pretrained_output_dims=8,
+        normalizer="layer-norm",
+        d_output=64,
+        masking="mlm",
+    )
+    transformer_config = tconf.XLNetConfig.build(64, 4, 2, 20)
+    task = tr.NextItemPredictionTask(weight_tying=True)
+    model = transformer_config.to_torch_model(inputs, task, max_sequence_length=max_length)
+
+    assert isinstance(model.input_schema, Schema)
+
+    args = trainer.T4RecTrainingArguments(
+        output_dir=".",
+        max_steps=5,
+        num_train_epochs=1,
+        per_device_train_batch_size=batch_size,
+        per_device_eval_batch_size=batch_size // 2,
+        max_sequence_length=max_length,
+        fp16=False,
+        report_to=[],
+        debug=["r"],
+    )
+    # Explicitly pass the merlin dataloader with pre-trained embeddings
+    recsys_trainer = tr.Trainer(
+        model=model,
+        args=args,
+        schema=schema,
+        train_dataloader=data_loader,
+        eval_dataloader=data_loader,
+        compute_metrics=True,
+    )
+
+    recsys_trainer.train()
+    eval_metrics = recsys_trainer.evaluate(eval_dataset=data.path, metric_key_prefix="eval")
+
+    assert isinstance(eval_metrics, dict)
+    assert eval_metrics["eval_/loss"] is not None

transformers4rec/torch/__init__.py

@@ -40,6 +40,7 @@
 from .features.embedding import (
     EmbeddingFeatures,
     FeatureConfig,
+    PretrainedEmbeddingFeatures,
     PretrainedEmbeddingsInitializer,
     SoftEmbedding,
     SoftEmbeddingFeatures,
@@ -104,6 +105,7 @@
     "EmbeddingFeatures",
     "SoftEmbeddingFeatures",
     "PretrainedEmbeddingsInitializer",
+    "PretrainedEmbeddingFeatures",
     "TabularSequenceFeatures",
     "SequenceEmbeddingFeatures",
     "FeatureConfig",