dragon-ml-toolbox 12.13.0__py3-none-any.whl → 13.1.0__py3-none-any.whl

ml_tools/ML_models.py

@@ -8,6 +8,7 @@ from ._logger import _LOGGER
 from .path_manager import make_fullpath
 from ._script_info import _script_info
 from .keys import PytorchModelArchitectureKeys
+from ._schema import FeatureSchema
 
 
 __all__ = [
@@ -298,76 +299,59 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
     """
     A Transformer-based model for tabular data tasks.
     
-    This model uses a Feature Tokenizer to convert all input features into a sequence of embeddings, prepends a [CLS] token, and processes the
+    This model uses a Feature Tokenizer to convert all input features into a
+    sequence of embeddings, prepends a [CLS] token, and processes the
     sequence with a standard Transformer Encoder.
     """
     def __init__(self, *,
-                 in_features: int,
+                 schema: FeatureSchema,
                  out_targets: int,
-                 categorical_index_map: Dict[int, int],
                  embedding_dim: int = 32,
                  num_heads: int = 8,
                  num_layers: int = 6,
                  dropout: float = 0.1):
         """
         Args:
-            in_features (int): The total number of columns in the input data (features).
-            out_targets (int): Number of output targets (1 for regression).
-            categorical_index_map (Dict[int, int]): Maps categorical column index to its cardinality (number of unique categories).
-            embedding_dim (int): The dimension for all feature embeddings. Must be divisible by num_heads.
-            num_heads (int): The number of heads in the multi-head attention mechanism.
-            num_layers (int): The number of sub-encoder-layers in the transformer encoder.
-            dropout (float): The dropout value.
-            
-        Note: 
-        - All arguments are keyword-only to promote clarity.
-        - Column indices start at 0.
-        
-        ### Data Preparation
-        The model requires a specific input format. All columns in the input DataFrame must be numerical, but they are treated differently based on the 
-        provided index lists.
-
-        **Nominal Categorical Features** (e.g., 'City', 'Color'): Should **NOT** be one-hot encoded. 
-        Instead, convert them to integer codes (label encoding). You must then provide a dictionary mapping their column indices to 
-        their cardinality (the number of unique categories) via the `categorical_map` parameter.
-
-        **Ordinal & Binary Features** (e.g., 'Low/Medium/High', 'True/False'): Should be treated as **numerical**. Map them to numbers that 
-        represent their state (e.g., `{'Low': 0, 'Medium': 1}` or `{False: 0, True: 1}`). Their column indices should **NOT** be included in the 
-        `categorical_map` parameter.
-
-        **Standard Numerical and Continuous Features** (e.g., 'Age', 'Price'): It is highly recommended to scale them before training.
+            schema (FeatureSchema): 
+                The definitive schema object created by `data_exploration.finalize_feature_schema()`.
+            out_targets (int): 
+                Number of output targets (1 for regression).
+            embedding_dim (int): 
+                The dimension for all feature embeddings. Must be divisible 
+                by num_heads.
+            num_heads (int): 
+                The number of heads in the multi-head attention mechanism.
+            num_layers (int): 
+                The number of sub-encoder-layers in the transformer encoder.
+            dropout (float): 
+                The dropout value.
         """
         super().__init__()
         
+         # --- Get info from schema ---
+        in_features = len(schema.feature_names)
+        categorical_index_map = schema.categorical_index_map
+
          # --- Validation ---
-        if categorical_index_map and max(categorical_index_map.keys()) >= in_features:
+        if categorical_index_map and (max(categorical_index_map.keys()) >= in_features):
             _LOGGER.error(f"A categorical index ({max(categorical_index_map.keys())}) is out of bounds for the provided input features ({in_features}).")
             raise ValueError()
         
-        # --- Derive numerical indices ---
-        all_indices = set(range(in_features))
-        categorical_indices_set = set(categorical_index_map.keys())
-        numerical_indices = sorted(list(all_indices - categorical_indices_set))
-        
         # --- Save configuration ---
-        self.in_features = in_features
+        self.schema = schema # <-- Save the whole schema
         self.out_targets = out_targets
-        self.numerical_indices = numerical_indices
-        self.categorical_map = categorical_index_map
         self.embedding_dim = embedding_dim
         self.num_heads = num_heads
         self.num_layers = num_layers
         self.dropout = dropout
 
-        # --- 1. Feature Tokenizer ---
+        # --- 1. Feature Tokenizer (now takes the schema) ---
         self.tokenizer = _FeatureTokenizer(
-            numerical_indices=numerical_indices,
-            categorical_map=categorical_index_map,
+            schema=schema,
             embedding_dim=embedding_dim
         )
         
         # --- 2. CLS Token ---
-        # A learnable token that will be prepended to the sequence.
         self.cls_token = nn.Parameter(torch.randn(1, 1, embedding_dim))
         
         # --- 3. Transformer Encoder ---
@@ -416,21 +400,87 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
     
     def get_architecture_config(self) -> Dict[str, Any]:
         """Returns the full configuration of the model."""
+        # Deconstruct schema into a JSON-friendly dict
+        # Tuples are saved as lists
+        schema_dict = {
+            'feature_names': self.schema.feature_names,
+            'continuous_feature_names': self.schema.continuous_feature_names,
+            'categorical_feature_names': self.schema.categorical_feature_names,
+            'categorical_index_map': self.schema.categorical_index_map,
+            'categorical_mappings': self.schema.categorical_mappings
+        }
+
         return {
-            'in_features': self.in_features,
+            'schema_dict': schema_dict,
             'out_targets': self.out_targets,
-            'categorical_map': self.categorical_map,
             'embedding_dim': self.embedding_dim,
             'num_heads': self.num_heads,
             'num_layers': self.num_layers,
             'dropout': self.dropout
         }
+    
+    @classmethod
+    def load(cls: type, file_or_dir: Union[str, Path], verbose: bool = True) -> nn.Module:
+        """Loads a model architecture from a JSON file."""
+        user_path = make_fullpath(file_or_dir)
+        
+        if user_path.is_dir():
+            json_filename = PytorchModelArchitectureKeys.SAVENAME + ".json"
+            target_path = make_fullpath(user_path / json_filename, enforce="file")
+        elif user_path.is_file():
+            target_path = user_path
+        else:
+            _LOGGER.error(f"Invalid path: '{file_or_dir}'")
+            raise IOError()
+
+        with open(target_path, 'r') as f:
+            saved_data = json.load(f)
+
+        saved_class_name = saved_data[PytorchModelArchitectureKeys.MODEL]
+        config = saved_data[PytorchModelArchitectureKeys.CONFIG]
+
+        if saved_class_name != cls.__name__:
+            _LOGGER.error(f"Model class mismatch. File specifies '{saved_class_name}', but '{cls.__name__}' was expected.")
+            raise ValueError()
+
+        # --- RECONSTRUCTION LOGIC ---
+        if 'schema_dict' not in config:
+            _LOGGER.error("Invalid architecture file: missing 'schema_dict'. This file may be from an older version.")
+            raise ValueError("Missing 'schema_dict' in config.")
+            
+        schema_data = config.pop('schema_dict')
+        
+        # Re-hydrate the categorical_index_map
+        # JSON saves all dict keys as strings, so we must convert them back to int.
+        raw_index_map = schema_data['categorical_index_map']
+        if raw_index_map is not None:
+            rehydrated_index_map = {int(k): v for k, v in raw_index_map.items()}
+        else:
+            rehydrated_index_map = None
+
+        # Re-hydrate the FeatureSchema object
+        # JSON deserializes tuples as lists, so we must convert them back.
+        schema = FeatureSchema(
+            feature_names=tuple(schema_data['feature_names']),
+            continuous_feature_names=tuple(schema_data['continuous_feature_names']),
+            categorical_feature_names=tuple(schema_data['categorical_feature_names']),
+            categorical_index_map=rehydrated_index_map,
+            categorical_mappings=schema_data['categorical_mappings']
+        )
+        
+        config['schema'] = schema
+        # --- End Reconstruction ---
+
+        model = cls(**config)
+        if verbose:
+            _LOGGER.info(f"Successfully loaded architecture for '{saved_class_name}'")
+        return model
         
     def __repr__(self) -> str:
         """Returns the developer-friendly string representation of the model."""
         # Build the architecture string part-by-part
         parts = [
-            f"Tokenizer(features={self.in_features}, dim={self.embedding_dim})",
+            f"Tokenizer(features={len(self.schema.feature_names)}, dim={self.embedding_dim})",
             "[CLS]",
             f"TransformerEncoder(layers={self.num_layers}, heads={self.num_heads})",
             f"PredictionHead(outputs={self.out_targets})"
@@ -443,29 +493,41 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
 
 class _FeatureTokenizer(nn.Module):
     """
-    Transforms raw numerical and categorical features from any column order into a sequence of embeddings.
+    Transforms raw numerical and categorical features from any column order 
+    into a sequence of embeddings.
     """
     def __init__(self,
-                 numerical_indices: List[int],
-                 categorical_map: Dict[int, int],
+                 schema: FeatureSchema,
                  embedding_dim: int):
         """
         Args:
-            numerical_indices (List[int]): A list of column indices for the numerical features.
-            categorical_map (Dict[int, int]): A dictionary mapping each categorical column index to its cardinality (number of unique categories).
-            embedding_dim (int): The dimension for all feature embeddings.
+            schema (FeatureSchema): 
+                The definitive schema object from data_exploration.
+            embedding_dim (int): 
+                The dimension for all feature embeddings.
         """
         super().__init__()
         
-        # Unpack the dictionary into separate lists for indices and cardinalities
-        self.categorical_indices = list(categorical_map.keys())
-        cardinalities = list(categorical_map.values())
+        # --- Get info from schema ---
+        categorical_map = schema.categorical_index_map
+        
+        if categorical_map:
+            # Unpack the dictionary into separate lists
+            self.categorical_indices = list(categorical_map.keys())
+            cardinalities = list(categorical_map.values())
+        else:
+            self.categorical_indices = []
+            cardinalities = []
+        
+        # Derive numerical indices by finding what's not categorical
+        all_indices = set(range(len(schema.feature_names)))
+        categorical_indices_set = set(self.categorical_indices)
+        self.numerical_indices = sorted(list(all_indices - categorical_indices_set))
         
-        self.numerical_indices = numerical_indices
         self.embedding_dim = embedding_dim
         
         # A learnable embedding for each numerical feature
-        self.numerical_embeddings = nn.Parameter(torch.randn(len(numerical_indices), embedding_dim))
+        self.numerical_embeddings = nn.Parameter(torch.randn(len(self.numerical_indices), embedding_dim))
         
         # A standard embedding layer for each categorical feature
         self.categorical_embeddings = nn.ModuleList(
@@ -487,6 +549,8 @@ class _FeatureTokenizer(nn.Module):
         # Process categorical features
         categorical_tokens = []
         for i, embed_layer in enumerate(self.categorical_embeddings):
+            # x_categorical[:, i] selects the i-th categorical column
+            # (e.g., all values for the 'color' feature)
             token = embed_layer(x_categorical[:, i]).unsqueeze(1)
             categorical_tokens.append(token)
         

ml_tools/ML_optimization.py

@@ -17,9 +17,10 @@ from ._script_info import _script_info
 from .ML_inference import PyTorchInferenceHandler
 from .keys import PyTorchInferenceKeys
 from .SQL import DatabaseManager
-from .optimization_tools import _save_result
+from .optimization_tools import _save_result, create_optimization_bounds
 from .utilities import save_dataframe_filename
 from .math_utilities import discretize_categorical_values
+from ._schema import FeatureSchema
 
 
 __all__ = [
@@ -40,66 +41,76 @@ class MLOptimizer:
     SNES and CEM algorithms do not accept bounds, the given bounds will be used as an initial starting point.
 
     Example:
-        >>> # 1. Get categorical info from preprocessing steps
-        >>> # e.g., from data_exploration.encode_categorical_features
-        >>> cat_mappings = {'feature_C': {'A': 0, 'B': 1}, 'feature_D': {'X': 0, 'Y': 1}}
-        >>> # e.g., from data_exploration.create_transformer_categorical_map
-        >>> # Assumes feature_C is at index 2 (cardinality 2) and feature_D is at index 3 (cardinality 2)
-        >>> cat_index_map = {2: 2, 3: 2}
+        >>> # 1. Get the final schema from data exploration
+        >>> schema = data_exploration.finalize_feature_schema(...)
+        >>> # 2. Define bounds for continuous features
+        >>> cont_bounds = {'feature_A': (0, 100), 'feature_B': (-10, 10)}
         >>>
-        >>> # 2. Initialize the optimizer
+        >>> # 3. Initialize the optimizer
         >>> optimizer = MLOptimizer(
         ...     inference_handler=my_handler,
-        ...     bounds=(lower_bounds, upper_bounds), # Bounds for ALL features
+        ...     schema=schema,
+        ...     continuous_bounds_map=cont_bounds,
         ...     task="max",
         ...     algorithm="Genetic",
-        ...     categorical_index_map=cat_index_map,
-        ...     categorical_mappings=cat_mappings,
         ... )
-        >>> # 3. Run the optimization
+        >>> # 4. Run the optimization
         >>> best_result = optimizer.run(
         ...     num_generations=100,
         ...     target_name="my_target",
-        ...     feature_names=my_feature_names,
         ...     save_dir="/path/to/results",
         ...     save_format="csv"
         ... )
     """
     def __init__(self,
                  inference_handler: PyTorchInferenceHandler,
-                 bounds: Tuple[List[float], List[float]],
+                 schema: FeatureSchema,
+                 continuous_bounds_map: Dict[str, Tuple[float, float]],
                  task: Literal["min", "max"],
                  algorithm: Literal["SNES", "CEM", "Genetic"] = "Genetic",
                  population_size: int = 200,
-                 categorical_index_map: Optional[Dict[int, int]] = None,
-                 categorical_mappings: Optional[Dict[str, Dict[str, int]]] = None,
                  discretize_start_at_zero: bool = True,
                  **searcher_kwargs):
         """
         Initializes the optimizer by creating the EvoTorch problem and searcher.
 
         Args:
-            inference_handler (PyTorchInferenceHandler): An initialized inference handler containing the model and weights.
-            bounds (tuple[list[float], list[float]]): A tuple containing the lower and upper bounds for ALL solution features. 
-                Use the `optimization_tools.create_optimization_bounds()` helper to easily generate this and ensure unbiased categorical bounds.
+            inference_handler (PyTorchInferenceHandler): 
+                An initialized inference handler containing the model.
+            schema (FeatureSchema): 
+                The definitive schema object from data_exploration.
+            continuous_bounds_map (Dict[str, Tuple[float, float]]):
+                A dictionary mapping the *name* of each **continuous** feature
+                to its (min_bound, max_bound) tuple.
             task (str): The optimization goal, either "min" or "max".
             algorithm (str): The search algorithm to use ("SNES", "CEM", "Genetic").
             population_size (int): Population size for CEM and GeneticAlgorithm.
-            categorical_index_map (Dict[int, int] | None): Used to discretize values after optimization. Maps {column_index: cardinality}.
-            categorical_mappings (Dict[str, Dict[str, int]] | None): Used to map discrete integer values back to strings (e.g., {0: 'Category_A'}) before saving.
             discretize_start_at_zero (bool): 
                 True if the discrete encoding starts at 0 (e.g., [0, 1, 2]).
                 False if it starts at 1 (e.g., [1, 2, 3]).
-            **searcher_kwargs: Additional keyword arguments for the selected search algorithm's constructor.
+            **searcher_kwargs: Additional keyword arguments for the selected 
+                               search algorithm's constructor.
         """
-        # Make a fitness function
+        # --- Store schema ---
+        self.schema = schema
+        
+        # --- 1. Create bounds from schema ---
+        # This is the new, robust way to get bounds
+        bounds = create_optimization_bounds(
+            schema=schema,
+            continuous_bounds_map=continuous_bounds_map,
+            start_at_zero=discretize_start_at_zero
+        )
+
+        # --- 2. Make a fitness function ---
         self.evaluator = FitnessEvaluator(
             inference_handler=inference_handler,
-            categorical_index_map=categorical_index_map,
+            # Get categorical info from the schema
+            categorical_index_map=schema.categorical_index_map,
             discretize_start_at_zero=discretize_start_at_zero
         )
         
-        # Call the existing factory function to get the problem and searcher factory
+        # --- 3. Create the problem and searcher factory ---
         self.problem, self.searcher_factory = create_pytorch_problem(
             evaluator=self.evaluator,
             bounds=bounds,
@@ -108,36 +119,36 @@ class MLOptimizer:
             population_size=population_size,
             **searcher_kwargs
         )
-        # Store categorical info to pass to the run function
-        self.categorical_map = categorical_index_map
-        self.categorical_mappings = categorical_mappings
+        
+        # --- 4. Store other info needed by run() ---
         self.discretize_start_at_zero = discretize_start_at_zero
 
     def run(self,
             num_generations: int,
             target_name: str,
             save_dir: Union[str, Path],
-            feature_names: Optional[List[str]],
             save_format: Literal['csv', 'sqlite', 'both'],
             repetitions: int = 1,
             verbose: bool = True) -> Optional[dict]:
         """
         Runs the evolutionary optimization process using the pre-configured settings.
 
+        The `feature_names` are automatically pulled from the `FeatureSchema`
+        provided during initialization.
+
         Args:
             num_generations (int): The total number of generations for each repetition.
             target_name (str): Target name used for the CSV filename and/or SQL table.
             save_dir (str | Path): The directory where result files will be saved.
-            feature_names (List[str] | None): Names of the solution features for labeling output.
-                If None, generic names like 'feature_0', 'feature_1', ... , will be created.
             save_format (Literal['csv', 'sqlite', 'both']): The format for saving results.
             repetitions (int): The number of independent times to run the optimization.
             verbose (bool): If True, enables detailed logging.
 
         Returns:
-            Optional[dict]: A dictionary with the best result if repetitions is 1, otherwise None.
+            Optional[dict]: A dictionary with the best result if repetitions is 1, 
+                            otherwise None.
         """
-        # Call the existing run function with the stored problem, searcher, and categorical info
+        # Call the existing run function, passing info from the schema
         return run_optimization(
             problem=self.problem,
             searcher_factory=self.searcher_factory,
@@ -145,11 +156,13 @@ class MLOptimizer:
             target_name=target_name,
             save_dir=save_dir,
             save_format=save_format,
-            feature_names=feature_names,
+            # Get the definitive feature names (as a list) from the schema
+            feature_names=list(self.schema.feature_names),
+            # Get categorical info from the schema
+            categorical_map=self.schema.categorical_index_map,
+            categorical_mappings=self.schema.categorical_mappings,
             repetitions=repetitions,
             verbose=verbose,
-            categorical_map=self.categorical_map,
-            categorical_mappings=self.categorical_mappings,
             discretize_start_at_zero=self.discretize_start_at_zero
         )
         

ml_tools/ML_trainer.py

@@ -5,12 +5,13 @@ import torch
 from torch import nn
 import numpy as np
 
-from .ML_callbacks import Callback, History, TqdmProgressBar
+from .ML_callbacks import Callback, History, TqdmProgressBar, ModelCheckpoint
 from .ML_evaluation import classification_metrics, regression_metrics, plot_losses, shap_summary_plot, plot_attention_importance
 from .ML_evaluation_multi import multi_target_regression_metrics, multi_label_classification_metrics, multi_target_shap_summary_plot
 from ._script_info import _script_info
-from .keys import PyTorchLogKeys
+from .keys import PyTorchLogKeys, PyTorchCheckpointKeys
 from ._logger import _LOGGER
+from .path_manager import make_fullpath
 
 
 __all__ = [
@@ -55,6 +56,7 @@ class MLTrainer:
         self.kind = kind
         self.criterion = criterion
         self.optimizer = optimizer
+        self.scheduler = None
         self.device = self._validate_device(device)
         self.dataloader_workers = dataloader_workers
         
@@ -70,6 +72,7 @@ class MLTrainer:
         self.history = {}
         self.epoch = 0
         self.epochs = 0 # Total epochs for the fit run
+        self.start_epoch = 1
         self.stop_training = False
 
     def _validate_device(self, device: str) -> torch.device:
@@ -109,8 +112,66 @@ class MLTrainer:
             num_workers=loader_workers, 
             pin_memory=("cuda" in self.device.type)
         )
+        
+    def _load_checkpoint(self, path: Union[str, Path]):
+        """Loads a training checkpoint to resume training."""
+        p = make_fullpath(path, enforce="file")
+        _LOGGER.info(f"Loading checkpoint from '{p.name}' to resume training...")
+        
+        try:
+            checkpoint = torch.load(p, map_location=self.device)
+            
+            if PyTorchCheckpointKeys.MODEL_STATE not in checkpoint or PyTorchCheckpointKeys.OPTIMIZER_STATE not in checkpoint:
+                _LOGGER.error(f"Checkpoint file '{p.name}' is invalid. Missing 'model_state_dict' or 'optimizer_state_dict'.")
+                raise KeyError()
 
-    def fit(self, epochs: int = 10, batch_size: int = 10, shuffle: bool = True):
+            self.model.load_state_dict(checkpoint[PyTorchCheckpointKeys.MODEL_STATE])
+            self.optimizer.load_state_dict(checkpoint[PyTorchCheckpointKeys.OPTIMIZER_STATE])
+            self.start_epoch = checkpoint.get(PyTorchCheckpointKeys.EPOCH, 0) + 1 # Resume on the *next* epoch
+            
+            # --- Scheduler State Loading Logic ---
+            scheduler_state_exists = PyTorchCheckpointKeys.SCHEDULER_STATE in checkpoint
+            scheduler_object_exists = self.scheduler is not None
+
+            if scheduler_object_exists and scheduler_state_exists:
+                # Case 1: Both exist. Attempt to load.
+                try:
+                    self.scheduler.load_state_dict(checkpoint[PyTorchCheckpointKeys.SCHEDULER_STATE]) # type: ignore
+                    scheduler_name = self.scheduler.__class__.__name__
+                    _LOGGER.info(f"Restored LR scheduler state for: {scheduler_name}")
+                except Exception as e:
+                    # Loading failed, likely a mismatch
+                    scheduler_name = self.scheduler.__class__.__name__
+                    _LOGGER.error(f"Failed to load scheduler state for '{scheduler_name}'. A different scheduler type might have been used.")
+                    raise e
+
+            elif scheduler_object_exists and not scheduler_state_exists:
+                # Case 2: Scheduler provided, but no state in checkpoint.
+                scheduler_name = self.scheduler.__class__.__name__
+                _LOGGER.warning(f"'{scheduler_name}' was provided, but no scheduler state was found in the checkpoint. The scheduler will start from its initial state.")
+            
+            elif not scheduler_object_exists and scheduler_state_exists:
+                # Case 3: State in checkpoint, but no scheduler provided.
+                _LOGGER.error("Checkpoint contains an LR scheduler state, but no LRScheduler callback was provided.")
+                raise ValueError()
+            
+            # Restore callback states
+            for cb in self.callbacks:
+                if isinstance(cb, ModelCheckpoint) and PyTorchCheckpointKeys.BEST_SCORE in checkpoint:
+                    cb.best = checkpoint[PyTorchCheckpointKeys.BEST_SCORE]
+                    _LOGGER.info(f"Restored {cb.__class__.__name__} 'best' score to: {cb.best:.4f}")
+            
+            _LOGGER.info(f"Checkpoint loaded. Resuming training from epoch {self.start_epoch}.")
+            
+        except Exception as e:
+            _LOGGER.error(f"Failed to load checkpoint from '{p}': {e}")
+            raise
+
+    def fit(self, 
+            epochs: int = 10, 
+            batch_size: int = 10, 
+            shuffle: bool = True,
+            resume_from_checkpoint: Optional[Union[str, Path]] = None):
         """
         Starts the training-validation process of the model.
         
@@ -120,6 +181,7 @@ class MLTrainer:
             epochs (int): The total number of epochs to train for.
             batch_size (int): The number of samples per batch.
             shuffle (bool): Whether to shuffle the training data at each epoch.
+            resume_from_checkpoint (str | Path | None): Optional path to a checkpoint to resume training.
             
         Note:
             For regression tasks using `nn.MSELoss` or `nn.L1Loss`, the trainer
@@ -132,15 +194,18 @@ class MLTrainer:
         self._create_dataloaders(batch_size, shuffle)
         self.model.to(self.device)
         
+        if resume_from_checkpoint:
+            self._load_checkpoint(resume_from_checkpoint)
+        
         # Reset stop_training flag on the trainer
         self.stop_training = False
 
-        self.callbacks_hook('on_train_begin')
+        self._callbacks_hook('on_train_begin')
         
-        for epoch in range(1, self.epochs + 1):
+        for epoch in range(self.start_epoch, self.epochs + 1):
             self.epoch = epoch
             epoch_logs = {}
-            self.callbacks_hook('on_epoch_begin', epoch, logs=epoch_logs)
+            self._callbacks_hook('on_epoch_begin', epoch, logs=epoch_logs)
 
             train_logs = self._train_step()
             epoch_logs.update(train_logs)
@@ -148,13 +213,13 @@ class MLTrainer:
             val_logs = self._validation_step()
             epoch_logs.update(val_logs)
             
-            self.callbacks_hook('on_epoch_end', epoch, logs=epoch_logs)
+            self._callbacks_hook('on_epoch_end', epoch, logs=epoch_logs)
             
             # Check the early stopping flag
             if self.stop_training:
                 break
 
-        self.callbacks_hook('on_train_end')
+        self._callbacks_hook('on_train_end')
         return self.history
     
     def _train_step(self):
@@ -166,7 +231,7 @@ class MLTrainer:
                 PyTorchLogKeys.BATCH_INDEX: batch_idx, 
                 PyTorchLogKeys.BATCH_SIZE: features.size(0)
             }
-            self.callbacks_hook('on_batch_begin', batch_idx, logs=batch_logs)
+            self._callbacks_hook('on_batch_begin', batch_idx, logs=batch_logs)
 
             features, target = features.to(self.device), target.to(self.device)
             self.optimizer.zero_grad()
@@ -188,7 +253,7 @@ class MLTrainer:
             
             # Add the batch loss to the logs and call the end-of-batch hook
             batch_logs[PyTorchLogKeys.BATCH_LOSS] = batch_loss
-            self.callbacks_hook('on_batch_end', batch_idx, logs=batch_logs)
+            self._callbacks_hook('on_batch_end', batch_idx, logs=batch_logs)
 
         return {PyTorchLogKeys.TRAIN_LOSS: running_loss / len(self.train_loader.dataset)} # type: ignore
 
@@ -538,11 +603,33 @@ class MLTrainer:
         else:
             _LOGGER.error("No attention weights were collected from the model.")
     
-    def callbacks_hook(self, method_name: str, *args, **kwargs):
+    def _callbacks_hook(self, method_name: str, *args, **kwargs):
         """Calls the specified method on all callbacks."""
         for callback in self.callbacks:
             method = getattr(callback, method_name)
             method(*args, **kwargs)
+            
+    def to_cpu(self):
+        """
+        Moves the model to the CPU and updates the trainer's device setting.
+        
+        This is useful for running operations that require the CPU.
+        """
+        self.device = torch.device('cpu')
+        self.model.to(self.device)
+        _LOGGER.info("Trainer and model moved to CPU.")
+    
+    def to_device(self, device: str):
+        """
+        Moves the model to the specified device and updates the trainer's device setting.
+
+        Args:
+            device (str): The target device (e.g., 'cuda', 'mps', 'cpu').
+        """
+        self.device = self._validate_device(device)
+        self.model.to(self.device)
+        _LOGGER.info(f"Trainer and model moved to {self.device}.")
+    
 
 def info():
     _script_info(__all__)

ml_tools/PSO_optimization.py

@@ -17,6 +17,10 @@ from ._script_info import _script_info
 from .SQL import DatabaseManager
 from .optimization_tools import _save_result
 
+"""
+DEPRECATED
+"""
+
 
 __all__ = [
     "ObjectiveFunction",
@@ -46,7 +50,7 @@ class ObjectiveFunction():
         self.binary_features = binary_features
         self.is_hybrid = False if binary_features <= 0 else True
         self.use_noise = add_noise
-        self._artifact = deserialize_object(trained_model_path, verbose=False, raise_on_error=True)
+        self._artifact = deserialize_object(trained_model_path, verbose=False)
         self.model = self._get_from_artifact(EnsembleKeys.MODEL)
         self.feature_names: Optional[list[str]] = self._get_from_artifact(EnsembleKeys.FEATURES) # type: ignore
         self.target_name: Optional[str] = self._get_from_artifact(EnsembleKeys.TARGET) # type: ignore