dragon-ml-toolbox 10.1.1__py3-none-any.whl → 14.2.0__py3-none-any.whl

ml_tools/ML_models.py

@@ -3,9 +3,13 @@ from torch import nn
 from typing import List, Union, Tuple, Dict, Any
 from pathlib import Path
 import json
+
 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__ = [
     "MultilayerPerceptron",
@@ -13,12 +17,67 @@ __all__ = [
     "MultiHeadAttentionMLP",
     "TabularTransformer",
     "SequencePredictorLSTM",
-    "save_architecture",
-    "load_architecture"
 ]
 
 
-class _BaseMLP(nn.Module):
+class _ArchitectureHandlerMixin:
+    """
+    A mixin class to provide save and load functionality for model architectures.
+    """
+    def save(self: nn.Module, directory: Union[str, Path], verbose: bool = True): # type: ignore
+        """Saves the model's architecture to a JSON file."""
+        if not hasattr(self, 'get_architecture_config'):
+            _LOGGER.error(f"Model '{self.__class__.__name__}' must have a 'get_architecture_config()' method to use this functionality.")
+            raise AttributeError()
+
+        path_dir = make_fullpath(directory, make=True, enforce="directory")
+        
+        json_filename = PytorchModelArchitectureKeys.SAVENAME + ".json"
+        
+        full_path = path_dir / json_filename
+
+        config = {
+            PytorchModelArchitectureKeys.MODEL: self.__class__.__name__,
+            PytorchModelArchitectureKeys.CONFIG: self.get_architecture_config() # type: ignore
+        }
+
+        with open(full_path, 'w') as f:
+            json.dump(config, f, indent=4)
+        
+        if verbose:
+            _LOGGER.info(f"Architecture for '{self.__class__.__name__}' saved as '{full_path.name}'")
+
+    @classmethod
+    def load(cls: type, file_or_dir: Union[str, Path], verbose: bool = True) -> nn.Module:
+        """Loads a model architecture from a JSON file. If a directory is provided, the function will attempt to load a JSON file inside."""
+        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()
+
+        model = cls(**config)
+        if verbose:
+            _LOGGER.info(f"Successfully loaded architecture for '{saved_class_name}'")
+        return model
+
+
+class _BaseMLP(nn.Module, _ArchitectureHandlerMixin):
     """
     A base class for Multilayer Perceptrons.
     
@@ -68,7 +127,7 @@ class _BaseMLP(nn.Module):
         # Set a customizable Prediction Head for flexibility, specially in transfer learning and fine-tuning
         self.output_layer = nn.Linear(current_features, out_targets)
 
-    def get_config(self) -> Dict[str, Any]:
+    def get_architecture_config(self) -> Dict[str, Any]:
         """Returns the base configuration of the model."""
         return {
             'in_features': self.in_features,
@@ -90,6 +149,31 @@ class _BaseMLP(nn.Module):
         return f"{name}(arch: {arch_str})"
 
 
+class _BaseAttention(_BaseMLP):
+    """
+    Abstract base class for MLP models that incorporate an attention mechanism
+    before the main MLP layers.
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # By default, models inheriting this do not have the flag.
+        self.attention = None
+        self.has_interpretable_attention = False
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Defines the standard forward pass."""
+        logits, _attention_weights = self.forward_attention(x)
+        return logits
+
+    def forward_attention(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Returns logits and attention weights."""
+        # This logic is now shared and defined in one place
+        x, attention_weights = self.attention(x) # type: ignore
+        x = self.mlp(x)
+        logits = self.output_layer(x)
+        return logits, attention_weights
+
+
 class MultilayerPerceptron(_BaseMLP):
     """
     Creates a versatile Multilayer Perceptron (MLP) for regression or classification tasks.
@@ -127,7 +211,7 @@ class MultilayerPerceptron(_BaseMLP):
         return self._repr_helper(name="MultilayerPerceptron", mlp_layers=layer_sizes)
 
 
-class AttentionMLP(_BaseMLP):
+class AttentionMLP(_BaseAttention):
     """
     A Multilayer Perceptron (MLP) that incorporates an Attention layer to dynamically weigh input features.
     
@@ -148,25 +232,7 @@ class AttentionMLP(_BaseMLP):
         super().__init__(in_features, out_targets, hidden_layers, drop_out)
         # Attention
         self.attention = _AttentionLayer(in_features)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Defines the standard forward pass.
-        """
-        logits, _attention_weights = self.forward_attention(x)
-        return logits
-    
-    def forward_attention(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Returns logits and attention weights
-        """
-        # The attention layer returns the processed x and the weights
-        x, attention_weights = self.attention(x)
-        
-        # Pass the attention-modified tensor through the MLP
-        logits = self.mlp(x)
-        
-        return logits, attention_weights
+        self.has_interpretable_attention = True
     
     def __repr__(self) -> str:
         """Returns the developer-friendly string representation of the model."""
@@ -181,7 +247,7 @@ class AttentionMLP(_BaseMLP):
         return self._repr_helper(name="AttentionMLP", mlp_layers=arch)
 
 
-class MultiHeadAttentionMLP(_BaseMLP):
+class MultiHeadAttentionMLP(_BaseAttention):
     """
     An MLP that incorporates a standard `nn.MultiheadAttention` layer to process
     the input features.
@@ -210,27 +276,9 @@ class MultiHeadAttentionMLP(_BaseMLP):
             dropout=attention_dropout
         )
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Defines the standard forward pass of the model."""
-        logits, _attention_weights = self.forward_attention(x)
-        return logits
-    
-    def forward_attention(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Returns logits and attention weights.
-        """
-        # The attention layer returns the processed x and the weights
-        x, attention_weights = self.attention(x)
-        
-        # Pass the attention-modified tensor through the MLP and prediction head
-        x = self.mlp(x)
-        logits = self.output_layer(x)
-        
-        return logits, attention_weights
-
-    def get_config(self) -> Dict[str, Any]:
+    def get_architecture_config(self) -> Dict[str, Any]:
         """Returns the full configuration of the model."""
-        config = super().get_config()
+        config = super().get_architecture_config()
         config['num_heads'] = self.num_heads
         config['attention_dropout'] = self.attention_dropout
         return config
@@ -247,70 +295,77 @@ class MultiHeadAttentionMLP(_BaseMLP):
         return f"MultiHeadAttentionMLP(arch: {arch_str})"
 
 
-class TabularTransformer(nn.Module):
+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, *,
+                 schema: FeatureSchema,
                  out_targets: int,
-                 numerical_indices: List[int],
-                 categorical_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:
-            out_targets (int): Number of output targets (1 for regression).
-            numerical_indices (List[int]): Column indices for numerical features.
-            categorical_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 be included in the 
-        `numerical_indices` list.
+        **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 Features** (e.g., 'Age', 'Price'): Should be included in the `numerical_indices` list. 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):
+            _LOGGER.error(f"A categorical index ({max(categorical_index_map.keys())}) is out of bounds for the provided input features ({in_features}).")
+            raise ValueError()
+        
         # --- Save configuration ---
+        self.schema = schema # <-- Save the whole schema
         self.out_targets = out_targets
-        self.numerical_indices = numerical_indices
-        self.categorical_map = categorical_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_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 ---
@@ -357,25 +412,89 @@ class TabularTransformer(nn.Module):
         
         return logits
     
-    def get_config(self) -> Dict[str, Any]:
+    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 {
+            'schema_dict': schema_dict,
             'out_targets': self.out_targets,
-            'numerical_indices': self.numerical_indices,
-            '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."""
-        num_features = len(self.numerical_indices) + len(self.categorical_map)
-
         # Build the architecture string part-by-part
         parts = [
-            f"Tokenizer(features={num_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})"
@@ -388,29 +507,41 @@ class TabularTransformer(nn.Module):
 
 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(
@@ -432,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)
         
@@ -529,7 +662,7 @@ class _MultiHeadAttentionLayer(nn.Module):
         return out, attn_weights.squeeze()
 
 
-class SequencePredictorLSTM(nn.Module):
+class SequencePredictorLSTM(nn.Module, _ArchitectureHandlerMixin):
     """
     A simple LSTM-based network for sequence-to-sequence prediction tasks.
 
@@ -597,7 +730,7 @@ class SequencePredictorLSTM(nn.Module):
         
         return predictions
     
-    def get_config(self) -> dict:
+    def get_architecture_config(self) -> dict:
         """Returns the configuration of the model."""
         return {
             'features': self.features,
@@ -615,76 +748,7 @@ class SequencePredictorLSTM(nn.Module):
         )
 
 
-def save_architecture(model: nn.Module, directory: Union[str, Path], verbose: bool=True):
-    """
-    Saves a model's architecture to a 'architecture.json' file.
-
-    This function relies on the model having a `get_config()` method that
-    returns a dictionary of the arguments needed to initialize it.
-
-    Args:
-        model (nn.Module): The PyTorch model instance to save.
-        directory (str | Path): The directory to save the JSON file.
-
-    Raises:
-        AttributeError: If the model does not have a `get_config()` method.
-    """
-    if not hasattr(model, 'get_config'):
-        _LOGGER.error(f"Model '{model.__class__.__name__}' does not have a 'get_config()' method.")
-        raise AttributeError() 
-
-    # Ensure the target directory exists
-    path_dir = make_fullpath(directory, make=True, enforce="directory")
-    full_path = path_dir / "architecture.json"
-
-    config = {
-        'model_class': model.__class__.__name__,
-        'config': model.get_config() # type: ignore
-    }
-
-    with open(full_path, 'w') as f:
-        json.dump(config, f, indent=4)
-    
-    if verbose:
-        _LOGGER.info(f"Architecture for '{model.__class__.__name__}' saved to '{path_dir.name}'")
-
-
-def load_architecture(filepath: Union[str, Path], expected_model_class: type, verbose: bool=True) -> nn.Module:
-    """
-    Loads a model architecture from a JSON file.
-
-    This function instantiates a model by providing an explicit class to use
-    and checking that it matches the class name specified in the file.
-
-    Args:
-        filepath (Union[str, Path]): The path of the JSON architecture file.
-        expected_model_class (type): The model class expected to load (e.g., MultilayerPerceptron).
-
-    Returns:
-        nn.Module: An instance of the model with a freshly initialized state.
-
-    Raises:
-        FileNotFoundError: If the filepath does not exist.
-        ValueError: If the class name in the file does not match the `expected_model_class`.
-    """
-    path_obj = make_fullpath(filepath, enforce="file")
-
-    with open(path_obj, 'r') as f:
-        saved_data = json.load(f)
-
-    saved_class_name = saved_data['model_class']
-    config = saved_data['config']
-
-    if saved_class_name != expected_model_class.__name__:
-        _LOGGER.error(f"Model class mismatch. File specifies '{saved_class_name}', but '{expected_model_class.__name__}' was expected.")
-        raise ValueError()
-
-    # Create an instance of the model using the provided class and config
-    model = expected_model_class(**config)
-    if verbose:
-        _LOGGER.info(f"Successfully loaded architecture for '{saved_class_name}'")
-    return model
-
+# ---- PyTorch models ---
 
 def info():
     _script_info(__all__)