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

ml_tools/ML_models/_base_mlp_attention.py

@@ -0,0 +1,198 @@
+import torch
+from torch import nn
+from typing import Any
+
+from .._core import get_logger
+
+from ._base_save_load import _ArchitectureHandlerMixin
+
+
+_LOGGER = get_logger("DragonModel: MLP")
+
+
+__all__ = [
+    "_BaseMLP",
+    "_BaseAttention",
+    "_AttentionLayer",
+    "_MultiHeadAttentionLayer",
+]
+
+
+class _BaseMLP(nn.Module, _ArchitectureHandlerMixin):
+    """
+    A base class for Multilayer Perceptrons.
+    
+    Handles validation, configuration, and the creation of the core MLP layers,
+    allowing subclasses to define their own pre-processing and forward pass.
+    """
+    def __init__(self, 
+                 in_features: int, 
+                 out_targets: int,
+                 hidden_layers: list[int], 
+                 drop_out: float) -> None:
+        super().__init__()
+
+        # --- Validation ---
+        if not isinstance(in_features, int) or in_features < 1:
+            _LOGGER.error("'in_features' must be a positive integer.")
+            raise ValueError()
+        if not isinstance(out_targets, int) or out_targets < 1:
+            _LOGGER.error("'out_targets' must be a positive integer.")
+            raise ValueError()
+        if not isinstance(hidden_layers, list) or not all(isinstance(n, int) for n in hidden_layers):
+            _LOGGER.error("'hidden_layers' must be a list of integers.")
+            raise TypeError()
+        if not (0.0 <= drop_out < 1.0):
+            _LOGGER.error("'drop_out' must be a float between 0.0 and 1.0.")
+            raise ValueError()
+        
+        # --- Save configuration ---
+        self.in_features = in_features
+        self.out_targets = out_targets
+        self.hidden_layers = hidden_layers
+        self.drop_out = drop_out
+
+        # --- Build the core MLP network ---
+        mlp_layers = []
+        current_features = in_features
+        for neurons in hidden_layers:
+            mlp_layers.extend([
+                nn.Linear(current_features, neurons),
+                nn.BatchNorm1d(neurons),
+                nn.ReLU(),
+                nn.Dropout(p=drop_out)
+            ])
+            current_features = neurons
+        
+        self.mlp = nn.Sequential(*mlp_layers)
+        # 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_architecture_config(self) -> dict[str, Any]:
+        """Returns the base configuration of the model."""
+        return {
+            'in_features': self.in_features,
+            'out_targets': self.out_targets,
+            'hidden_layers': self.hidden_layers,
+            'drop_out': self.drop_out
+        }
+        
+    def _repr_helper(self, name: str, mlp_layers: list[str]):
+        last_layer = self.output_layer
+        if isinstance(last_layer, nn.Linear):
+            mlp_layers.append(str(last_layer.out_features))
+        else:
+            mlp_layers.append("Custom Prediction Head")
+        
+        # Creates a string like: 10 -> 40 -> 80 -> 40 -> 2
+        arch_str = ' -> '.join(mlp_layers)
+        
+        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 _AttentionLayer(nn.Module):
+    """
+    Calculates attention weights and applies them to the input features, incorporating a residual connection for improved stability and performance.
+    
+    Returns both the final output and the weights for interpretability.
+    """
+    def __init__(self, num_features: int):
+        super().__init__()
+        # The hidden layer size is a hyperparameter
+        hidden_size = max(16, num_features // 4)
+        
+        # Learn to produce attention scores
+        self.attention_net = nn.Sequential(
+            nn.Linear(num_features, hidden_size),
+            nn.Tanh(),
+            nn.Linear(hidden_size, num_features) # Output one score per feature
+        )
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # x shape: (batch_size, num_features)
+        
+        # Get one raw "importance" score per feature
+        attention_scores = self.attention_net(x)
+        
+        # Apply the softmax module to get weights that sum to 1
+        attention_weights = self.softmax(attention_scores)
+        
+        # Weighted features (attention mechanism's output)
+        weighted_features = x * attention_weights
+        
+        # Residual connection
+        residual_connection = x + weighted_features
+        
+        return residual_connection, attention_weights
+
+
+class _MultiHeadAttentionLayer(nn.Module):
+    """
+    A wrapper for the standard `torch.nn.MultiheadAttention` layer.
+
+    This layer treats the entire input feature vector as a single item in a
+    sequence and applies self-attention to it. It is followed by a residual
+    connection and layer normalization, which is a standard block in
+    Transformer-style models.
+    """
+    def __init__(self, num_features: int, num_heads: int, dropout: float):
+        super().__init__()
+        self.attention = nn.MultiheadAttention(
+            embed_dim=num_features,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True  # Crucial for (batch, seq, feature) input
+        )
+        self.layer_norm = nn.LayerNorm(num_features)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # x shape: (batch_size, num_features)
+
+        # nn.MultiheadAttention expects a sequence dimension.
+        # We add a sequence dimension of length 1.
+        # x_reshaped shape: (batch_size, 1, num_features)
+        x_reshaped = x.unsqueeze(1)
+
+        # Apply self-attention. query, key, and value are all the same.
+        # attn_output shape: (batch_size, 1, num_features)
+        # attn_weights shape: (batch_size, 1, 1)
+        attn_output, attn_weights = self.attention(
+            query=x_reshaped,
+            key=x_reshaped,
+            value=x_reshaped,
+            need_weights=True,
+            average_attn_weights=True # Average weights across heads
+        )
+
+        # Add residual connection and apply layer normalization (Post-LN)
+        out = self.layer_norm(x + attn_output.squeeze(1))
+
+        # Squeeze weights for a consistent output shape
+        return out, attn_weights.squeeze()
+

ml_tools/{_core/_models_advanced_base.py → ML_models/_base_save_load.py}

@@ -1,42 +1,40 @@
 from torch import nn
-from typing import Any, Dict, Union
-import json
+from typing import Union, Any
 from pathlib import Path
+import json
 from abc import ABC, abstractmethod
 
-from ._path_manager import make_fullpath
-from ._keys import PytorchModelArchitectureKeys
-from ._schema import FeatureSchema
-from ._logger import get_logger
+from ..schema import FeatureSchema
 
+from .._core import get_logger
+from ..path_manager import make_fullpath
+from ..keys._keys import PytorchModelArchitectureKeys, SchemaKeys
 
-_LOGGER = get_logger("DragonModel")
+from .._core._schema_load_ops import prepare_schema_from_json
 
 
+_LOGGER = get_logger("DragonModel: Save/Load")
+
+
+__all__ = [
+    "_ArchitectureHandlerMixin",
+    "_ArchitectureBuilder",
+]
+
 ##################################
-# Base class for loading and saving advanced models
-##################################
-class _ArchitectureBuilder(nn.Module, ABC):
+# Mixin class for saving and loading basic model architectures
+class _ArchitectureHandlerMixin:
     """
-    Base class for Dragon models that unifies architecture handling.
-    
-    Implements:
-    - JSON serialization and JSON deserialization with automatic FeatureSchema reconstruction.
-    - Standardized string representation (__repr__) showing hyperparameters.
+    A mixin class to provide save and load functionality for model architectures.
     """
-    def __init__(self):
-        super().__init__()
-        # Placeholder for hyperparameters, to be populated by child classes
-        self.model_hparams: Dict[str, Any] = {}
-        
     # abstract method that must be implemented by children
     @abstractmethod
-    def get_architecture_config(self) -> Dict[str, Any]:
+    def get_architecture_config(self) -> dict[str, Any]:
         "To be implemented by children"
         pass
-
-    def save(self: nn.Module, directory: Union[str, Path], verbose: bool = True): # type: ignore
-        """Saves the model's architecture to a JSON file."""
+    
+    def save_architecture(self: nn.Module, directory: Union[str, Path], verbose: bool = True): # type: ignore
+        """Saves the model's architecture to an "architecture.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()
@@ -59,9 +57,10 @@ class _ArchitectureBuilder(nn.Module, ABC):
             _LOGGER.info(f"Architecture for '{self.__class__.__name__}' saved as '{full_path.name}'")
 
     @classmethod
-    def load(cls, file_or_dir: Union[str, Path], verbose: bool = True) -> nn.Module:
+    def load_architecture(cls: type, file_or_dir: Union[str, Path], verbose: bool = True) -> nn.Module:
         """
-        Loads a model architecture from a JSON file.
+        Loads a model architecture from a JSON file. 
+        If a directory is provided, the function will attempt to load the JSON file "architecture.json" inside.
         """
         user_path = make_fullpath(file_or_dir)
         
@@ -84,35 +83,59 @@ class _ArchitectureBuilder(nn.Module, ABC):
             _LOGGER.error(f"Model class mismatch. File specifies '{saved_class_name}', but '{cls.__name__}' was expected.")
             raise ValueError()
 
-        # --- Schema Reconstruction Logic (Unified) ---
-        if 'schema_dict' not in config:
-            raise ValueError("Missing 'schema_dict' in config.")
-            
-        schema_data = config.pop('schema_dict')
-        
-        raw_index_map = schema_data['categorical_index_map']
-        if raw_index_map is not None:
-            # JSON keys are strings, convert back to int
-            rehydrated_index_map = {int(k): v for k, v in raw_index_map.items()}
-        else:
-            rehydrated_index_map = None
-
-        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 ---
+        # Hook to allow children classes to modify config before init (reconstruction)
+        config = cls._prepare_config_for_load(config)
 
         model = cls(**config)
         if verbose:
             _LOGGER.info(f"Successfully loaded architecture for '{saved_class_name}'")
         return model
 
+    @classmethod
+    def _prepare_config_for_load(cls, config: dict[str, Any]) -> dict[str, Any]:
+        """
+        Hook method to process configuration data before model instantiation.
+        Base implementation simply returns the config as-is.
+        """
+        return config
+
+
+##################################
+# Base class for loading and saving advanced models
+##################################
+class _ArchitectureBuilder(_ArchitectureHandlerMixin, nn.Module, ABC):
+    """
+    Base class for Dragon models that unifies architecture handling.
+    
+    Implements:
+    - JSON serialization and JSON deserialization with automatic FeatureSchema reconstruction.
+    - Standardized string representation (__repr__) showing hyperparameters.
+    """
+    def __init__(self):
+        super().__init__()
+        # Placeholder for hyperparameters, to be populated by child classes
+        self.model_hparams: dict[str, Any] = {}
+
+    @classmethod
+    def _prepare_config_for_load(cls, config: dict[str, Any]) -> dict[str, Any]:
+        """
+        Overrides the mixin hook to reconstruct the FeatureSchema object
+        from the raw dictionary data found in the JSON.
+        """
+        if SchemaKeys.SCHEMA_DICT not in config:
+            _LOGGER.error(f"The model architecture is missing the '{SchemaKeys.SCHEMA_DICT}' key.")
+            raise ValueError()
+            
+        schema_data = config.pop(SchemaKeys.SCHEMA_DICT)
+        
+        # Use shared helper to prepare arguments (handles tuple/int conversion)
+        schema_kwargs = prepare_schema_from_json(schema_data)
+        
+        schema = FeatureSchema(**schema_kwargs)
+        
+        config['schema'] = schema
+        return config
+
     def __repr__(self):
         # 1. Format hyperparameters
         hparams_str = ",\n  ".join([f"{k}={v}" for k, v in self.model_hparams.items()])
@@ -130,3 +153,4 @@ class _ArchitectureBuilder(nn.Module, ABC):
             main_str += "\n".join(child_lines) + "\n"
         main_str += ")"
         return main_str
+

ml_tools/ML_models/_dragon_tabular.py

@@ -0,0 +1,248 @@
+import torch
+from torch import nn
+from typing import Any
+
+from ..schema import FeatureSchema
+
+from .._core import get_logger
+from ..keys._keys import SchemaKeys
+
+from ._base_save_load import _ArchitectureBuilder
+
+
+_LOGGER = get_logger("DragonTabularTransformer")
+
+
+__all__ = [
+    "DragonTabularTransformer"
+]
+
+
+class DragonTabularTransformer(_ArchitectureBuilder):
+    """
+    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
+    sequence with a standard Transformer Encoder.
+    """
+    def __init__(self, *,
+                 schema: FeatureSchema,
+                 out_targets: int,
+                 embedding_dim: int = 256,
+                 num_heads: int = 8,
+                 num_layers: int = 6,
+                 dropout: float = 0.2):
+        """
+        Args:
+            schema (FeatureSchema): 
+                The definitive FeatureSchema object.
+            out_targets (int): 
+                Number of output targets.
+            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:
+        
+        **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.
+
+        **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).
+            
+        """
+        # _ArchitectureBuilder init sets up self.model_hparams
+        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.embedding_dim = embedding_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.dropout = dropout
+
+        # --- 1. Feature Tokenizer (now takes the schema) ---
+        self.tokenizer = _FeatureTokenizer(
+            schema=schema,
+            embedding_dim=embedding_dim
+        )
+        
+        # --- 2. CLS Token ---
+        self.cls_token = nn.Parameter(torch.randn(1, 1, embedding_dim))
+        
+        # --- 3. Transformer Encoder ---
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=embedding_dim,
+            nhead=num_heads,
+            dropout=dropout,
+            batch_first=True # Crucial for (batch, seq, feature) input
+        )
+        self.transformer_encoder = nn.TransformerEncoder(
+            encoder_layer=encoder_layer,
+            num_layers=num_layers
+        )
+        
+        # --- 4. Prediction Head ---
+        self.output_layer = nn.Linear(embedding_dim, out_targets)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Defines the forward pass of the model."""
+        # Get the batch size for later use
+        batch_size = x.shape[0]
+        
+        # 1. Get feature tokens from the tokenizer
+        # -> tokens shape: (batch_size, num_features, embedding_dim)
+        tokens = self.tokenizer(x)
+        
+        # 2. Prepend the [CLS] token to the sequence
+        # -> cls_tokens shape: (batch_size, 1, embedding_dim)
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+        # -> full_sequence shape: (batch_size, num_features + 1, embedding_dim)
+        full_sequence = torch.cat([cls_tokens, tokens], dim=1)
+
+        # 3. Pass the full sequence through the Transformer Encoder
+        # -> transformer_out shape: (batch_size, num_features + 1, embedding_dim)
+        transformer_out = self.transformer_encoder(full_sequence)
+        
+        # 4. Isolate the output of the [CLS] token (it's the first one)
+        # -> cls_output shape: (batch_size, embedding_dim)
+        cls_output = transformer_out[:, 0]
+        
+        # 5. Pass the [CLS] token's output through the prediction head
+        # -> logits shape: (batch_size, out_targets)
+        logits = self.output_layer(cls_output)
+        
+        return logits
+    
+    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 {
+            SchemaKeys.SCHEMA_DICT: schema_dict,
+            'out_targets': self.out_targets,
+            'embedding_dim': self.embedding_dim,
+            'num_heads': self.num_heads,
+            'num_layers': self.num_layers,
+            'dropout': self.dropout
+        }
+        
+    def __repr__(self) -> str:
+        """Returns the developer-friendly string representation of the model."""
+        # Build the architecture string part-by-part
+        parts = [
+            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})"
+        ]
+        
+        arch_str = " -> ".join(parts)
+        
+        return f"DragonTabularTransformer(arch: {arch_str})"
+
+
+class _FeatureTokenizer(nn.Module):
+    """
+    Transforms raw numerical and categorical features from any column order 
+    into a sequence of embeddings.
+    """
+    def __init__(self,
+                 schema: FeatureSchema,
+                 embedding_dim: int):
+        """
+        Args:
+            schema (FeatureSchema): 
+                The definitive schema object from data_exploration.
+            embedding_dim (int): 
+                The dimension for all feature embeddings.
+        """
+        super().__init__()
+        
+        # --- 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.embedding_dim = embedding_dim
+        
+        # A learnable embedding for each numerical feature
+        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(
+            [nn.Embedding(num_embeddings=c, embedding_dim=embedding_dim) for c in cardinalities]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Processes features from a single input tensor and concatenates them
+        into a sequence of tokens.
+        """
+        # Select the correct columns for each type using the stored indices
+        x_numerical = x[:, self.numerical_indices].float()
+        x_categorical = x[:, self.categorical_indices].long()
+
+        # Process numerical features
+        numerical_tokens = x_numerical.unsqueeze(-1) * self.numerical_embeddings
+        
+        # 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)
+        
+        # Concatenate all tokens into a single sequence
+        if not self.categorical_indices:
+             all_tokens = numerical_tokens
+        elif not self.numerical_indices:
+             all_tokens = torch.cat(categorical_tokens, dim=1)
+        else:
+             all_categorical_tokens = torch.cat(categorical_tokens, dim=1)
+             all_tokens = torch.cat([numerical_tokens, all_categorical_tokens], dim=1)
+        
+        return all_tokens
+

ml_tools/ML_models/_imprimir.py

@@ -0,0 +1,18 @@
+from .._core import _imprimir_disponibles
+
+_GRUPOS = [
+    # MLP and Attention Models
+    "DragonMLP",
+    "DragonAttentionMLP",
+    "DragonMultiHeadAttentionNet",
+    # Tabular Transformer Model
+    "DragonTabularTransformer",
+    # Advanced Models
+    "DragonGateModel",
+    "DragonNodeModel",
+    "DragonAutoInt",
+    "DragonTabNet",
+]
+
+def info():
+    _imprimir_disponibles(_GRUPOS)