Added tests with pretrained embeddings for DLRM and DCN

Author: gabrielspmoreira

merlin/datasets/entertainment/music_streaming/schema.json

@@ -42,6 +42,16 @@
         "tag": [
           "categorical",
           "item"
+        ],
+        "extraMetadata": [
+          {
+            "_dims": [
+              [
+                0.0,
+                null
+              ]
+            ]
+          }
         ]
       }
     },

merlin/datasets/testing/schema.json

@@ -107,6 +107,16 @@
           "item_id",
           "categorical",
           "item"
+        ],
+        "extraMetadata": [
+          {
+            "_dims": [
+              [
+                0.0,
+                null
+              ]
+            ]
+          }
         ]
       }
     },

tests/unit/tf/blocks/test_dlrm.py

@@ -14,9 +14,11 @@
 # limitations under the License.
 #
 
+import numpy as np
 import pytest
 
 import merlin.models.tf as mm
+from merlin.dataloader.ops.embeddings import EmbeddingOperator
 from merlin.io import Dataset
 from merlin.schema import Tags
 
@@ -29,11 +31,11 @@ def test_dlrm_block(testing_data: Dataset):
         bottom_block=mm.MLPBlock([64]),
         top_block=mm.DenseResidualBlock(),
     )
-    features = mm.sample_batch(testing_data, batch_size=100, include_targets=False)
+    features = mm.sample_batch(testing_data, batch_size=10, include_targets=False)
     outputs = dlrm(features)
     num_features = len(schema.select_by_tag(Tags.CATEGORICAL)) + 1
     dot_product_dim = (num_features - 1) * num_features // 2
-    assert list(outputs.shape) == [100, dot_product_dim + 64]
+    assert list(outputs.shape) == [10, dot_product_dim + 64]
 
 
 def test_dlrm_block_no_top_block(testing_data: Dataset):
@@ -43,19 +45,19 @@ def test_dlrm_block_no_top_block(testing_data: Dataset):
         embedding_dim=64,
         bottom_block=mm.MLPBlock([64]),
     )
-    outputs = dlrm(mm.sample_batch(testing_data, batch_size=100, include_targets=False))
+    outputs = dlrm(mm.sample_batch(testing_data, batch_size=10, include_targets=False))
     num_features = len(schema.select_by_tag(Tags.CATEGORICAL)) + 1
     dot_product_dim = (num_features - 1) * num_features // 2
 
-    assert list(outputs.shape) == [100, dot_product_dim]
+    assert list(outputs.shape) == [10, dot_product_dim]
 
 
 def test_dlrm_block_no_continuous_features(testing_data: Dataset):
     schema = testing_data.schema.remove_by_tag(Tags.CONTINUOUS)
     dlrm = mm.DLRMBlock(schema, embedding_dim=64, top_block=mm.MLPBlock([32]))
-    outputs = dlrm(mm.sample_batch(testing_data, batch_size=100, include_targets=False))
+    outputs = dlrm(mm.sample_batch(testing_data, batch_size=10, include_targets=False))
 
-    assert list(outputs.shape) == [100, 32]
+    assert list(outputs.shape) == [10, 32]
 
 
 def test_dlrm_block_no_categ_features(testing_data: Dataset):
@@ -70,9 +72,9 @@ def test_dlrm_block_no_categ_features(testing_data: Dataset):
 def test_dlrm_block_single_categ_feature(testing_data: Dataset):
     schema = testing_data.schema.select_by_tag([Tags.ITEM_ID])
     dlrm = mm.DLRMBlock(schema, embedding_dim=64, top_block=mm.MLPBlock([32]))
-    outputs = dlrm(mm.sample_batch(testing_data, batch_size=100, include_targets=False))
+    outputs = dlrm(mm.sample_batch(testing_data, batch_size=10, include_targets=False))
 
-    assert list(outputs.shape) == [100, 32]
+    assert list(outputs.shape) == [10, 32]
 
 
 def test_dlrm_block_no_schema():
@@ -120,6 +122,43 @@ def test_dlrm_with_embeddings(testing_data: Dataset):
         bottom_block=mm.MLPBlock([embedding_dim]),
         top_block=mm.MLPBlock([top_dim]),
     )
-    outputs = dlrm(mm.sample_batch(testing_data, batch_size=100, include_targets=False))
+    outputs = dlrm(mm.sample_batch(testing_data, batch_size=10, include_targets=False))
 
-    assert list(outputs.shape) == [100, 4]
+    assert list(outputs.shape) == [10, 4]
+
+
+def test_dlrm_with_pretrained_embeddings(testing_data: Dataset):
+    embedding_dim = 12
+    top_dim = 4
+
+    item_cardinality = testing_data.schema["item_id"].int_domain.max + 1
+    pretrained_embedding = np.random.rand(item_cardinality, 12)
+
+    loader = mm.Loader(
+        testing_data,
+        batch_size=10,
+        transforms=[
+            EmbeddingOperator(
+                pretrained_embedding,
+                lookup_key="item_id",
+                embedding_name="pretrained_item_embeddings",
+            ),
+        ],
+    )
+    schema = loader.output_schema
+
+    embeddings = mm.Embeddings(schema.select_by_tag(Tags.CATEGORICAL), dim=embedding_dim)
+    pretrained_embeddings = mm.PretrainedEmbeddings(
+        schema.select_by_tag(Tags.EMBEDDING),
+        output_dims=embedding_dim,
+    )
+
+    dlrm = mm.DLRMBlock(
+        schema,
+        embeddings=mm.ParallelBlock(embeddings, pretrained_embeddings),
+        bottom_block=mm.MLPBlock([embedding_dim]),
+        top_block=mm.MLPBlock([top_dim]),
+    )
+    outputs = dlrm(mm.sample_batch(loader, include_targets=False))
+
+    assert list(outputs.shape) == [10, 4]

tests/unit/tf/models/test_ranking.py

@@ -20,6 +20,7 @@
 from tensorflow.keras import regularizers
 
 import merlin.models.tf as mm
+from merlin.dataloader.ops.embeddings import EmbeddingOperator
 from merlin.datasets.synthetic import generate_data
 from merlin.io import Dataset
 from merlin.models.tf.transforms.features import expected_input_cols_from_schema
@@ -164,6 +165,47 @@ def test_dcn_model(music_streaming_data, stacked, run_eagerly):
     testing_utils.model_test(model, music_streaming_data, run_eagerly=run_eagerly)
 
 
+@pytest.mark.parametrize("run_eagerly", [True, False])
+def test_dcn_model_with_pretrained_embeddings(music_streaming_data: Dataset, run_eagerly):
+    music_streaming_data.schema = music_streaming_data.schema.select_by_name(
+        ["item_id", "item_category", "user_age", "click"]
+    )
+
+    cardinality = music_streaming_data.schema["item_category"].int_domain.max + 1
+    pretrained_embedding = np.random.rand(cardinality, 12)
+
+    loader = mm.Loader(
+        music_streaming_data,
+        batch_size=10,
+        transforms=[
+            EmbeddingOperator(
+                pretrained_embedding,
+                lookup_key="item_category",
+                embedding_name="pretrained_category_embeddings",
+            ),
+        ],
+    )
+    schema = loader.output_schema
+
+    pretrained_embeddings = mm.PretrainedEmbeddings(
+        schema.select_by_tag(Tags.EMBEDDING),
+        output_dims=16,
+    )
+
+    input_block = mm.InputBlockV2(schema, pretrained_embeddings=pretrained_embeddings)
+
+    model = mm.DCNModel(
+        schema,
+        input_block=input_block,
+        depth=1,
+        deep_block=mm.MLPBlock([2]),
+        stacked=True,
+        prediction_tasks=mm.BinaryOutput("click"),
+    )
+
+    testing_utils.model_test(model, loader, run_eagerly=run_eagerly)
+
+
 @pytest.mark.parametrize("run_eagerly", [True, False])
 def test_deepfm_model_only_categ_feats(music_streaming_data, run_eagerly):
     music_streaming_data.schema = music_streaming_data.schema.select_by_name(