Updating public-API to use the lib with a single import (#205)

Author: marcromeyn

tests/tf/block/test_base.py

@@ -16,24 +16,24 @@
 
 import pytest
 
-tf4rec = pytest.importorskip("transformers4rec.tf")
+tr = pytest.importorskip("transformers4rec.tf")
 
 
 def test_sequential_block_yoochoose(yoochoose_schema, tf_yoochoose_like):
-    inputs = tf4rec.TabularFeatures.from_schema(yoochoose_schema, aggregation="concat")
+    inputs = tr.TabularFeatures.from_schema(yoochoose_schema, aggregation="concat")
 
-    body = tf4rec.SequentialBlock([inputs, tf4rec.MLPBlock([64])])
+    body = tr.SequentialBlock([inputs, tr.MLPBlock([64])])
 
     outputs = body(tf_yoochoose_like)
 
     assert list(outputs.shape) == [100, 64]
 
 
 def test_sequential_block_yoochoose_without_aggregation(yoochoose_schema, tf_yoochoose_like):
-    inputs = tf4rec.TabularFeatures.from_schema(yoochoose_schema)
+    inputs = tr.TabularFeatures.from_schema(yoochoose_schema)
 
     with pytest.raises(TypeError) as excinfo:
-        body = tf4rec.SequentialBlock([inputs, tf4rec.MLPBlock([64])])
+        body = tr.SequentialBlock([inputs, tr.MLPBlock([64])])
 
         body(tf_yoochoose_like)
 

tests/tf/block/test_dlrm.py

@@ -16,13 +16,13 @@
 
 import pytest
 
-tf4rec = pytest.importorskip("transformers4rec.tf")
+tr = pytest.importorskip("transformers4rec.tf")
 
 
 def test_dlrm_block_yoochoose(yoochoose_schema, tf_yoochoose_like):
-    dlrm = tf4rec.DLRMBlock.from_schema(yoochoose_schema, bottom_mlp=tf4rec.MLPBlock([64]))
+    dlrm = tr.DLRMBlock.from_schema(yoochoose_schema, bottom_mlp=tr.MLPBlock([64]))
 
-    body = tf4rec.SequentialBlock([dlrm, tf4rec.MLPBlock([64])])
+    body = tr.SequentialBlock([dlrm, tr.MLPBlock([64])])
 
     outputs = body(tf_yoochoose_like)
 

tests/tf/block/test_mlp.py

@@ -16,8 +16,8 @@
 
 import pytest
 
-tf4rec = pytest.importorskip("transformers4rec.tf")
 tf = pytest.importorskip("tensorflow")
+tr = pytest.importorskip("transformers4rec.tf")
 
 
 @pytest.mark.parametrize("dim", [32, 64])
@@ -29,12 +29,10 @@
 def test_mlp_block_yoochoose(
     yoochoose_schema, tf_yoochoose_like, dim, activation, dropout, normalization
 ):
-    inputs = tf4rec.TabularFeatures.from_schema(yoochoose_schema, aggregation="concat")
+    inputs = tr.TabularFeatures.from_schema(yoochoose_schema, aggregation="concat")
 
-    mlp = tf4rec.MLPBlock(
-        [dim], activation=activation, dropout=dropout, normalization=normalization
-    )
-    body = tf4rec.SequentialBlock([inputs, mlp])
+    mlp = tr.MLPBlock([dim], activation=activation, dropout=dropout, normalization=normalization)
+    body = tr.SequentialBlock([inputs, mlp])
 
     outputs = body(tf_yoochoose_like)
 

tests/tf/block/test_transformer.py

@@ -18,7 +18,7 @@
 
 from transformers4rec.config import transformer as tconf
 
-tf4rec = pytest.importorskip("transformers4rec.tf")
+tr = pytest.importorskip("transformers4rec.tf")
 
 config_classes = [
     tconf.XLNetConfig,
@@ -27,7 +27,7 @@
 ]
 
 # fixed parameters for tests
-lm_tasks = list(tf4rec.masking.masking_registry.keys())
+lm_tasks = list(tr.masking.masking_registry.keys())
 # lm_tasks.remove("permutation")
 
 
@@ -36,7 +36,7 @@
 def test_transformer_block(yoochoose_schema, tf_yoochoose_like, task):
 
     col_group = yoochoose_schema
-    tab_module = tf4rec.TabularSequenceFeatures.from_schema(
+    tab_module = tr.TabularSequenceFeatures.from_schema(
         col_group,
         max_sequence_length=20,
         aggregation="sequential-concat",
@@ -47,11 +47,11 @@ def test_transformer_block(yoochoose_schema, tf_yoochoose_like, task):
         d_model=64, n_head=4, n_layer=2, total_seq_length=20
     )
 
-    block = tf4rec.SequentialBlock(
+    block = tr.SequentialBlock(
         [
             tab_module,
-            tf4rec.MLPBlock([64]),
-            tf4rec.TransformerBlock(transformer_config, masking=tab_module.masking),
+            tr.MLPBlock([64]),
+            tr.TransformerBlock(transformer_config, masking=tab_module.masking),
         ]
     )
 

tests/tf/conftest.py

@@ -25,7 +25,7 @@
 from merlin_standard_lib.utils.proto_utils import has_field, proto_text_to_better_proto
 
 tf = pytest.importorskip("tensorflow")
-tf4rec = pytest.importorskip("transformers4rec.tf")
+tr = pytest.importorskip("transformers4rec.tf")
 
 NUM_EXAMPLES = 1000
 MAX_CARDINALITY = 100
@@ -56,7 +56,7 @@ def tf_cat_features():
 
 @pytest.fixture
 def tf_yoochoose_tabular_features(yoochoose_schema):
-    return tf4rec.TabularFeatures.from_schema(
+    return tr.TabularFeatures.from_schema(
         yoochoose_schema,
         max_sequence_length=20,
         continuous_projection=64,
@@ -66,7 +66,7 @@ def tf_yoochoose_tabular_features(yoochoose_schema):
 
 @pytest.fixture
 def tf_yoochoose_tabular_sequence_features(yoochoose_schema):
-    return tf4rec.TabularSequenceFeatures.from_schema(
+    return tr.TabularSequenceFeatures.from_schema(
         yoochoose_schema,
         max_sequence_length=20,
         continuous_projection=64,

tests/tf/features/test_continuous.py

@@ -17,29 +17,29 @@
 import pytest
 
 from merlin_standard_lib import Tag
-from tests.tf import _utils as test_utils
 
-tf4rec = pytest.importorskip("transformers4rec.tf")
+tr = pytest.importorskip("transformers4rec.tf")
+test_utils = pytest.importorskip("transformers4rec.tf.utils.testing_utils")
 
 
 def test_continuous_features(tf_con_features):
     features = ["a", "b"]
-    con = tf4rec.ContinuousFeatures(features)(tf_con_features)
+    con = tr.ContinuousFeatures(features)(tf_con_features)
 
     assert list(con.keys()) == features
 
 
 def test_continuous_features_yoochoose(yoochoose_schema, tf_yoochoose_like):
     schema = yoochoose_schema.select_by_tag(Tag.CONTINUOUS)
 
-    inputs = tf4rec.ContinuousFeatures.from_schema(schema)
+    inputs = tr.ContinuousFeatures.from_schema(schema)
     outputs = inputs(tf_yoochoose_like)
 
     assert list(outputs.keys()) == schema.column_names
 
 
 def test_serialization_continuous_features(yoochoose_schema, tf_yoochoose_like):
-    inputs = tf4rec.ContinuousFeatures.from_schema(yoochoose_schema)
+    inputs = tr.ContinuousFeatures.from_schema(yoochoose_schema)
 
     copy_layer = test_utils.assert_serialization(inputs)
 
@@ -50,7 +50,7 @@ def test_serialization_continuous_features(yoochoose_schema, tf_yoochoose_like):
 def test_continuous_features_yoochoose_model(yoochoose_schema, tf_yoochoose_like, run_eagerly):
     schema = yoochoose_schema.select_by_tag(Tag.CONTINUOUS)
 
-    inputs = tf4rec.ContinuousFeatures.from_schema(schema, aggregation="concat")
-    body = tf4rec.SequentialBlock([inputs, tf4rec.MLPBlock([64])])
+    inputs = tr.ContinuousFeatures.from_schema(schema, aggregation="concat")
+    body = tr.SequentialBlock([inputs, tr.MLPBlock([64])])
 
     test_utils.assert_body_works_in_model(tf_yoochoose_like, inputs, body, run_eagerly)

tests/tf/features/test_embedding.py

@@ -18,18 +18,18 @@
 from tensorflow.python.ops import init_ops_v2
 
 from merlin_standard_lib import Tag
-from tests.tf import _utils as test_utils
 
-tf4rec = pytest.importorskip("transformers4rec.tf")
+tr = pytest.importorskip("transformers4rec.tf")
+test_utils = pytest.importorskip("transformers4rec.tf.utils.testing_utils")
 
 
 def test_embedding_features(tf_cat_features):
     dim = 15
     feature_config = {
-        f: tf4rec.FeatureConfig(tf4rec.TableConfig(100, dim, name=f, initializer=None))
+        f: tr.FeatureConfig(tr.TableConfig(100, dim, name=f, initializer=None))
         for f in tf_cat_features.keys()
     }
-    embeddings = tf4rec.EmbeddingFeatures(feature_config)(tf_cat_features)
+    embeddings = tr.EmbeddingFeatures(feature_config)(tf_cat_features)
 
     assert list(embeddings.keys()) == list(feature_config.keys())
     assert all([emb.shape[-1] == dim for emb in embeddings.values()])
@@ -38,7 +38,7 @@ def test_embedding_features(tf_cat_features):
 def test_embedding_features_yoochoose(yoochoose_schema, tf_yoochoose_like):
     schema = yoochoose_schema.select_by_tag(Tag.CATEGORICAL)
 
-    emb_module = tf4rec.EmbeddingFeatures.from_schema(schema)
+    emb_module = tr.EmbeddingFeatures.from_schema(schema)
     embeddings = emb_module(tf_yoochoose_like)
 
     assert list(embeddings.keys()) == schema.column_names
@@ -48,7 +48,7 @@ def test_embedding_features_yoochoose(yoochoose_schema, tf_yoochoose_like):
 
 
 def test_serialization_embedding_features(yoochoose_schema, tf_yoochoose_like):
-    inputs = tf4rec.EmbeddingFeatures.from_schema(yoochoose_schema)
+    inputs = tr.EmbeddingFeatures.from_schema(yoochoose_schema)
 
     copy_layer = test_utils.assert_serialization(inputs)
 
@@ -66,16 +66,16 @@ def test_serialization_embedding_features(yoochoose_schema, tf_yoochoose_like):
 def test_embedding_features_yoochoose_model(yoochoose_schema, tf_yoochoose_like, run_eagerly):
     schema = yoochoose_schema.select_by_tag(Tag.CATEGORICAL)
 
-    inputs = tf4rec.EmbeddingFeatures.from_schema(schema, aggregation="concat")
-    body = tf4rec.SequentialBlock([inputs, tf4rec.MLPBlock([64])])
+    inputs = tr.EmbeddingFeatures.from_schema(schema, aggregation="concat")
+    body = tr.SequentialBlock([inputs, tr.MLPBlock([64])])
 
     test_utils.assert_body_works_in_model(tf_yoochoose_like, inputs, body, run_eagerly)
 
 
 def test_embedding_features_yoochoose_custom_dims(yoochoose_schema, tf_yoochoose_like):
     schema = yoochoose_schema.select_by_tag(Tag.CATEGORICAL)
 
-    emb_module = tf4rec.EmbeddingFeatures.from_schema(
+    emb_module = tr.EmbeddingFeatures.from_schema(
         schema, embedding_dims={"item_id/list": 100}, embedding_dim_default=64
     )
 
@@ -91,7 +91,7 @@ def test_embedding_features_yoochoose_custom_dims(yoochoose_schema, tf_yoochoose
 def test_embedding_features_yoochoose_infer_embedding_sizes(yoochoose_schema, tf_yoochoose_like):
     schema = yoochoose_schema.select_by_tag(Tag.CATEGORICAL)
 
-    emb_module = tf4rec.EmbeddingFeatures.from_schema(
+    emb_module = tr.EmbeddingFeatures.from_schema(
         schema, infer_embedding_sizes=True, infer_embedding_sizes_multiplier=3.0
     )
 
@@ -112,7 +112,7 @@ def test_embedding_features_yoochoose_custom_initializers(yoochoose_schema, tf_y
     CATEGORY_STD = 0.1
 
     schema = yoochoose_schema.select_by_tag(Tag.CATEGORICAL)
-    emb_module = tf4rec.EmbeddingFeatures.from_schema(
+    emb_module = tr.EmbeddingFeatures.from_schema(
         schema,
         embeddings_initializers={
             "item_id/list": init_ops_v2.TruncatedNormal(mean=ITEM_MEAN, stddev=ITEM_STD),

tests/tf/features/test_sequence.py

@@ -17,14 +17,14 @@
 import pytest
 
 from merlin_standard_lib import Tag
-from tests.tf import _utils as test_utils
 
-tf4rec = pytest.importorskip("transformers4rec.tf")
+tr = pytest.importorskip("transformers4rec.tf")
+test_utils = pytest.importorskip("transformers4rec.tf.utils.testing_utils")
 
 
 def test_sequence_embedding_features(yoochoose_schema, tf_yoochoose_like):
     schema = yoochoose_schema.select_by_tag(Tag.CATEGORICAL)
-    emb_module = tf4rec.SequenceEmbeddingFeatures.from_schema(schema)
+    emb_module = tr.SequenceEmbeddingFeatures.from_schema(schema)
 
     outputs = emb_module(tf_yoochoose_like)
 
@@ -35,7 +35,7 @@ def test_sequence_embedding_features(yoochoose_schema, tf_yoochoose_like):
 
 
 def test_serialization_sequence_embedding_features(yoochoose_schema, tf_yoochoose_like):
-    inputs = tf4rec.SequenceEmbeddingFeatures.from_schema(yoochoose_schema)
+    inputs = tr.SequenceEmbeddingFeatures.from_schema(yoochoose_schema)
 
     copy_layer = test_utils.assert_serialization(inputs)
 
@@ -53,17 +53,17 @@ def test_serialization_sequence_embedding_features(yoochoose_schema, tf_yoochoos
 def test_sequence_embedding_features_yoochoose_model(
     yoochoose_schema, tf_yoochoose_like, run_eagerly
 ):
-    inputs = tf4rec.TabularSequenceFeatures.from_schema(
+    inputs = tr.TabularSequenceFeatures.from_schema(
         yoochoose_schema, max_sequence_length=20, aggregation="sequential-concat"
     )
 
-    body = tf4rec.SequentialBlock([inputs, tf4rec.MLPBlock([64])])
+    body = tr.SequentialBlock([inputs, tr.MLPBlock([64])])
 
     test_utils.assert_body_works_in_model(tf_yoochoose_like, inputs, body, run_eagerly)
 
 
 def test_sequence_tabular_features_with_projection(yoochoose_schema, tf_yoochoose_like):
-    tab_module = tf4rec.TabularSequenceFeatures.from_schema(
+    tab_module = tr.TabularSequenceFeatures.from_schema(
         yoochoose_schema, max_sequence_length=20, continuous_projection=64
     )
 
@@ -75,7 +75,7 @@ def test_sequence_tabular_features_with_projection(yoochoose_schema, tf_yoochoos
 
 
 def test_serialization_sequence_tabular_features(yoochoose_schema, tf_yoochoose_like):
-    inputs = tf4rec.TabularSequenceFeatures.from_schema(yoochoose_schema)
+    inputs = tr.TabularSequenceFeatures.from_schema(yoochoose_schema)
 
     copy_layer = test_utils.assert_serialization(inputs)
 
@@ -88,12 +88,12 @@ def test_tabular_features_yoochoose_direct(
     tf_yoochoose_like,
     run_eagerly,
 ):
-    continuous_layer = tf4rec.ContinuousFeatures.from_schema(yoochoose_schema, tags=["continuous"])
-    categorical_layer = tf4rec.SequenceEmbeddingFeatures.from_schema(
+    continuous_layer = tr.ContinuousFeatures.from_schema(yoochoose_schema, tags=["continuous"])
+    categorical_layer = tr.SequenceEmbeddingFeatures.from_schema(
         yoochoose_schema, tags=["categorical"]
     )
 
-    inputs = tf4rec.TabularSequenceFeatures(
+    inputs = tr.TabularSequenceFeatures(
         continuous_layer=continuous_layer,
         categorical_layer=categorical_layer,
         aggregation="sequential-concat",
@@ -141,12 +141,12 @@ def test_tabular_features_yoochoose_direct(
 
 def test_sequence_tabular_features_with_projection_and_d_output(yoochoose_schema):
     with pytest.raises(ValueError) as excinfo:
-        tf4rec.TabularSequenceFeatures.from_schema(
+        tr.TabularSequenceFeatures.from_schema(
             yoochoose_schema,
             max_sequence_length=20,
             continuous_projection=64,
             d_output=100,
-            projection=tf4rec.MLPBlock([64]),
+            projection=tr.MLPBlock([64]),
             masking="causal",
         )