adapt first only regressor and classifier to handle embeddings

Author: AnFreTh

mambular/models/sklearn_base_classifier.py

@@ -8,15 +8,19 @@
 import torch
 from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint, ModelSummary
 from sklearn.base import BaseEstimator
-from sklearn.metrics import accuracy_score, log_loss, mean_squared_error
+from sklearn.metrics import accuracy_score, log_loss
 from skopt import gp_minimize
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
 from ..base_models.lightning_wrapper import TaskModel
 from ..data_utils.datamodule import MambularDataModule
 from ..preprocessing import Preprocessor
-from ..utils.config_mapper import activation_mapper, get_search_space, round_to_nearest_16
+from ..utils.config_mapper import (
+    activation_mapper,
+    get_search_space,
+    round_to_nearest_16,
+)
 
 
 class SklearnBaseClassifier(BaseEstimator):
@@ -39,11 +43,15 @@ def __init__(self, model, config, **kwargs):
         ]
 
         self.config_kwargs = {
-            k: v for k, v in kwargs.items() if k not in self.preprocessor_arg_names and not k.startswith("optimizer")
+            k: v
+            for k, v in kwargs.items()
+            if k not in self.preprocessor_arg_names and not k.startswith("optimizer")
         }
         self.config = config(**self.config_kwargs)
 
-        preprocessor_kwargs = {k: v for k, v in kwargs.items() if k in self.preprocessor_arg_names}
+        preprocessor_kwargs = {
+            k: v for k, v in kwargs.items() if k in self.preprocessor_arg_names
+        }
 
         self.preprocessor = Preprocessor(**preprocessor_kwargs)
         self.task_model = None
@@ -63,7 +71,8 @@ def __init__(self, model, config, **kwargs):
         self.optimizer_kwargs = {
             k: v
             for k, v in kwargs.items()
-            if k not in ["lr", "weight_decay", "patience", "lr_patience", "optimizer_type"]
+            if k
+            not in ["lr", "weight_decay", "patience", "lr_patience", "optimizer_type"]
             and k.startswith("optimizer_")
         }
 
@@ -84,7 +93,10 @@ def get_params(self, deep=True):
         params.update(self.config_kwargs)
 
         if deep:
-            preprocessor_params = {"prepro__" + key: value for key, value in self.preprocessor.get_params().items()}
+            preprocessor_params = {
+                "prepro__" + key: value
+                for key, value in self.preprocessor.get_params().items()
+            }
             params.update(preprocessor_params)
 
         return params
@@ -102,8 +114,14 @@ def set_params(self, **parameters):
         self : object
             Estimator instance.
         """
-        config_params = {k: v for k, v in parameters.items() if not k.startswith("prepro__")}
-        preprocessor_params = {k.split("__")[1]: v for k, v in parameters.items() if k.startswith("prepro__")}
+        config_params = {
+            k: v for k, v in parameters.items() if not k.startswith("prepro__")
+        }
+        preprocessor_params = {
+            k.split("__")[1]: v
+            for k, v in parameters.items()
+            if k.startswith("prepro__")
+        }
 
         if config_params:
             self.config_kwargs.update(config_params)
@@ -125,6 +143,8 @@ def build_model(
         val_size: float = 0.2,
         X_val=None,
         y_val=None,
+        embeddings=None,
+        embeddings_val=None,
         random_state: int = 101,
         batch_size: int = 128,
         shuffle: bool = True,
@@ -201,20 +221,36 @@ def build_model(
             **dataloader_kwargs,
         )
 
-        self.data_module.preprocess_data(X, y, X_val, y_val, val_size=val_size, random_state=random_state)
+        self.data_module.preprocess_data(
+            X,
+            y,
+            X_val=X_val,
+            y_val=y_val,
+            embeddings_train=embeddings,
+            embeddings_val=embeddings_val,
+            val_size=val_size,
+            random_state=random_state,
+        )
 
         num_classes = len(np.unique(np.array(y)))
 
         self.task_model = TaskModel(
             model_class=self.base_model,  # type: ignore
             num_classes=num_classes,
             config=self.config,
-            cat_feature_info=self.data_module.cat_feature_info,
-            num_feature_info=self.data_module.num_feature_info,
-            lr_patience=(lr_patience if lr_patience is not None else self.config.lr_patience),
+            feature_information=(
+                self.data_module.num_feature_info,
+                self.data_module.cat_feature_info,
+                self.data_module.embedding_feature_info,
+            ),
+            lr_patience=(
+                lr_patience if lr_patience is not None else self.config.lr_patience
+            ),
             lr=lr if lr is not None else self.config.lr,
             lr_factor=lr_factor if lr_factor is not None else self.config.lr_factor,
-            weight_decay=(weight_decay if weight_decay is not None else self.config.weight_decay),
+            weight_decay=(
+                weight_decay if weight_decay is not None else self.config.weight_decay
+            ),
             train_metrics=train_metrics,
             val_metrics=val_metrics,
             optimizer_type=self.optimizer_type,
@@ -245,7 +281,9 @@ def get_number_of_params(self, requires_grad=True):
             If the model has not been built prior to calling this method.
         """
         if not self.built:
-            raise ValueError("The model must be built before the number of parameters can be estimated")
+            raise ValueError(
+                "The model must be built before the number of parameters can be estimated"
+            )
         else:
             if requires_grad:
                 return sum(p.numel() for p in self.task_model.parameters() if p.requires_grad)  # type: ignore
@@ -259,6 +297,8 @@ def fit(
         val_size: float = 0.2,
         X_val=None,
         y_val=None,
+        embeddings=None,
+        embeddings_val=None,
         max_epochs: int = 100,
         random_state: int = 101,
         batch_size: int = 128,
@@ -340,6 +380,8 @@ def fit(
                 val_size=val_size,
                 X_val=X_val,
                 y_val=y_val,
+                embeddings=embeddings,
+                embeddings_val=embeddings_val,
                 random_state=random_state,
                 batch_size=batch_size,
                 shuffle=shuffle,
@@ -390,7 +432,7 @@ def fit(
 
         return self
 
-    def predict(self, X, device=None):
+    def predict(self, X, embeddings=None, device=None):
         """Predicts target labels for the given input samples.
 
         Parameters
@@ -408,7 +450,7 @@ def predict(self, X, device=None):
             raise ValueError("The model or data module has not been fitted yet.")
 
         # Preprocess the data using the data module
-        self.data_module.assign_predict_dataset(X)
+        self.data_module.assign_predict_dataset(X, embeddings)
 
         # Set model to evaluation mode
         self.task_model.eval()
@@ -438,7 +480,7 @@ def predict(self, X, device=None):
         # Convert predictions to NumPy array and return
         return predictions.cpu().numpy()
 
-    def predict_proba(self, X, device=None):
+    def predict_proba(self, X, embeddings=None, device=None):
         """Predicts class probabilities for the given input samples.
 
         Parameters
@@ -482,7 +524,7 @@ def predict_proba(self, X, device=None):
         # Convert probabilities to NumPy array and return
         return probabilities.cpu().numpy()
 
-    def evaluate(self, X, y_true, metrics=None):
+    def evaluate(self, X, y_true, embeddings=None, metrics=None):
         """Evaluate the model on the given data using specified metrics.
 
         Parameters
@@ -491,6 +533,8 @@ def evaluate(self, X, y_true, metrics=None):
             The input samples to predict.
         y_true : array-like of shape (n_samples,)
             The true class labels against which to evaluate the predictions.
+        embneddings : array-like or list of shape(n_samples, dimension)
+            List or array with embeddings for unstructured data inputs
         metrics : dict
             A dictionary where keys are metric names and values are tuples containing the metric function
             and a boolean indicating whether the metric requires probability scores (True) or class labels (False).
@@ -518,11 +562,11 @@ def evaluate(self, X, y_true, metrics=None):
 
         # Generate class probabilities if any metric requires them
         if any(use_proba for _, use_proba in metrics.values()):
-            probabilities = self.predict_proba(X)
+            probabilities = self.predict_proba(X, embeddings)
 
         # Generate class labels if any metric requires them
         if any(not use_proba for _, use_proba in metrics.values()):
-            predictions = self.predict(X)
+            predictions = self.predict(X, embeddings)
 
         # Compute each metric
         for metric_name, (metric_func, use_proba) in metrics.items():
@@ -533,7 +577,7 @@ def evaluate(self, X, y_true, metrics=None):
 
         return scores
 
-    def score(self, X, y, metric=(log_loss, True)):
+    def score(self, X, y, embeddings=None, metric=(log_loss, True)):
         """Calculate the score of the model using the specified metric.
 
         Parameters
@@ -557,13 +601,13 @@ def score(self, X, y, metric=(log_loss, True)):
             X = pd.DataFrame(X)
 
         if use_proba:
-            probabilities = self.predict_proba(X)
+            probabilities = self.predict_proba(X, embeddings)
             return metric_func(y, probabilities)
         else:
-            predictions = self.predict(X)
+            predictions = self.predict(X, embeddings)
             return metric_func(y, predictions)
 
-    def encode(self, X, batch_size=64):
+    def encode(self, X, embeddings=None, batch_size=64):
         """
         Encodes input data using the trained model's embedding layer.
 
@@ -587,14 +631,16 @@ def encode(self, X, batch_size=64):
         # Ensure model and data module are initialized
         if self.task_model is None or self.data_module is None:
             raise ValueError("The model or data module has not been fitted yet.")
-        encoded_dataset = self.data_module.preprocess_new_data(X)
+        encoded_dataset = self.data_module.preprocess_new_data(X, embeddings)
 
         data_loader = DataLoader(encoded_dataset, batch_size=batch_size, shuffle=False)
 
         # Process data in batches
         encoded_outputs = []
-        for num_features, cat_features in tqdm(data_loader):
-            embeddings = self.task_model.base_model.encode(num_features, cat_features)  # Call your encode function
+        for batch in tqdm(data_loader):
+            embeddings = self.task_model.base_model.encode(
+                batch
+            )  # Call your encode function
             encoded_outputs.append(embeddings)
 
         # Concatenate all encoded outputs
@@ -608,6 +654,8 @@ def optimize_hparams(
         y,
         X_val=None,
         y_val=None,
+        embeddings=None,
+        embeddings_val=None,
         time=100,
         max_epochs=200,
         prune_by_epoch=True,
@@ -658,13 +706,25 @@ def optimize_hparams(
         )
 
         # Initial model fitting to get the baseline validation loss
-        self.fit(X, y, X_val=X_val, y_val=y_val, max_epochs=max_epochs)
+        self.fit(
+            X,
+            y,
+            X_val=X_val,
+            y_val=y_val,
+            embeddings=embeddings,
+            embeddings_val=embeddings_val,
+            max_epochs=max_epochs,
+        )
         best_val_loss = float("inf")
 
         if X_val is not None and y_val is not None:
-            val_loss = self.evaluate(X_val, y_val, metrics={"Accuracy": (accuracy_score, False)})["Accuracy"]
+            val_loss = self.evaluate(
+                X_val, y_val, metrics={"Accuracy": (accuracy_score, False)}
+            )["Accuracy"]
         else:
-            val_loss = self.trainer.validate(self.task_model, self.data_module)[0]["val_loss"]
+            val_loss = self.trainer.validate(self.task_model, self.data_module)[0][
+                "val_loss"
+            ]
 
         best_val_loss = val_loss
         best_epoch_val_loss = self.task_model.epoch_val_loss_at(  # type: ignore
@@ -690,7 +750,9 @@ def _objective(hyperparams):
                         if param_value in activation_mapper:
                             setattr(self.config, key, activation_mapper[param_value])
                         else:
-                            raise ValueError(f"Unknown activation function: {param_value}")
+                            raise ValueError(
+                                f"Unknown activation function: {param_value}"
+                            )
                     else:
                         setattr(self.config, key, param_value)
 
@@ -699,11 +761,15 @@ def _objective(hyperparams):
                 self.config.head_layer_sizes = head_layer_sizes[:head_layer_size_length]
 
             # Build the model with updated hyperparameters
-            self.build_model(X, y, X_val=X_val, y_val=y_val, lr=self.config.lr, **optimize_kwargs)
+            self.build_model(
+                X, y, X_val=X_val, y_val=y_val, lr=self.config.lr, **optimize_kwargs
+            )
 
             # Dynamically set the early pruning threshold
             if prune_by_epoch:
-                early_pruning_threshold = best_epoch_val_loss * 1.5  # Prune based on specific epoch loss
+                early_pruning_threshold = (
+                    best_epoch_val_loss * 1.5
+                )  # Prune based on specific epoch loss
             else:
                 # Prune based on the best overall validation loss
                 early_pruning_threshold = best_val_loss * 1.5
@@ -715,15 +781,26 @@ def _objective(hyperparams):
             # Fit the model (limit epochs for faster optimization)
             try:
                 # Wrap the risky operation (model fitting) in a try-except block
-                self.fit(X, y, X_val=X_val, y_val=y_val, max_epochs=max_epochs, rebuild=False)
+                self.fit(
+                    X,
+                    y,
+                    X_val=X_val,
+                    y_val=y_val,
+                    embeddings=embeddings,
+                    embeddings_val=embeddings_val,
+                    max_epochs=max_epochs,
+                    rebuild=False,
+                )
 
                 # Evaluate validation loss
                 if X_val is not None and y_val is not None:
-                    val_loss = self.evaluate(X_val, y_val, metrics={"Mean Squared Error": mean_squared_error})[  # type: ignore
+                    val_loss = self.evaluate(X_val, y_val, metrics={"Accuracy": (accuracy_score, False)})[  # type: ignore
                         "Mean Squared Error"
                     ]
                 else:
-                    val_loss = self.trainer.validate(self.task_model, self.data_module)[0]["val_loss"]
+                    val_loss = self.trainer.validate(self.task_model, self.data_module)[
+                        0
+                    ]["val_loss"]
 
                 # Pruning based on validation loss at specific epoch
                 epoch_val_loss = self.task_model.epoch_val_loss_at(  # type: ignore
@@ -740,15 +817,21 @@ def _objective(hyperparams):
 
             except Exception as e:
                 # Penalize the hyperparameter configuration with a large value
-                print(f"Error encountered during fit with hyperparameters {hyperparams}: {e}")
-                return best_val_loss * 100  # Large value to discourage this configuration
+                print(
+                    f"Error encountered during fit with hyperparameters {hyperparams}: {e}"
+                )
+                return (
+                    best_val_loss * 100
+                )  # Large value to discourage this configuration
 
         # Perform Bayesian optimization using scikit-optimize
         result = gp_minimize(_objective, param_space, n_calls=time, random_state=42)
 
         # Update the model with the best-found hyperparameters
         best_hparams = result.x  # type: ignore
-        head_layer_sizes = [] if "head_layer_sizes" in self.config.__dataclass_fields__ else None
+        head_layer_sizes = (
+            [] if "head_layer_sizes" in self.config.__dataclass_fields__ else None
+        )
         layer_sizes = [] if "layer_sizes" in self.config.__dataclass_fields__ else None
 
         # Iterate over the best hyperparameters found by optimization