include encoding function to create embeddings

Author: AnFreTh

mambular/base_models/basemodel.py

@@ -33,7 +33,11 @@ def save_hyperparameters(self, ignore=[]):
             List of keys to ignore while saving hyperparameters, by default [].
         """
         # Filter the config and extra hparams for ignored keys
-        config_hparams = {k: v for k, v in vars(self.config).items() if k not in ignore} if self.config else {}
+        config_hparams = (
+            {k: v for k, v in vars(self.config).items() if k not in ignore}
+            if self.config
+            else {}
+        )
         extra_hparams = {k: v for k, v in self.extra_hparams.items() if k not in ignore}
         config_hparams.update(extra_hparams)
 
@@ -148,7 +152,9 @@ def initialize_pooling_layers(self, config, n_inputs):
         """Initializes the layers needed for learnable pooling methods based on self.hparams.pooling_method."""
         if self.hparams.pooling_method == "learned_flatten":
             # Flattening + Linear layer
-            self.learned_flatten_pooling = nn.Linear(n_inputs * config.dim_feedforward, config.dim_feedforward)
+            self.learned_flatten_pooling = nn.Linear(
+                n_inputs * config.dim_feedforward, config.dim_feedforward
+            )
 
         elif self.hparams.pooling_method == "attention":
             # Attention-based pooling with learnable attention weights
@@ -216,3 +222,29 @@ def pool_sequence(self, out):
             return out
         else:
             raise ValueError(f"Invalid pooling method: {self.hparams.pooling_method}")
+
+    def encode(self, num_features, cat_features):
+        if not hasattr(self, "embedding_layer"):
+            raise ValueError("The model does not have an embedding layer")
+
+        # Check if at least one of the contextualized embedding methods exists
+        valid_layers = ["mamba", "rnn", "lstm", "encoder"]
+        available_layer = next(
+            (attr for attr in valid_layers if hasattr(self, attr)), None
+        )
+
+        if not available_layer:
+            raise ValueError("The model does not generate contextualized embeddings")
+
+        # Get the actual layer and call it
+        x = self.embedding_layer(num_features=num_features, cat_features=cat_features)
+
+        if getattr(self.hparams, "shuffle_embeddings", False):
+            x = x[:, self.perm, :]
+
+        layer = getattr(self, available_layer)
+        if available_layer == "rnn":
+            embeddings, _ = layer(x)
+        else:
+            embeddings = layer(x)
+        return embeddings