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

ml_tools/ML_evaluation.py

@@ -19,11 +19,12 @@ import torch
 import shap
 from pathlib import Path
 from typing import Union, Optional, List, Literal
+import warnings
 
 from .path_manager import make_fullpath
 from ._logger import _LOGGER
 from ._script_info import _script_info
-from .keys import SHAPKeys
+from .keys import SHAPKeys, PyTorchLogKeys
 
 
 __all__ = [
@@ -43,8 +44,8 @@ def plot_losses(history: dict, save_dir: Union[str, Path]):
         history (dict): A dictionary containing 'train_loss' and 'val_loss'.
         save_dir (str | Path): Directory to save the plot image.
     """
-    train_loss = history.get('train_loss', [])
-    val_loss = history.get('val_loss', [])
+    train_loss = history.get(PyTorchLogKeys.TRAIN_LOSS, [])
+    val_loss = history.get(PyTorchLogKeys.VAL_LOSS, [])
     
     if not train_loss and not val_loss:
         print("Warning: Loss history is empty or incomplete. Cannot plot.")
@@ -257,7 +258,7 @@ def shap_summary_plot(model,
                       feature_names: Optional[list[str]], 
                       save_dir: Union[str, Path],
                       device: torch.device = torch.device('cpu'),
-                      explainer_type: Literal['deep', 'kernel'] = 'deep'):
+                      explainer_type: Literal['deep', 'kernel'] = 'kernel'):
     """
     Calculates SHAP values and saves summary plots and data.
 
@@ -269,7 +270,7 @@ def shap_summary_plot(model,
         save_dir (str | Path): Directory to save SHAP artifacts.
         device (torch.device): The torch device for SHAP calculations.
         explainer_type (Literal['deep', 'kernel']): The explainer to use.
-            - 'deep': (Default) Uses shap.DeepExplainer. Fast and efficient for
+            - 'deep': Uses shap.DeepExplainer. Fast and efficient for
               PyTorch models.
             - 'kernel': Uses shap.KernelExplainer. Model-agnostic but EXTREMELY
               slow and memory-intensive.
@@ -284,7 +285,7 @@ def shap_summary_plot(model,
     instances_to_explain_np = None
 
     if explainer_type == 'deep':
-        # --- 1. Use DeepExplainer (Preferred) ---
+        # --- 1. Use DeepExplainer  ---
         
         # Ensure data is torch.Tensor
         if isinstance(background_data, np.ndarray):
@@ -298,17 +299,19 @@ def shap_summary_plot(model,
 
         background_data = background_data.to(device)
         instances_to_explain = instances_to_explain.to(device)
-
-        explainer = shap.DeepExplainer(model, background_data)
+        
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=UserWarning)
+            explainer = shap.DeepExplainer(model, background_data)
+            
         # print("Calculating SHAP values with DeepExplainer...")
         shap_values = explainer.shap_values(instances_to_explain)
         instances_to_explain_np = instances_to_explain.cpu().numpy()
 
     elif explainer_type == 'kernel':
-        # --- 2. Use KernelExplainer (Slow Fallback) ---
+        # --- 2. Use KernelExplainer ---
         _LOGGER.warning(
-            "Using KernelExplainer. This is memory-intensive and slow. "
-            "Consider reducing 'n_samples' if the process terminates unexpectedly."
+            "KernelExplainer is memory-intensive and slow. Consider reducing the number of instances to explain if the process terminates unexpectedly."
         )
 
         # Ensure data is np.ndarray
@@ -344,14 +347,26 @@ def shap_summary_plot(model,
     else:
         _LOGGER.error(f"Invalid explainer_type: '{explainer_type}'. Must be 'deep' or 'kernel'.")
         raise ValueError()
+    
+    if not isinstance(shap_values, list) and shap_values.ndim == 3 and shap_values.shape[2] == 1:
+        # _LOGGER.info("Squeezing SHAP values from (N, F, 1) to (N, F) for regression plot.")
+        shap_values = shap_values.squeeze(-1)
 
     # --- 3. Plotting and Saving ---
     save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
     plt.ioff()
     
+    # Convert instances to a DataFrame. robust way to ensure SHAP correctly maps values to feature names.
+    if feature_names is None:
+        # Create generic names if none were provided
+        num_features = instances_to_explain_np.shape[1]
+        feature_names = [f'feature_{i}' for i in range(num_features)]
+        
+    instances_df = pd.DataFrame(instances_to_explain_np, columns=feature_names)
+    
     # Save Bar Plot
     bar_path = save_dir_path / "shap_bar_plot.svg"
-    shap.summary_plot(shap_values, instances_to_explain_np, feature_names=feature_names, plot_type="bar", show=False)
+    shap.summary_plot(shap_values, instances_df, plot_type="bar", show=False)
     ax = plt.gca()
     ax.set_xlabel("SHAP Value Impact", labelpad=10)
     plt.title("SHAP Feature Importance")
@@ -362,7 +377,7 @@ def shap_summary_plot(model,
 
     # Save Dot Plot
     dot_path = save_dir_path / "shap_dot_plot.svg"
-    shap.summary_plot(shap_values, instances_to_explain_np, feature_names=feature_names, plot_type="dot", show=False)
+    shap.summary_plot(shap_values, instances_df, plot_type="dot", show=False)
     ax = plt.gca()
     ax.set_xlabel("SHAP Value Impact", labelpad=10)
     if plt.gcf().axes and len(plt.gcf().axes) > 1:
@@ -385,9 +400,6 @@ def shap_summary_plot(model,
         mean_abs_shap = np.abs(shap_values).mean(axis=0)
 
     mean_abs_shap = mean_abs_shap.flatten()
-    
-    if feature_names is None:
-        feature_names = [f'feature_{i}' for i in range(len(mean_abs_shap))]
         
     summary_df = pd.DataFrame({
         SHAPKeys.FEATURE_COLUMN: feature_names,
@@ -397,7 +409,7 @@ def shap_summary_plot(model,
     summary_df.to_csv(summary_path, index=False)
     
     _LOGGER.info(f"📝 SHAP summary data saved as '{summary_path.name}'")
-    plt.ion()    
+    plt.ion()
 
 
 def plot_attention_importance(weights: List[torch.Tensor], feature_names: Optional[List[str]], save_dir: Union[str, Path], top_n: int = 10):

ml_tools/ML_evaluation_multi.py

@@ -20,6 +20,7 @@ from sklearn.metrics import (
 )
 from pathlib import Path
 from typing import Union, List, Literal
+import warnings
 
 from .path_manager import make_fullpath, sanitize_filename
 from ._logger import _LOGGER
@@ -234,7 +235,7 @@ def multi_target_shap_summary_plot(
     target_names: List[str],
     save_dir: Union[str, Path],
     device: torch.device = torch.device('cpu'),
-    explainer_type: Literal['deep', 'kernel'] = 'deep'
+    explainer_type: Literal['deep', 'kernel'] = 'kernel'
 ):
     """
     Calculates SHAP values for a multi-target model and saves summary plots and data for each target.
@@ -248,7 +249,7 @@ def multi_target_shap_summary_plot(
         save_dir (str | Path): Directory to save SHAP artifacts.
         device (torch.device): The torch device for SHAP calculations.
         explainer_type (Literal['deep', 'kernel']): The explainer to use.
-            - 'deep': (Default) Uses shap.DeepExplainer. Fast and efficient.
+            - 'deep': Uses shap.DeepExplainer. Fast and efficient.
             - 'kernel': Uses shap.KernelExplainer. Model-agnostic but slow and memory-intensive.
     """
     _LOGGER.info(f"--- Multi-Target SHAP Value Explanation (Using: {explainer_type.upper()}Explainer) ---")
@@ -259,7 +260,7 @@ def multi_target_shap_summary_plot(
     instances_to_explain_np = None
 
     if explainer_type == 'deep':
-        # --- 1. Use DeepExplainer (Preferred) ---
+        # --- 1. Use DeepExplainer ---
         
         # Ensure data is torch.Tensor
         if isinstance(background_data, np.ndarray):
@@ -273,18 +274,20 @@ def multi_target_shap_summary_plot(
 
         background_data = background_data.to(device)
         instances_to_explain = instances_to_explain.to(device)
-
-        explainer = shap.DeepExplainer(model, background_data)
-        print("Calculating SHAP values with DeepExplainer...")
+        
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=UserWarning)
+            explainer = shap.DeepExplainer(model, background_data)
+            
+        # print("Calculating SHAP values with DeepExplainer...")
         # DeepExplainer returns a list of arrays for multi-output models
         shap_values_list = explainer.shap_values(instances_to_explain)
         instances_to_explain_np = instances_to_explain.cpu().numpy()
 
     elif explainer_type == 'kernel':
-        # --- 2. Use KernelExplainer (Slow Fallback) ---
+        # --- 2. Use KernelExplainer  ---
         _LOGGER.warning(
-            "Using KernelExplainer. This is memory-intensive and slow. "
-            "Consider reducing 'n_samples' if the process terminates."
+            "KernelExplainer is memory-intensive and slow. Consider reducing the number of instances to explain if the process terminates unexpectedly."
         )
         
         # Convert all data to numpy
@@ -304,7 +307,7 @@ def multi_target_shap_summary_plot(
             return output.cpu().numpy() # Return full multi-output array
 
         explainer = shap.KernelExplainer(prediction_wrapper, background_summary)
-        print("Calculating SHAP values with KernelExplainer...")
+        # print("Calculating SHAP values with KernelExplainer...")
         # KernelExplainer also returns a list of arrays for multi-output models
         shap_values_list = explainer.shap_values(instances_to_explain_np, l1_reg="aic")
         # instances_to_explain_np is already set

ml_tools/ML_inference.py

@@ -9,7 +9,7 @@ from .ML_scaler import PytorchScaler
 from ._script_info import _script_info
 from ._logger import _LOGGER
 from .path_manager import make_fullpath
-from .keys import PyTorchInferenceKeys
+from .keys import PyTorchInferenceKeys, PyTorchCheckpointKeys
 
 
 __all__ = [
@@ -56,11 +56,21 @@ class _BaseInferenceHandler(ABC):
         model_p = make_fullpath(state_dict, enforce="file")
 
         try:
-            # Load the state dictionary and apply it to the model structure
-            self.model.load_state_dict(torch.load(model_p, map_location=self.device))
+            # Load whatever is in the file
+            loaded_data = torch.load(model_p, map_location=self.device)
+
+            # Check if it's the new checkpoint dictionary or an old weights-only file
+            if isinstance(loaded_data, dict) and PyTorchCheckpointKeys.MODEL_STATE in loaded_data:
+                # It's a new training checkpoint, extract the weights
+                self.model.load_state_dict(loaded_data[PyTorchCheckpointKeys.MODEL_STATE])
+            else:
+                # It's an old-style file (or just a state_dict), load it directly
+                self.model.load_state_dict(loaded_data)
+            
+            _LOGGER.info(f"Model state loaded from '{model_p.name}'.")
+                
             self.model.to(self.device)
             self.model.eval()  # Set the model to evaluation mode
-            _LOGGER.info(f"Model state loaded from '{model_p.name}' and set to evaluation mode.")
         except Exception as e:
             _LOGGER.error(f"Failed to load model state from '{model_p}': {e}")
             raise
@@ -72,7 +82,6 @@ class _BaseInferenceHandler(ABC):
             _LOGGER.warning("CUDA not available, switching to CPU.")
             device_lower = "cpu"
         elif device_lower == "mps" and not torch.backends.mps.is_available():
-            # Your M-series Mac will appreciate this check!
             _LOGGER.warning("Apple Metal Performance Shaders (MPS) not available, switching to CPU.")
             device_lower = "cpu"
         return torch.device(device_lower)

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,73 @@ 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,
+                 embedding_dim: int = 256,
                  num_heads: int = 8,
                  num_layers: int = 6,
-                 dropout: float = 0.1):
+                 dropout: float = 0.2):
         """
         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.
+            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. Common values: (64, 128, 192, 256, etc.)
+            num_heads (int): 
+                The number of heads in the multi-head attention mechanism. Common values: (4, 8, 16)
+            num_layers (int): 
+                The number of sub-encoder-layers in the transformer encoder. Common values: (4, 8, 12)
+            dropout (float): 
+                The dropout value.
+                
+        ## Note:
         
-        ### 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.
+        **Embedding Dimension:** "Width" of the model. It's the N-dimension vector that will be used to represent each one of the features.
+            - Each continuous feature gets its own learnable N-dimension vector.
+            - Each categorical feature gets an embedding table that maps every category (e.g., "color=red", "color=blue") to a unique N-dimension vector.
+            
+        **Attention Heads:** Controls the "Multi-Head Attention" mechanism. Instead of looking at all the feature interactions at once, the model splits its attention into N parallel heads.
+            - Embedding Dimensions get divided by the number of Attention Heads, resulting in the dimensions assigned per head.
 
-        **Standard Numerical and Continuous Features** (e.g., 'Age', 'Price'): It is highly recommended to scale them before training.
+        **Number of Layers:** "Depth" of the model. Number of identical `TransformerEncoderLayer` blocks that are stacked on top of each other.
+            - Layer 1: The attention heads find simple, direct interactions between the features.
+            - Layer 2: Takes the output of Layer 1 and finds interactions between those interactions and so on.
+            - Trade-off: More layers are more powerful but are slower to train and more prone to overfitting. If the training loss goes down but the validation loss goes up, you might have too many layers (or need more dropout).
+            
         """
         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 +414,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 +507,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 +563,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)
         
@@ -670,5 +748,7 @@ class SequencePredictorLSTM(nn.Module, _ArchitectureHandlerMixin):
         )
 
 
+# ---- PyTorch models ---
+
 def info():
     _script_info(__all__)