dragon-ml-toolbox 19.14.0__py3-none-any.whl → 20.0.0__py3-none-any.whl

ml_tools/ML_configuration/_models.py

@@ -0,0 +1,206 @@
+from typing import Optional, Literal
+
+from ..schema import FeatureSchema
+
+from ._base_model_config import _BaseModelParams
+
+
+__all__ = [    
+    # --- Model Parameter Configs ---
+    "DragonMLPParams",
+    "DragonAttentionMLPParams",
+    "DragonMultiHeadAttentionNetParams",
+    "DragonTabularTransformerParams",
+    "DragonGateParams",
+    "DragonNodeParams",
+    "DragonTabNetParams",
+    "DragonAutoIntParams",
+]
+
+
+# ----------------------------
+# Model Parameters Configurations
+# ----------------------------
+
+# --- Standard Models ---
+
+class DragonMLPParams(_BaseModelParams):
+    def __init__(self, 
+                 in_features: int, 
+                 out_targets: int,
+                 hidden_layers: list[int], 
+                 drop_out: float = 0.2) -> None:
+        self.in_features = in_features
+        self.out_targets = out_targets
+        self.hidden_layers = hidden_layers
+        self.drop_out = drop_out
+
+
+class DragonAttentionMLPParams(_BaseModelParams):
+    def __init__(self, 
+                 in_features: int, 
+                 out_targets: int,
+                 hidden_layers: list[int], 
+                 drop_out: float = 0.2) -> None:
+        self.in_features = in_features
+        self.out_targets = out_targets
+        self.hidden_layers = hidden_layers
+        self.drop_out = drop_out
+
+
+class DragonMultiHeadAttentionNetParams(_BaseModelParams):
+    def __init__(self, 
+                 in_features: int, 
+                 out_targets: int,
+                 hidden_layers: list[int], 
+                 drop_out: float = 0.2,
+                 num_heads: int = 4, 
+                 attention_dropout: float = 0.1) -> None:
+        self.in_features = in_features
+        self.out_targets = out_targets
+        self.hidden_layers = hidden_layers
+        self.drop_out = drop_out
+        self.num_heads = num_heads
+        self.attention_dropout = attention_dropout
+
+
+class DragonTabularTransformerParams(_BaseModelParams):
+    def __init__(self, *,
+                 schema: FeatureSchema,
+                 out_targets: int,
+                 embedding_dim: int = 256,
+                 num_heads: int = 8,
+                 num_layers: int = 6,
+                 dropout: float = 0.2) -> None:
+        self.schema = schema
+        self.out_targets = out_targets
+        self.embedding_dim = embedding_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.dropout = dropout
+
+# --- Advanced Models ---
+
+class DragonGateParams(_BaseModelParams):
+    def __init__(self, *,
+                 schema: FeatureSchema,
+                 out_targets: int,
+                 embedding_dim: int = 16,
+                 gflu_stages: int = 6,
+                 gflu_dropout: float = 0.1,
+                 num_trees: int = 20,
+                 tree_depth: int = 4,
+                 tree_dropout: float = 0.1,
+                 chain_trees: bool = False,
+                 tree_wise_attention: bool = True,
+                 tree_wise_attention_dropout: float = 0.1,
+                 binning_activation: Literal['entmoid', 'sparsemoid', 'sigmoid'] = "entmoid",
+                 feature_mask_function: Literal['entmax', 'sparsemax', 'softmax', 't-softmax'] = "entmax",
+                 share_head_weights: bool = True,
+                 batch_norm_continuous: bool = True) -> None:
+        self.schema = schema
+        self.out_targets = out_targets
+        self.embedding_dim = embedding_dim
+        self.gflu_stages = gflu_stages
+        self.gflu_dropout = gflu_dropout
+        self.num_trees = num_trees
+        self.tree_depth = tree_depth
+        self.tree_dropout = tree_dropout
+        self.chain_trees = chain_trees
+        self.tree_wise_attention = tree_wise_attention
+        self.tree_wise_attention_dropout = tree_wise_attention_dropout
+        self.binning_activation = binning_activation
+        self.feature_mask_function = feature_mask_function
+        self.share_head_weights = share_head_weights
+        self.batch_norm_continuous = batch_norm_continuous
+
+
+class DragonNodeParams(_BaseModelParams):
+    def __init__(self, *,
+                 schema: FeatureSchema,
+                 out_targets: int,
+                 embedding_dim: int = 24,
+                 num_trees: int = 1024,
+                 num_layers: int = 2,
+                 tree_depth: int = 6,
+                 additional_tree_output_dim: int = 3,
+                 max_features: Optional[int] = None,
+                 input_dropout: float = 0.0,
+                 embedding_dropout: float = 0.0,
+                 choice_function: Literal['entmax', 'sparsemax', 'softmax'] = 'entmax',
+                 bin_function: Literal['entmoid', 'sparsemoid', 'sigmoid'] = 'entmoid',
+                 batch_norm_continuous: bool = False) -> None:
+        self.schema = schema
+        self.out_targets = out_targets
+        self.embedding_dim = embedding_dim
+        self.num_trees = num_trees
+        self.num_layers = num_layers
+        self.tree_depth = tree_depth
+        self.additional_tree_output_dim = additional_tree_output_dim
+        self.max_features = max_features
+        self.input_dropout = input_dropout
+        self.embedding_dropout = embedding_dropout
+        self.choice_function = choice_function
+        self.bin_function = bin_function
+        self.batch_norm_continuous = batch_norm_continuous
+
+
+class DragonAutoIntParams(_BaseModelParams):
+    def __init__(self, *,
+                 schema: FeatureSchema,
+                 out_targets: int,
+                 embedding_dim: int = 32,
+                 attn_embed_dim: int = 32,
+                 num_heads: int = 2,
+                 num_attn_blocks: int = 3,
+                 attn_dropout: float = 0.1,
+                 has_residuals: bool = True,
+                 attention_pooling: bool = True,
+                 deep_layers: bool = True,
+                 layers: str = "128-64-32",
+                 activation: str = "ReLU",
+                 embedding_dropout: float = 0.0,
+                 batch_norm_continuous: bool = False) -> None:
+        self.schema = schema
+        self.out_targets = out_targets
+        self.embedding_dim = embedding_dim
+        self.attn_embed_dim = attn_embed_dim
+        self.num_heads = num_heads
+        self.num_attn_blocks = num_attn_blocks
+        self.attn_dropout = attn_dropout
+        self.has_residuals = has_residuals
+        self.attention_pooling = attention_pooling
+        self.deep_layers = deep_layers
+        self.layers = layers
+        self.activation = activation
+        self.embedding_dropout = embedding_dropout
+        self.batch_norm_continuous = batch_norm_continuous
+
+
+class DragonTabNetParams(_BaseModelParams):
+    def __init__(self, *,
+                 schema: FeatureSchema,
+                 out_targets: int,
+                 n_d: int = 8,
+                 n_a: int = 8,
+                 n_steps: int = 3,
+                 gamma: float = 1.3,
+                 n_independent: int = 2,
+                 n_shared: int = 2,
+                 virtual_batch_size: int = 128,
+                 momentum: float = 0.02,
+                 mask_type: Literal['sparsemax', 'entmax', 'softmax'] = 'sparsemax',
+                 batch_norm_continuous: bool = False) -> None:
+        self.schema = schema
+        self.out_targets = out_targets
+        self.n_d = n_d
+        self.n_a = n_a
+        self.n_steps = n_steps
+        self.gamma = gamma
+        self.n_independent = n_independent
+        self.n_shared = n_shared
+        self.virtual_batch_size = virtual_batch_size
+        self.momentum = momentum
+        self.mask_type = mask_type
+        self.batch_norm_continuous = batch_norm_continuous
+

ml_tools/ML_configuration/_training.py

@@ -0,0 +1,124 @@
+from typing import Union, Optional, Any, Literal
+from pathlib import Path
+
+from .._core import get_logger
+from ..path_manager import make_fullpath
+
+from ._base_model_config import _BaseModelParams
+
+
+_LOGGER = get_logger("ML Configuration")
+
+
+__all__ = [    
+    # --- Training Config ---
+    "DragonTrainingConfig",
+    "DragonParetoConfig"
+]
+
+
+class DragonTrainingConfig(_BaseModelParams):
+    """
+    Configuration object for the training process.
+    
+    Can be unpacked as a dictionary for logging or accessed as an object.
+    
+    Accepts arbitrary keyword arguments which are set as instance attributes.
+    """
+    def __init__(self,
+                 validation_size: float,
+                 test_size: float,
+                 initial_learning_rate: float,
+                 batch_size: int,
+                 random_state: int = 101,
+                 **kwargs: Any) -> None:
+        """  
+        Args:
+            validation_size (float): Proportion of data for validation set.
+            test_size (float): Proportion of data for test set.
+            initial_learning_rate (float): Starting learning rate.
+            batch_size (int): Number of samples per training batch.
+            random_state (int): Seed for reproducibility.
+            **kwargs: Additional training parameters as key-value pairs.
+        """
+        self.validation_size = validation_size
+        self.test_size = test_size
+        self.initial_learning_rate = initial_learning_rate
+        self.batch_size = batch_size
+        self.random_state = random_state
+        
+        # Process kwargs with validation
+        for key, value in kwargs.items():
+            # Python guarantees 'key' is a string for **kwargs
+            
+            # Allow None in value
+            if value is None:
+                setattr(self, key, value)
+                continue
+            
+            if isinstance(value, dict):
+                _LOGGER.error("Nested dictionaries are not supported, unpack them first.")
+                raise TypeError()
+            
+            # Check if value is a number or a string or a JSON supported type, except dict
+            if not isinstance(value, (str, int, float, bool, list, tuple)):
+                _LOGGER.error(f"Invalid type for configuration '{key}': {type(value).__name__}")
+                raise TypeError()
+            
+            setattr(self, key, value)
+
+
+class DragonParetoConfig(_BaseModelParams):
+    """
+    Configuration object for the Pareto Optimization process.
+    """
+    def __init__(self,
+                 save_directory: Union[str, Path],
+                 target_objectives: dict[str, Literal["min", "max"]],
+                 continuous_bounds_map: Union[dict[str, tuple[float, float]], dict[str, list[float]], str, Path],
+                 columns_to_round: Optional[list[str]] = None,
+                 population_size: int = 500,
+                 generations: int = 1000,
+                 solutions_filename: str = "NonDominatedSolutions",
+                 float_precision: int = 4,
+                 log_interval: int = 10,
+                 plot_size: tuple[int, int] = (10, 7),
+                 plot_font_size: int = 16,
+                 discretize_start_at_zero: bool = True):
+        """  
+        Configure the Pareto Optimizer.
+
+        Args:
+            save_directory (str | Path): Directory to save artifacts.
+            target_objectives (Dict[str, "min"|"max"]): Dictionary mapping target names to optimization direction.
+                Example: {"price": "max", "error": "min"}
+            continuous_bounds_map (Dict): Bounds for continuous features {name: (min, max)}. Or a path/str to a directory containing the "optimization_bounds.json" file.
+            columns_to_round (List[str] | None): List of continuous column names that should be rounded to the nearest integer.
+            population_size (int): Size of the genetic population.
+            generations (int): Number of generations to run.
+            solutions_filename (str): Filename for saving Pareto solutions.
+            float_precision (int): Number of decimal places to round standard float columns.
+            log_interval (int): Interval for logging progress.
+            plot_size (Tuple[int, int]): Size of the 2D plots.
+            plot_font_size (int): Font size for plot text.
+            discretize_start_at_zero (bool): Categorical encoding start index. True=0, False=1.
+        """
+        # Validate string or Path
+        valid_save_dir = make_fullpath(save_directory, make=True, enforce="directory")
+        
+        if isinstance(continuous_bounds_map, (str, Path)):
+            continuous_bounds_map = make_fullpath(continuous_bounds_map, make=False, enforce="directory")
+        
+        self.save_directory = valid_save_dir
+        self.target_objectives = target_objectives
+        self.continuous_bounds_map = continuous_bounds_map
+        self.columns_to_round = columns_to_round
+        self.population_size = population_size
+        self.generations = generations
+        self.solutions_filename = solutions_filename
+        self.float_precision = float_precision
+        self.log_interval = log_interval
+        self.plot_size = plot_size
+        self.plot_font_size = plot_font_size
+        self.discretize_start_at_zero = discretize_start_at_zero
+

ml_tools/ML_datasetmaster/__init__.py

@@ -0,0 +1,28 @@
+from ._datasetmaster import (
+    DragonDataset,
+    DragonDatasetMulti,
+)
+
+from ._sequence_datasetmaster import (
+    DragonDatasetSequence
+)
+
+from ._vision_datasetmaster import (
+    DragonDatasetVision,
+    DragonDatasetSegmentation,
+    DragonDatasetObjectDetection
+)
+
+from ._imprimir import info
+
+
+__all__ = [
+    "DragonDataset",
+    "DragonDatasetMulti",
+    # sequence
+    "DragonDatasetSequence",
+    # vision
+    "DragonDatasetVision",
+    "DragonDatasetSegmentation",
+    "DragonDatasetObjectDetection",
+]

ml_tools/ML_datasetmaster/_base_datasetmaster.py

@@ -0,0 +1,337 @@
+import torch
+from torch.utils.data import Dataset
+import pandas
+import numpy
+from typing import Union, Optional
+from abc import ABC
+from pathlib import Path
+
+from ..IO_tools import save_list_strings, custom_logger
+from ..ML_scaler import DragonScaler
+from ..schema import FeatureSchema
+
+from ..path_manager import make_fullpath, sanitize_filename
+from .._core import get_logger
+from ..keys._keys import DatasetKeys, ScalerKeys
+
+
+_LOGGER = get_logger("DragonDataset")
+
+
+__all__ = [
+    "_BaseDatasetMaker",
+    "_PytorchDataset",
+]
+
+
+# --- Internal Helper Class ---
+class _PytorchDataset(Dataset):
+    """
+    Internal helper class to create a PyTorch Dataset.
+    Converts numpy/pandas data into tensors for model consumption.
+    """
+    def __init__(self, features: Union[numpy.ndarray, pandas.DataFrame], 
+                 labels: Union[numpy.ndarray, pandas.Series, pandas.DataFrame],
+                 labels_dtype: torch.dtype,
+                 features_dtype: torch.dtype = torch.float32,
+                 feature_names: Optional[list[str]] = None,
+                 target_names: Optional[list[str]] = None):
+        
+        if isinstance(features, numpy.ndarray):
+            self.features = torch.tensor(features, dtype=features_dtype)
+        else: # It's a pandas.DataFrame
+            self.features = torch.tensor(features.to_numpy(), dtype=features_dtype)
+
+        if isinstance(labels, numpy.ndarray):
+            self.labels = torch.tensor(labels, dtype=labels_dtype)
+        elif isinstance(labels, (pandas.Series, pandas.DataFrame)):
+            self.labels = torch.tensor(labels.to_numpy(), dtype=labels_dtype)
+        else:
+            self.labels = torch.tensor(labels, dtype=labels_dtype)
+            
+        self._feature_names = feature_names
+        self._target_names = target_names
+        self._classes: list[str] = []
+        self._class_map: dict[str,int] = dict()
+        self._feature_scaler: Optional[DragonScaler] = None
+        self._target_scaler: Optional[DragonScaler] = None
+        
+    def __len__(self):
+        return len(self.features)
+
+    def __getitem__(self, index):
+        return self.features[index], self.labels[index]
+    
+    @property
+    def feature_names(self):
+        if self._feature_names is not None:
+            return self._feature_names
+        else:
+            _LOGGER.error(f"Dataset {self.__class__} has not been initialized with any feature names.")
+            raise ValueError()
+        
+    @property
+    def target_names(self):
+        if self._target_names is not None:
+            return self._target_names
+        else:
+            _LOGGER.error(f"Dataset {self.__class__} has not been initialized with any target names.")
+            raise ValueError()
+
+    @property
+    def classes(self):
+        return self._classes
+    
+    @property
+    def class_map(self):
+        return self._class_map
+    
+    @property
+    def feature_scaler(self):
+        return self._feature_scaler
+    
+    @property
+    def target_scaler(self):
+        return self._target_scaler
+
+
+# --- Abstract Base Class ---
+class _BaseDatasetMaker(ABC):
+    """
+    Abstract base class for dataset makers. Contains shared logic.
+    """
+    def __init__(self):
+        self._train_ds: Optional[Dataset] = None
+        self._val_ds: Optional[Dataset] = None
+        self._test_ds: Optional[Dataset] = None
+        
+        self.feature_scaler: Optional[DragonScaler] = None
+        self.target_scaler: Optional[DragonScaler] = None
+        
+        self._id: Optional[str] = None
+        self._feature_names: list[str] = []
+        self._target_names: list[str] = []
+        self._X_train_shape = (0,0)
+        self._X_val_shape = (0,0)
+        self._X_test_shape = (0,0)
+        self._y_train_shape = (0,)
+        self._y_val_shape = (0,)
+        self._y_test_shape = (0,)
+        self.class_map: dict[str, int] = dict()
+        self.classes: list[str] = list()
+        
+    def _prepare_feature_scaler(self, 
+                        X_train: pandas.DataFrame, 
+                        y_train: Union[pandas.Series, pandas.DataFrame], 
+                        X_val: pandas.DataFrame,
+                        X_test: pandas.DataFrame, 
+                        label_dtype: torch.dtype, 
+                        schema: FeatureSchema) -> tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
+        """Internal helper to fit and apply a DragonScaler for FEATURES using a FeatureSchema."""
+        continuous_feature_indices: Optional[list[int]] = None
+
+        # Get continuous feature indices *from the schema*
+        if schema.continuous_feature_names:
+            _LOGGER.info("Getting continuous feature indices from schema.")
+            try:
+                # Convert columns to a standard list for .index()
+                train_cols_list = X_train.columns.to_list()
+                # Map names from schema to column indices in the training DataFrame
+                continuous_feature_indices = [train_cols_list.index(name) for name in schema.continuous_feature_names]
+            except ValueError as e: 
+                _LOGGER.error(f"Feature name from schema not found in training data columns:\n{e}")
+                raise ValueError()
+        else:
+            _LOGGER.info("No continuous features listed in schema. Feature scaler will not be fitted.")
+
+        X_train_values = X_train.to_numpy()
+        X_val_values = X_val.to_numpy()
+        X_test_values = X_test.to_numpy()
+
+        # continuous_feature_indices is derived
+        if self.feature_scaler is None and continuous_feature_indices:
+            _LOGGER.info("Fitting a new DragonScaler on training features.")
+            temp_train_ds = _PytorchDataset(X_train_values, y_train, label_dtype) 
+            self.feature_scaler = DragonScaler.fit(temp_train_ds, continuous_feature_indices)
+
+        if self.feature_scaler and self.feature_scaler.mean_ is not None:
+            _LOGGER.info("Applying scaler transformation to train, validation, and test feature sets.")
+            X_train_tensor = self.feature_scaler.transform(torch.tensor(X_train_values, dtype=torch.float32))
+            X_val_tensor = self.feature_scaler.transform(torch.tensor(X_val_values, dtype=torch.float32))
+            X_test_tensor = self.feature_scaler.transform(torch.tensor(X_test_values, dtype=torch.float32))
+            return X_train_tensor.numpy(), X_val_tensor.numpy(), X_test_tensor.numpy()
+
+        return X_train_values, X_val_values, X_test_values
+    
+    def _prepare_target_scaler(self,
+                               y_train: Union[pandas.Series, pandas.DataFrame],
+                               y_val: Union[pandas.Series, pandas.DataFrame],
+                               y_test: Union[pandas.Series, pandas.DataFrame]) -> tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
+        """Internal helper to fit and apply a DragonScaler for TARGETS."""
+        
+        y_train_arr = y_train.to_numpy() if isinstance(y_train, (pandas.Series, pandas.DataFrame)) else y_train
+        y_val_arr = y_val.to_numpy() if isinstance(y_val, (pandas.Series, pandas.DataFrame)) else y_val
+        y_test_arr = y_test.to_numpy() if isinstance(y_test, (pandas.Series, pandas.DataFrame)) else y_test
+
+        if self.target_scaler is None:
+            _LOGGER.info("Fitting a new DragonScaler on training targets.")
+            # Convert to float tensor for calculation
+            y_train_tensor = torch.tensor(y_train_arr, dtype=torch.float32)
+            self.target_scaler = DragonScaler.fit_tensor(y_train_tensor)
+            
+        if self.target_scaler and self.target_scaler.mean_ is not None:
+             _LOGGER.info("Applying scaler transformation to train, validation, and test targets.")
+             y_train_tensor = self.target_scaler.transform(torch.tensor(y_train_arr, dtype=torch.float32))
+             y_val_tensor = self.target_scaler.transform(torch.tensor(y_val_arr, dtype=torch.float32))
+             y_test_tensor = self.target_scaler.transform(torch.tensor(y_test_arr, dtype=torch.float32))
+             return y_train_tensor.numpy(), y_val_tensor.numpy(), y_test_tensor.numpy()
+
+        return y_train_arr, y_val_arr, y_test_arr
+    
+    def _attach_scalers_to_datasets(self):
+        """Helper to attach the master scalers to the child datasets."""
+        for ds in [self._train_ds, self._val_ds, self._test_ds]:
+            if ds is not None:
+                ds._feature_scaler = self.feature_scaler # type: ignore
+                ds._target_scaler = self.target_scaler # type: ignore
+
+    @property
+    def train_dataset(self) -> Dataset:
+        if self._train_ds is None: 
+            _LOGGER.error("Train Dataset not yet created.")
+            raise RuntimeError()
+        return self._train_ds
+    
+    @property
+    def validation_dataset(self) -> Dataset:
+        if self._val_ds is None: 
+            _LOGGER.error("Validation Dataset not yet created.")
+            raise RuntimeError()
+        return self._val_ds
+
+    @property
+    def test_dataset(self) -> Dataset:
+        if self._test_ds is None: 
+            _LOGGER.error("Test Dataset not yet created.")
+            raise RuntimeError()
+        return self._test_ds
+
+    @property
+    def feature_names(self) -> list[str]:
+        return self._feature_names
+    
+    @property
+    def target_names(self) -> list[str]:
+        return self._target_names
+    
+    @property
+    def number_of_features(self) -> int:
+        return len(self._feature_names)
+    
+    @property
+    def number_of_targets(self) -> int:
+        return len(self._target_names)
+
+    @property
+    def id(self) -> Optional[str]:
+        return self._id
+
+    @id.setter
+    def id(self, dataset_id: str):
+        if not isinstance(dataset_id, str): raise ValueError("ID must be a string.")
+        self._id = dataset_id
+
+    def dataframes_info(self) -> None:
+        print("--- DataFrame Shapes After Split ---")
+        print(f"  X_train shape: {self._X_train_shape}, y_train shape: {self._y_train_shape}")
+        print(f"  X_val shape:   {self._X_val_shape}, y_val shape:   {self._y_val_shape}")
+        print(f"  X_test shape:  {self._X_test_shape}, y_test shape:  {self._y_test_shape}")
+        print("------------------------------------")
+    
+    def save_feature_names(self, directory: Union[str, Path], verbose: bool=True) -> None:
+        save_list_strings(list_strings=self._feature_names,
+                          directory=directory,
+                          filename=DatasetKeys.FEATURE_NAMES,
+                          verbose=verbose)
+        
+    def save_target_names(self, directory: Union[str, Path], verbose: bool=True) -> None:
+        save_list_strings(list_strings=self._target_names,
+                          directory=directory,
+                          filename=DatasetKeys.TARGET_NAMES,
+                          verbose=verbose)
+
+    def save_scaler(self, directory: Union[str, Path], verbose: bool=True) -> None:
+        """
+        Saves both feature and target scalers (if they exist) to a single .pth file
+        using a dictionary structure.
+        """
+        if self.feature_scaler is None and self.target_scaler is None:
+            _LOGGER.warning("No scalers (feature or target) were fitted. Nothing to save.")
+            return
+
+        if not self.id: 
+            _LOGGER.error("Must set the dataset `id` before saving scaler.")
+            raise ValueError()
+        
+        save_path = make_fullpath(directory, make=True, enforce="directory")
+        sanitized_id = sanitize_filename(self.id)
+        filename = f"{DatasetKeys.SCALER_PREFIX}{sanitized_id}.pth"
+        filepath = save_path / filename
+        
+        # Construct the consolidated dictionary
+        combined_state = {}
+        
+        print_message = "Saved "
+        
+        if self.feature_scaler:
+            combined_state[ScalerKeys.FEATURE_SCALER] = self.feature_scaler._get_state()
+            print_message += "feature scaler "
+            
+        if self.target_scaler:
+            if self.feature_scaler:
+                print_message += "and "
+            combined_state[ScalerKeys.TARGET_SCALER] = self.target_scaler._get_state()
+            print_message += "target scaler "
+            
+        torch.save(combined_state, filepath)
+        
+        if verbose:
+            _LOGGER.info(f"{print_message}to '{filepath.name}'.")
+            
+    def save_class_map(self, directory: Union[str,Path], verbose: bool=True) -> None:
+        """  
+        Saves the class map dictionary to a JSON file.
+        
+        Args:
+            directory (str | Path): Directory to save the class map.
+            verbose (bool): Whether to print log messages.
+        """
+        if not self.class_map:
+            _LOGGER.warning(f"No class_map defined. Skipping.")
+            return
+        
+        log_name = f"Class_to_Index_{self.id}" if self.id else "Class_to_Index"
+        
+        custom_logger(data=self.class_map,
+                      save_directory=directory,
+                      log_name=log_name,
+                      add_timestamp=False,
+                      dict_as="json")
+        if verbose:
+            _LOGGER.info(f"Class map for '{self.id}' saved as '{log_name}.json'.")
+
+    def save_artifacts(self, directory: Union[str, Path], verbose: bool=True) -> None:
+        """
+        Saves all dataset artifacts: feature names, target names, scalers, and class map (if applicable).
+        
+        Args:
+            directory (str | Path): Directory to save artifacts.
+            verbose (bool): Whether to print log messages.
+        """
+        self.save_feature_names(directory=directory, verbose=verbose)
+        self.save_target_names(directory=directory, verbose=verbose)
+        if self.feature_scaler is not None or self.target_scaler is not None:
+            self.save_scaler(directory=directory, verbose=verbose)
+        if self.class_map:
+            self.save_class_map(directory=directory, verbose=verbose)
+