renaming sklearn class attributes

Author: AnFreTh

mambular/models/sklearn_base_classifier.py

@@ -37,7 +37,7 @@ def __init__(self, model, config, **kwargs):
         }
 
         self.preprocessor = Preprocessor(**preprocessor_kwargs)
-        self.model = None
+        self.task_model = None
 
         # Raise a warning if task is set to 'classification'
         if preprocessor_kwargs.get("task") == "regression":
@@ -194,7 +194,7 @@ def build_model(
 
         num_classes = len(np.unique(y))
 
-        self.model = TaskModel(
+        self.task_model = TaskModel(
             model_class=self.base_model,
             num_classes=num_classes,
             config=self.config,
@@ -237,10 +237,10 @@ def get_number_of_params(self, requires_grad=True):
         else:
             if requires_grad:
                 return sum(
-                    p.numel() for p in self.model.parameters() if p.requires_grad
+                    p.numel() for p in self.task_model.parameters() if p.requires_grad
                 )
             else:
-                return sum(p.numel() for p in self.model.parameters())
+                return sum(p.numel() for p in self.task_model.parameters())
 
     def fit(
         self,
@@ -345,7 +345,7 @@ def fit(
 
             num_classes = len(np.unique(y))
 
-            self.model = TaskModel(
+            self.task_model = TaskModel(
                 model_class=self.base_model,
                 num_classes=num_classes,
                 config=self.config,
@@ -379,12 +379,12 @@ def fit(
             ],
             **trainer_kwargs
         )
-        self.trainer.fit(self.model, self.data_module)
+        self.trainer.fit(self.task_model, self.data_module)
 
         best_model_path = checkpoint_callback.best_model_path
         if best_model_path:
             checkpoint = torch.load(best_model_path)
-            self.model.load_state_dict(checkpoint["state_dict"])
+            self.task_model.load_state_dict(checkpoint["state_dict"])
 
         return self
 
@@ -404,14 +404,14 @@ def predict(self, X):
             The predicted target values.
         """
         # Ensure model and data module are initialized
-        if self.model is None or self.data_module is None:
+        if self.task_model is None or self.data_module is None:
             raise ValueError("The model or data module has not been fitted yet.")
 
         # Preprocess the data using the data module
         cat_tensors, num_tensors = self.data_module.preprocess_test_data(X)
 
         # Move tensors to appropriate device
-        device = next(self.model.parameters()).device
+        device = next(self.task_model.parameters()).device
         if isinstance(cat_tensors, list):
             cat_tensors = [tensor.to(device) for tensor in cat_tensors]
         else:
@@ -423,11 +423,11 @@ def predict(self, X):
             num_tensors = num_tensors.to(device)
 
         # Set model to evaluation mode
-        self.model.eval()
+        self.task_model.eval()
 
         # Perform inference
         with torch.no_grad():
-            logits = self.model(num_features=num_tensors, cat_features=cat_tensors)
+            logits = self.task_model(num_features=num_tensors, cat_features=cat_tensors)
 
             # Check the shape of the logits to determine binary or multi-class classification
             if logits.shape[1] == 1:
@@ -484,7 +484,7 @@ def predict_proba(self, X):
         # Preprocess the data
         if not isinstance(X, pd.DataFrame):
             X = pd.DataFrame(X)
-        device = next(self.model.parameters()).device
+        device = next(self.task_model.parameters()).device
         cat_tensors, num_tensors = self.data_module.preprocess_test_data(X)
         if isinstance(cat_tensors, list):
             cat_tensors = [tensor.to(device) for tensor in cat_tensors]
@@ -497,11 +497,11 @@ def predict_proba(self, X):
             num_tensors = num_tensors.to(device)
 
         # Set the model to evaluation mode
-        self.model.eval()
+        self.task_model.eval()
 
         # Perform inference
         with torch.no_grad():
-            logits = self.model(num_features=num_tensors, cat_features=cat_tensors)
+            logits = self.task_model(num_features=num_tensors, cat_features=cat_tensors)
             if logits.shape[1] > 1:
                 probabilities = torch.softmax(logits, dim=1)
             else:

mambular/models/sklearn_base_lss.py

@@ -58,7 +58,7 @@ def __init__(self, model, config, **kwargs):
         }
 
         self.preprocessor = Preprocessor(**preprocessor_kwargs)
-        self.model = None
+        self.task_model = None
 
         # Raise a warning if task is set to 'classification'
         if preprocessor_kwargs.get("task") == "classification":
@@ -212,7 +212,7 @@ def build_model(
 
         num_classes = len(np.unique(y))
 
-        self.model = TaskModel(
+        self.task_model = TaskModel(
             model_class=self.base_model,
             num_classes=num_classes,
             config=self.config,
@@ -255,10 +255,10 @@ def get_number_of_params(self, requires_grad=True):
         else:
             if requires_grad:
                 return sum(
-                    p.numel() for p in self.model.parameters() if p.requires_grad
+                    p.numel() for p in self.task_model.parameters() if p.requires_grad
                 )
             else:
-                return sum(p.numel() for p in self.model.parameters())
+                return sum(p.numel() for p in self.task_model.parameters())
 
     def fit(
         self,
@@ -383,7 +383,7 @@ def fit(
             X, y, X_val, y_val, val_size=val_size, random_state=random_state
         )
 
-        self.model = TaskModel(
+        self.task_model = TaskModel(
             model_class=self.base_model,
             num_classes=self.family.param_count,
             family=self.family,
@@ -419,12 +419,12 @@ def fit(
             ],
             **trainer_kwargs
         )
-        self.trainer.fit(self.model, self.data_module)
+        self.trainer.fit(self.task_model, self.data_module)
 
         best_model_path = checkpoint_callback.best_model_path
         if best_model_path:
             checkpoint = torch.load(best_model_path)
-            self.model.load_state_dict(checkpoint["state_dict"])
+            self.task_model.load_state_dict(checkpoint["state_dict"])
 
         return self
 
@@ -444,14 +444,14 @@ def predict(self, X, raw=False):
             The predicted target values.
         """
         # Ensure model and data module are initialized
-        if self.model is None or self.data_module is None:
+        if self.task_model is None or self.data_module is None:
             raise ValueError("The model or data module has not been fitted yet.")
 
         # Preprocess the data using the data module
         cat_tensors, num_tensors = self.data_module.preprocess_test_data(X)
 
         # Move tensors to appropriate device
-        device = next(self.model.parameters()).device
+        device = next(self.task_model.parameters()).device
         if isinstance(cat_tensors, list):
             cat_tensors = [tensor.to(device) for tensor in cat_tensors]
         else:
@@ -463,14 +463,14 @@ def predict(self, X, raw=False):
             num_tensors = num_tensors.to(device)
 
         # Set model to evaluation mode
-        self.model.eval()
+        self.task_model.eval()
 
         # Perform inference
         with torch.no_grad():
-            predictions = self.model(num_features=num_tensors, cat_features=cat_tensors)
+            predictions = self.task_model(num_features=num_tensors, cat_features=cat_tensors)
 
         if not raw:
-            return self.model.family(predictions).cpu().numpy()
+            return self.task_model.family(predictions).cpu().numpy()
 
         # Convert predictions to NumPy array and return
         else:
@@ -506,7 +506,7 @@ def evaluate(self, X, y_true, metrics=None, distribution_family=None):
         """
         # Infer distribution family from model settings if not provided
         if distribution_family is None:
-            distribution_family = getattr(self.model, "distribution_family", "normal")
+            distribution_family = getattr(self.task_model, "distribution_family", "normal")
 
         # Setup default metrics if none are provided
         if metrics is None:

mambular/models/sklearn_base_regressor.py

@@ -37,7 +37,7 @@ def __init__(self, model, config, **kwargs):
 
         self.preprocessor = Preprocessor(**preprocessor_kwargs)
         self.base_model = model
-        self.model = None
+        self.task_model = None
         self.built = False
 
         # Raise a warning if task is set to 'classification'
@@ -190,7 +190,7 @@ def build_model(
             X, y, X_val, y_val, val_size=val_size, random_state=random_state
         )
 
-        self.model = TaskModel(
+        self.task_model = TaskModel(
             model_class=self.base_model,
             config=self.config,
             cat_feature_info=self.data_module.cat_feature_info,
@@ -232,10 +232,10 @@ def get_number_of_params(self, requires_grad=True):
         else:
             if requires_grad:
                 return sum(
-                    p.numel() for p in self.model.parameters() if p.requires_grad
+                    p.numel() for p in self.task_model.parameters() if p.requires_grad
                 )
             else:
-                return sum(p.numel() for p in self.model.parameters())
+                return sum(p.numel() for p in self.task_model.parameters())
 
     def fit(
         self,
@@ -336,7 +336,7 @@ def fit(
                 X, y, X_val, y_val, val_size=val_size, random_state=random_state
             )
 
-            self.model = TaskModel(
+            self.task_model = TaskModel(
                 model_class=self.base_model,
                 config=self.config,
                 cat_feature_info=self.data_module.cat_feature_info,
@@ -372,12 +372,12 @@ def fit(
             ],
             **trainer_kwargs
         )
-        self.trainer.fit(self.model, self.data_module)
+        self.trainer.fit(self.task_model, self.data_module)
 
         best_model_path = checkpoint_callback.best_model_path
         if best_model_path:
             checkpoint = torch.load(best_model_path)
-            self.model.load_state_dict(checkpoint["state_dict"])
+            self.task_model.load_state_dict(checkpoint["state_dict"])
 
         return self
 
@@ -397,14 +397,14 @@ def predict(self, X):
             The predicted target values.
         """
         # Ensure model and data module are initialized
-        if self.model is None or self.data_module is None:
+        if self.task_model is None or self.data_module is None:
             raise ValueError("The model or data module has not been fitted yet.")
 
         # Preprocess the data using the data module
         cat_tensors, num_tensors = self.data_module.preprocess_test_data(X)
 
         # Move tensors to appropriate device
-        device = next(self.model.parameters()).device
+        device = next(self.task_model.parameters()).device
         if isinstance(cat_tensors, list):
             cat_tensors = [tensor.to(device) for tensor in cat_tensors]
         else:
@@ -416,11 +416,13 @@ def predict(self, X):
             num_tensors = num_tensors.to(device)
 
         # Set model to evaluation mode
-        self.model.eval()
+        self.task_model.eval()
 
         # Perform inference
         with torch.no_grad():
-            predictions = self.model(num_features=num_tensors, cat_features=cat_tensors)
+            predictions = self.task_model(
+                num_features=num_tensors, cat_features=cat_tensors
+            )
 
         # Convert predictions to NumPy array and return
         return predictions.cpu().numpy()