PyTorch Tabular integrations (#1559)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Jirka <jirka.borovec@seznam.cz>
Lightning-Universe · May 14, 2023 · 14c2755 · 14c2755
1 parent 59dab9b
commit 14c2755
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 15 deletions.
diff --git a/src/flash/core/integrations/pytorch_tabular/adapter.py b/src/flash/core/integrations/pytorch_tabular/adapter.py
@@ -50,6 +50,7 @@ def from_task(
             "categorical_dim": len(categorical_fields),
             "continuous_dim": num_features - len(categorical_fields),
             "output_dim": output_dim,
+            "embedded_cat_dim": sum([embd_dim for _, embd_dim in embedding_sizes]),
         }
         return cls(
             task_type,

diff --git a/src/flash/core/integrations/pytorch_tabular/backbones.py b/src/flash/core/integrations/pytorch_tabular/backbones.py
@@ -30,6 +30,7 @@
         AutoIntConfig,
         CategoryEmbeddingModelConfig,
         FTTransformerConfig,
+        GatedAdditiveTreeEnsembleConfig,
         NodeConfig,
         TabNetModelConfig,
         TabTransformerConfig,
@@ -88,8 +89,9 @@ def load_pytorch_tabular(
             AutoIntConfig,
             NodeConfig,
             CategoryEmbeddingModelConfig,
+            GatedAdditiveTreeEnsembleConfig,
         ],
-        ["tabnet", "tabtransformer", "fttransformer", "autoint", "node", "category_embedding"],
+        ["tabnet", "tabtransformer", "fttransformer", "autoint", "node", "category_embedding", "gate"],
     ):
         PYTORCH_TABULAR_BACKBONES(
             functools.partial(load_pytorch_tabular, model_config_class),

diff --git a/tests/tabular/classification/test_data_model_integration.py b/tests/tabular/classification/test_data_model_integration.py
@@ -39,8 +39,7 @@
         ("fttransformer", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
         ("autoint", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
         ("node", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
-        # ("category_embedding",  # todo: seems to be bug in tabular
-        #  {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
+        ("category_embedding", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
         # No categorical / numerical fields
         ("tabnet", {"categorical_fields": ["category"]}),
         ("tabnet", {"numerical_fields": ["scalar_a", "scalar_b"]}),

diff --git a/tests/tabular/classification/test_model.py b/tests/tabular/classification/test_model.py
@@ -55,7 +55,7 @@ class TestTabularClassifier(TaskTester):
                     {"backbone": "fttransformer"},
                     {"backbone": "autoint"},
                     {"backbone": "node"},
-                    # {"backbone": "category_embedding"},  # todo: seems to be bug in tabular
+                    {"backbone": "category_embedding"},
                 ],
             )
         ],
@@ -68,7 +68,7 @@ class TestTabularClassifier(TaskTester):
                     {"backbone": "fttransformer"},
                     {"backbone": "autoint"},
                     {"backbone": "node"},
-                    # {"backbone": "category_embedding"},  # todo: seems to be bug in tabular
+                    {"backbone": "category_embedding"},
                 ],
             )
         ],
@@ -81,7 +81,7 @@ class TestTabularClassifier(TaskTester):
                     {"backbone": "fttransformer"},
                     {"backbone": "autoint"},
                     {"backbone": "node"},
-                    # {"backbone": "category_embedding"},  # todo: seems to be bug in tabular
+                    {"backbone": "category_embedding"},
                 ],
             )
         ],

diff --git a/tests/tabular/regression/test_data_model_integration.py b/tests/tabular/regression/test_data_model_integration.py
@@ -48,8 +48,7 @@
         ("fttransformer", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
         ("autoint", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
         ("node", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
-        # ("category_embedding",  # todo: seems to be bug in tabular
-        #  {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
+        ("category_embedding", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
         # No categorical / numerical fields
         ("tabnet", {"categorical_fields": ["category"]}),
         ("tabnet", {"numerical_fields": ["scalar_a", "scalar_b"]}),
@@ -82,8 +81,7 @@ def test_regression_data_frame(backbone, fields, tmpdir):
         ("fttransformer", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
         ("autoint", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
         ("node", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
-        # ("category_embedding",  # todo: seems to be bug in tabular
-        #  {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
+        ("category_embedding", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
         # No categorical / numerical fields
         ("tabnet", {"categorical_fields": ["category"]}),
         ("tabnet", {"numerical_fields": ["scalar_a", "scalar_b"]}),
@@ -113,8 +111,7 @@ def test_regression_dicts(backbone, fields, tmpdir):
         ("fttransformer", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
         ("autoint", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
         ("node", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
-        # ("category_embedding",  # todo: seems to be bug in tabular
-        #  {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
+        ("category_embedding", {"categorical_fields": ["category"], "numerical_fields": ["scalar_a", "scalar_b"]}),
         # No categorical / numerical fields
         ("tabnet", {"categorical_fields": ["category"]}),
         ("tabnet", {"numerical_fields": ["scalar_a", "scalar_b"]}),

diff --git a/tests/tabular/regression/test_model.py b/tests/tabular/regression/test_model.py
@@ -53,7 +53,7 @@ class TestTabularRegressor(TaskTester):
                     {"backbone": "fttransformer"},
                     {"backbone": "autoint"},
                     {"backbone": "node"},
-                    # {"backbone": "category_embedding"},  # todo: seems to be bug in tabular
+                    {"backbone": "category_embedding"},
                 ],
             )
         ],
@@ -66,7 +66,7 @@ class TestTabularRegressor(TaskTester):
                     {"backbone": "fttransformer"},
                     {"backbone": "autoint"},
                     {"backbone": "node"},
-                    # {"backbone": "category_embedding"},  # todo: seems to be bug in tabular
+                    {"backbone": "category_embedding"},
                 ],
             )
         ],
@@ -79,7 +79,7 @@ class TestTabularRegressor(TaskTester):
                     {"backbone": "fttransformer"},
                     {"backbone": "autoint"},
                     {"backbone": "node"},
-                    # {"backbone": "category_embedding"},  # todo: seems to be bug in tabular
+                    {"backbone": "category_embedding"},
                 ],
             )
         ],