Merge pull request #258 from bishnukhadka/tabr

TabR integration + compatibility check for PLR embedding with pre-processing types.

Author: AnFreTh

.gitignore

@@ -83,6 +83,7 @@ Mambular is a Python package that brings the power of advanced deep learning arc
 | `Trompt`        | Trompt: Towards a Better Deep Neural Network for Tabular Data introduced [here](https://arxiv.org/abs/2305.18446).                                  |
 | `Tangos`        | Tangos: Regularizing Tabular Neural Networks through Gradient Orthogonalization and Specialization introduced [here](https://openreview.net/pdf?id=n6H86gW8u0d).                                  |
 | `ModernNCA`        | Revisiting Nearest Neighbor for Tabular Data: A Deep Tabular Baseline Two Decades Later introduced [here](https://arxiv.org/abs/2407.03257).                                  |
+| `TabR` | TabR: Tabular Deep Learning Meets Nearest Neighbors in 2023 [here](https://arxiv.org/abs/2307.14338) |
 
 
 

README.md

@@ -125,6 +125,8 @@ def __init__(self, num_feature_info, cat_feature_info, emb_feature_info, config)
         if self.layer_norm_after_embedding:
             self.embedding_norm = nn.LayerNorm(self.d_model)
 
+        self.feature_info = (num_feature_info, cat_feature_info, emb_feature_info)
+
     def forward(self, num_features, cat_features, emb_features):
         """Defines the forward pass of the model.
 
@@ -171,6 +173,8 @@ def forward(self, num_features, cat_features, emb_features):
 
         # Process numerical embeddings based on embedding_type
         if self.embedding_type == "plr":
+            # check pre-processing type compatibility with plr
+            self.check_plr_embedding_compatibility(self.feature_info)
             # For PLR, pass all numerical features together
             if num_features is not None:
                 num_features = torch.stack(num_features, dim=1).squeeze(
@@ -226,6 +230,21 @@ def forward(self, num_features, cat_features, emb_features):
             x = self.embedding_dropout(x)
 
         return x
+    
+    def check_plr_embedding_compatibility(self, feature_info:tuple):
+        # List of incompatible preprocessing terms for PLR embedding
+        incompatible_terms = ['ple', 'one-hot', 'polynomial', 'splines', 'sigmoid', 'rbf']
+        
+        # Iterate through each dictionary in the tuple (data)
+        for sub_dict in feature_info:
+            # Iterate through each feature in the current dictionary
+            for feature, properties in sub_dict.items():
+                preprocessing = properties.get('preprocessing', '')
+                
+                # Check for incompatible terms in the preprocessing string
+                for term in incompatible_terms:
+                    if term in preprocessing:
+                        raise ValueError(f"PLR embedding type doesn't work with the '{term}' pre-processing method.\n")
 
 
 class OneHotEncoding(nn.Module):

mambular/arch_utils/layer_utils/embedding_layer.py

@@ -22,7 +22,7 @@ def __init__(
         self.save_hyperparameters(ignore=["feature_information"])
 
         self.returns_ensemble = False
-        self.uses_nca_candidates = True
+        self.uses_candidates = True
 
         self.T = config.temperature
         self.sample_rate = config.sample_rate
@@ -31,6 +31,7 @@ def __init__(
                 *feature_information,
                 config=config,
             )
+            
             input_dim = np.sum(
                 [len(info) * self.hparams.d_model for info in feature_information]
             )
@@ -75,7 +76,7 @@ def forward(self, *data):
             x = self.post_encoder(x)
         return self.tabular_head(x)
 
-    def nca_train(self, *data, targets, candidate_x, candidate_y):
+    def train_with_candidates(self, *data, targets, candidate_x, candidate_y):
         """NCA-style training forward pass selecting candidates."""
         if self.hparams.use_embeddings:
             x = self.embedding_layer(*data)
@@ -85,6 +86,7 @@ def nca_train(self, *data, targets, candidate_x, candidate_y):
             B, S, D = candidate_x.shape
             candidate_x = candidate_x.reshape(B, S * D)
         else:
+
             x = torch.cat([t for tensors in data for t in tensors], dim=1)
             candidate_x = torch.cat(
                 [t for tensors in candidate_x for t in tensors], dim=1
@@ -129,7 +131,7 @@ def nca_train(self, *data, targets, candidate_x, candidate_y):
 
         return logits
 
-    def nca_validate(self, *data, candidate_x, candidate_y):
+    def validate_with_candidates(self, *data, candidate_x, candidate_y):
         """Validation forward pass with NCA-style candidate selection."""
         if self.hparams.use_embeddings:
             x = self.embedding_layer(*data)
@@ -172,7 +174,7 @@ def nca_validate(self, *data, candidate_x, candidate_y):
 
         return logits
 
-    def nca_predict(self, *data, candidate_x, candidate_y):
+    def predict_with_candidates(self, *data, candidate_x, candidate_y):
         """Prediction forward pass with candidate selection."""
         if self.hparams.use_embeddings:
             x = self.embedding_layer(*data)