Add music-streaming synthetic data to test the support of all predictions tasks with the Trainer class

tests/torch/test_trainer.py

@@ -20,6 +20,7 @@
 import pytest
 
 from transformers4rec.config import trainer
+from transformers4rec.config import transformer as tconf
 
 pytorch = pytest.importorskip("torch")
 tr = pytest.importorskip("transformers4rec.torch")
@@ -314,3 +315,64 @@ def test_evaluate_results(torch_yoochoose_next_item_prediction_model):
     result_2 = {k: result_2[k] for k in default_metric}
 
     assert result_1 == result_2
+
+
+@pytest.mark.parametrize(
+    "task",
+    [tr.NextItemPredictionTask(weight_tying=True)],
+)
+def test_trainer_music_streaming(task):
+    # TODO: Add binary and regression tasks
+    pytest.importorskip("pyarrow")
+    data = tr.data.music_streaming_testing_data
+    batch_size = 16
+
+    inputs = tr.TabularSequenceFeatures.from_schema(
+        data.schema,
+        max_sequence_length=20,
+        d_output=64,
+        masking="mlm",
+    )
+    transformer_config = tconf.XLNetConfig.build(64, 4, 2, 20)
+    model = transformer_config.to_torch_model(inputs, task)
+
+    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,
+        data_loader_engine="pyarrow",
+        max_sequence_length=20,
+        fp16=False,
+        no_cuda=True,
+        report_to=[],
+        debug=["r"],
+    )
+
+    recsys_trainer = tr.Trainer(
+        model=model,
+        args=args,
+        schema=data.schema,
+        train_dataset_or_path=data.path,
+        eval_dataset_or_path=data.path,
+        test_dataset_or_path=data.path,
+        compute_metrics=True,
+    )
+
+    eval_metrics = recsys_trainer.evaluate(eval_dataset=data.path, metric_key_prefix="eval")
+    predictions = recsys_trainer.predict(data.path)
+
+    assert isinstance(eval_metrics, dict)
+    default_metric = [
+        "eval_/next-item/ndcg_at_10",
+        "eval_/next-item/ndcg_at_20",
+        "eval_/next-item/avg_precision_at_10",
+        "eval_/next-item/avg_precision_at_20",
+        "eval_/next-item/recall_at_10",
+        "eval_/next-item/recall_at_20",
+    ]
+    assert set(default_metric).issubset(set(eval_metrics.keys()))
+    assert eval_metrics["eval_/loss"] is not None
+
+    assert predictions is not None

transformers4rec/data/__init__.py

@@ -16,6 +16,7 @@
 
 
 from .testing.dataset import tabular_sequence_testing_data
+from .testing.music_streaming.dataset import music_streaming_testing_data
 from .testing.tabular_data.dataset import tabular_testing_data
 
-__all__ = ["tabular_sequence_testing_data", "tabular_testing_data"]
+__all__ = ["tabular_sequence_testing_data", "tabular_testing_data", "music_streaming_testing_data"]

transformers4rec/data/testing/music_streaming/dataset.py

@@ -0,0 +1,5 @@
+import pathlib
+
+from transformers4rec.data.dataset import ParquetDataset
+
+music_streaming_testing_data: ParquetDataset = ParquetDataset(pathlib.Path(__file__).parent)

transformers4rec/data/testing/music_streaming/schema.json

@@ -0,0 +1,162 @@
+{
+    "feature": [
+      {
+        "name": "session_id",
+        "type": "INT",
+        "intDomain": {
+          "name": "session_id",
+          "max": "10000",
+          "isCategorical": true
+        },
+        "annotation": {
+          "tag": [
+            "categorical",
+            "session_id"
+          ]
+        }
+      },
+      {
+        "name": "item_id",
+        "type": "INT",
+        "valueCount": {
+          "min": "1",
+          "max": "20"
+        },
+        "intDomain": {
+          "name": "item_id",
+          "max": "10000",
+          "isCategorical": true,
+          "is_list": true
+        },
+        "annotation": {
+          "tag": [
+            "categorical",
+            "item_id",
+            "item", 
+            "list"
+          ]
+        }
+      },
+      {
+        "name": "item_category",
+        "type": "INT",
+        "is_list": true,
+        "valueCount": {
+          "min": "1",
+          "max": "20"
+        },
+        "intDomain": {
+          "name": "item_category",
+          "max": "100",
+          "isCategorical": true,
+          "is_list": true
+        },
+        "annotation": {
+          "tag": [
+            "categorical",
+            "item",
+            "list"
+
+          ]
+        }
+      },
+      {
+        "name": "item_recency",
+        "type": "FLOAT",
+        "valueCount": {
+          "min": "1",
+          "max": "20"
+        },
+        "floatDomain": {
+          "name": "item_recency",
+          "max": 1.0,
+          "is_list": true
+        },
+        "annotation": {
+          "tag": [
+            "continuous",
+            "item",
+            "list"
+          ]
+        }
+      },
+      {
+        "name": "item_genres",
+        "valueCount": {
+          "min": "1",
+          "max": "20"
+        },
+        "type": "INT",
+        "intDomain": {
+          "name": "genres",
+          "max": "100",
+          "isCategorical": true,
+          "is_list": true
+        },
+        "annotation": {
+          "tag": [
+            "categorical",
+            "item",
+            "list"
+          ]
+        }
+      },
+      {
+        "name": "user_id",
+        "type": "INT",
+        "intDomain": {
+          "name": "user_id",
+          "max": "10000",
+          "isCategorical": true
+        },
+        "annotation": {
+          "tag": [
+            "categorical",
+            "user_id"
+          ]
+        }
+      },
+      {
+        "name": "country",
+        "type": "INT",
+        "intDomain": {
+          "name": "country",
+          "max": "100",
+          "isCategorical": true
+        },
+        "annotation": {
+          "tag": [
+            "categorical",
+            "user"
+          ]
+        }
+      },
+      {
+        "name": "click",
+        "annotation": {
+          "tag": [
+            "binary_classification",
+            "target"
+          ]
+        }
+      },
+      {
+        "name": "play_percentage",
+        "annotation": {
+          "tag": [
+            "regression",
+            "target"
+          ]
+        }
+      },
+      {
+        "name": "like",
+        "annotation": {
+          "tag": [
+            "binary_classification",
+            "target"
+          ]
+        }
+      }
+    ]
+}