rename to self.estimator

Author: AnFreTh

mambular/base_models/utils/pretraining.py

@@ -20,8 +20,8 @@ def __init__(
         pool_sequence=True,
     ):
         super().__init__()
-        self.base_model = base_model
-        self.base_model.eval()
+        self.estimator = base_model
+        self.estimator.eval()
         self.k_neighbors = k_neighbors
         self.temperature = temperature
         self.lr = lr
@@ -33,9 +33,9 @@ def __init__(
         self.loss_fn = nn.CosineEmbeddingLoss(margin=margin, reduction="mean")
 
     def forward(self, x):
-        x = self.base_model.encode(x, grad=True)
+        x = self.estimator.encode(x, grad=True)
         if self.pool_sequence:
-            return self.base_model.pool_sequence(x)
+            return self.estimator.pool_sequence(x)
         return x  # Return unpooled sequence embeddings (N, S, D)
 
     def get_knn(self, labels):
@@ -140,7 +140,7 @@ def contrastive_loss(self, embeddings, knn_indices, neg_indices):
 
     def training_step(self, batch, batch_idx):
 
-        self.base_model.embedding_layer.train()
+        self.estimator.embedding_layer.train()
 
         data, labels = batch
         embeddings = self(data)
@@ -173,7 +173,7 @@ def validation_step(self, batch, batch_idx):
         return loss
 
     def configure_optimizers(self):
-        params = chain(self.base_model.parameters())
+        params = chain(self.estimator.parameters())
         return torch.optim.Adam(params, lr=self.lr)
 
 

mambular/models/utils/sklearn_base_classifier.py

@@ -248,7 +248,7 @@ def predict(self, X, embeddings=None, device=None):
         logits = torch.cat(logits_list, dim=0)  # type: ignore
 
         # Check if ensemble is used
-        if getattr(self.base_model, "returns_ensemble", False):  # If using ensemble
+        if getattr(self.estimator, "returns_ensemble", False):  # If using ensemble
             logits = logits.mean(dim=1)  # Average over ensemble dimension
             if logits.dim() == 1:  # Ensure correct shape
                 logits = logits.unsqueeze(1)
@@ -296,7 +296,7 @@ def predict_proba(self, X, embeddings=None, device=None):
         logits = torch.cat(logits_list, dim=0)
 
         # Check if ensemble is used
-        if getattr(self.base_model, "returns_ensemble", False):  # If using ensemble
+        if getattr(self.estimator, "returns_ensemble", False):  # If using ensemble
             logits = logits.mean(dim=1)  # Average over ensemble dimension
             if logits.dim() == 1:  # Ensure correct shape
                 logits = logits.unsqueeze(1)
@@ -439,7 +439,7 @@ def pretrain(
         Notes
         -----
         - This function requires that `self.build_model()` has been called beforehand.
-        - The pretraining method uses `self.task_model.base_model.embedding_layer`.
+        - The pretraining method uses `self.task_model.estimator.embedding_layer`.
         - The method invokes `super()._pretrain()` with regression mode enabled.
 
         """
@@ -448,13 +448,13 @@ def pretrain(
                 "The model has not been built yet. Call model.build_model(**args) first."
             )
 
-        if not hasattr(self.task_model.base_model, "embedding_layer"):
+        if not hasattr(self.task_model.estimator, "embedding_layer"):
             raise ValueError("The model does not have an embedding layer")
 
         self.data_module.setup("fit")
 
         super()._pretrain(
-            self.task_model.base_model,
+            self.task_model.estimator,
             self.data_module,
             pretrain_epochs=pretrain_epochs,
             k_neighbors=k_neighbors,

mambular/models/utils/sklearn_base_lss.py

@@ -9,18 +9,13 @@
 from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint, ModelSummary
 from sklearn.base import BaseEstimator
 from sklearn.metrics import accuracy_score, mean_squared_error
-from skopt import gp_minimize
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
 from ...base_models.utils.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.distributional_metrics import (
     beta_brier_score,
     dirichlet_error,
@@ -78,7 +73,7 @@ def __init__(self, model, config, **kwargs):
 
         self.preprocessor = Preprocessor(**preprocessor_kwargs)
         self.task_model = None
-        self.base_model = model
+        self.estimator = model
         self.built = False
 
         # Raise a warning if task is set to 'classification'
@@ -246,7 +241,7 @@ def build_model(
         )
 
         self.task_model = TaskModel(
-            model_class=self.base_model,  # type: ignore
+            model_class=self.estimator,  # type: ignore
             num_classes=self.family.param_count,
             family=self.family,
             config=self.config,
@@ -268,7 +263,7 @@ def build_model(
         )
 
         self.built = True
-        self.base_model = self.task_model.base_model
+        self.estimator = self.task_model.estimator
 
         return self
 
@@ -497,7 +492,7 @@ def predict(self, X, raw=False, device=None):
         predictions = torch.cat(predictions_list, dim=0)
 
         # Check if ensemble is used
-        if getattr(self.base_model, "returns_ensemble", False):  # If using ensemble
+        if getattr(self.estimator, "returns_ensemble", False):  # If using ensemble
             predictions = predictions.mean(dim=1)  # Average over ensemble dimension
 
         if not raw:
@@ -642,7 +637,7 @@ def encode(self, X, batch_size=64):
         # Process data in batches
         encoded_outputs = []
         for num_features, cat_features in tqdm(data_loader):
-            embeddings = self.task_model.base_model.encode(
+            embeddings = self.task_model.estimator.encode(
                 num_features, cat_features
             )  # Call your encode function
             encoded_outputs.append(embeddings)

mambular/models/utils/sklearn_base_regressor.py

@@ -249,7 +249,7 @@ def predict(self, X, embeddings=None, device=None):
 
         # Check if ensemble is used
         if getattr(
-            self.task_model.base_model, "returns_ensemble", False
+            self.task_model.estimator, "returns_ensemble", False
         ):  # If using ensemble
             predictions = predictions.mean(dim=1)  # Average over ensemble dimension
 
@@ -360,7 +360,7 @@ def pretrain(
         Notes
         -----
         - This function requires that `self.build_model()` has been called beforehand.
-        - The pretraining method uses `self.task_model.base_model.embedding_layer`.
+        - The pretraining method uses `self.task_model.estimator.embedding_layer`.
         - The method invokes `super()._pretrain()` with regression mode enabled.
 
         """
@@ -369,13 +369,13 @@ def pretrain(
                 "The model has not been built yet. Call model.build_model(**args) first."
             )
 
-        if not hasattr(self.task_model.base_model, "embedding_layer"):
+        if not hasattr(self.task_model.estimator, "embedding_layer"):
             raise ValueError("The model does not have an embedding layer")
 
         self.data_module.setup("fit")
 
         super()._pretrain(
-            self.task_model.base_model,
+            self.task_model.estimator,
             self.data_module,
             pretrain_epochs=pretrain_epochs,
             k_neighbors=k_neighbors,

mambular/models/utils/sklearn_parent.py

@@ -54,7 +54,7 @@ def __init__(self, model, config, **kwargs):
         }
 
         self.preprocessor = Preprocessor(**self.preprocessor_kwargs)
-        self.base_model = model
+        self.estimator = model
         self.task_model = None
         self.built = False
 
@@ -208,7 +208,7 @@ def _build_model(
         )
 
         self.task_model = TaskModel(
-            model_class=self.base_model,  # type: ignore
+            model_class=self.estimator,  # type: ignore
             config=self.config,
             feature_information=(
                 self.data_module.num_feature_info,
@@ -230,7 +230,7 @@ def _build_model(
         )
 
         self.built = True
-        self.base_model = self.task_model.base_model
+        self.estimator = self.task_model.estimator
 
         return self
 
@@ -399,7 +399,7 @@ def fit(
             **trainer_kwargs,
         )
         self.task_model.train()
-        self.task_model.base_model.train()
+        self.task_model.estimator.train()
         self.trainer.fit(self.task_model, self.data_module)  # type: ignore
 
         self.best_model_path = checkpoint_callback.best_model_path
@@ -458,7 +458,7 @@ def encode(self, X, embeddings=None, batch_size=64):
         # Process data in batches
         encoded_outputs = []
         for batch in tqdm(data_loader):
-            embeddings = self.task_model.base_model.encode(
+            embeddings = self.task_model.estimator.encode(
                 batch
             )  # Call your encode function
             encoded_outputs.append(embeddings)