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

ml_tools/ML_models/_mlp_attention.py

@@ -0,0 +1,134 @@
+import torch
+from torch import nn
+from typing import Any
+
+from ._base_mlp_attention import _BaseMLP, _BaseAttention, _AttentionLayer, _MultiHeadAttentionLayer
+
+
+__all__ = [
+    "DragonMLP",
+    "DragonAttentionMLP",
+    "DragonMultiHeadAttentionNet",
+]
+
+
+class DragonMLP(_BaseMLP):
+    """
+    Creates a versatile Multilayer Perceptron (MLP) for regression or classification tasks.
+    """
+    def __init__(self, in_features: int, out_targets: int,
+                 hidden_layers: list[int] = [256, 128], drop_out: float = 0.2) -> None:
+        """
+        Args:
+            in_features (int): The number of input features (e.g., columns in your data).
+            out_targets (int): The number of output targets. For regression, this is
+                typically 1. For classification, it's the number of classes.
+            hidden_layers (list[int]): A list where each integer represents the
+                number of neurons in a hidden layer.
+            drop_out (float): The dropout probability for neurons in each hidden
+                layer. Must be between 0.0 and 1.0.
+                
+        ### Rules of thumb:
+        - Choose a number of hidden neurons between the size of the input layer and the size of the output layer. 
+        - The number of hidden neurons should be 2/3 the size of the input layer, plus the size of the output layer. 
+        - The number of hidden neurons should be less than twice the size of the input layer.
+        """
+        super().__init__(in_features, out_targets, hidden_layers, drop_out)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Defines the forward pass of the model."""
+        x = self.mlp(x)
+        logits = self.output_layer(x)
+        return logits
+    
+    def __repr__(self) -> str:
+        """Returns the developer-friendly string representation of the model."""
+        # Extracts the number of neurons from each nn.Linear layer
+        layer_sizes = [str(layer.in_features) for layer in self.mlp if isinstance(layer, nn.Linear)]
+        
+        return self._repr_helper(name="DragonMLP", mlp_layers=layer_sizes)
+
+
+class DragonAttentionMLP(_BaseAttention):
+    """
+    A Multilayer Perceptron (MLP) that incorporates an Attention layer to dynamically weigh input features.
+    
+    In inference mode use `forward_attention()` to get a tuple with `(output, attention_weights)`
+    """
+    def __init__(self, in_features: int, out_targets: int,
+                 hidden_layers: list[int] = [256, 128], drop_out: float = 0.2) -> None:
+        """
+        Args:
+            in_features (int): The number of input features (e.g., columns in your data).
+            out_targets (int): The number of output targets. For regression, this is
+                typically 1. For classification, it's the number of classes.
+            hidden_layers (list[int]): A list where each integer represents the
+                number of neurons in a hidden layer.
+            drop_out (float): The dropout probability for neurons in each hidden
+                layer. Must be between 0.0 and 1.0.
+        """
+        super().__init__(in_features, out_targets, hidden_layers, drop_out)
+        # Attention
+        self.attention = _AttentionLayer(in_features)
+        self.has_interpretable_attention = True
+    
+    def __repr__(self) -> str:
+        """Returns the developer-friendly string representation of the model."""
+        # Start with the input features and the attention marker
+        arch = [str(self.in_features), "[Attention]"]
+
+        # Find all other linear layers in the MLP 
+        for layer in self.mlp[1:]: # type: ignore
+            if isinstance(layer, nn.Linear):
+                arch.append(str(layer.in_features))
+        
+        return self._repr_helper(name="DragonAttentionMLP", mlp_layers=arch)
+
+
+class DragonMultiHeadAttentionNet(_BaseAttention):
+    """
+    An MLP that incorporates a standard `nn.MultiheadAttention` layer to process
+    the input features.
+
+    In inference mode use `forward_attention()` to get a tuple with `(output, attention_weights)`.
+    """
+    def __init__(self, in_features: int, out_targets: int,
+                 hidden_layers: list[int] = [256, 128], drop_out: float = 0.2,
+                 num_heads: int = 4, attention_dropout: float = 0.1) -> None:
+        """
+        Args:
+            in_features (int): The number of input features.
+            out_targets (int): The number of output targets.
+            hidden_layers (list[int]): A list of neuron counts for each hidden layer.
+            drop_out (float): The dropout probability for the MLP layers.
+            num_heads (int): The number of attention heads.
+            attention_dropout (float): Dropout probability in the attention layer.
+        """
+        super().__init__(in_features, out_targets, hidden_layers, drop_out)
+        self.num_heads = num_heads
+        self.attention_dropout = attention_dropout
+        
+        self.attention = _MultiHeadAttentionLayer(
+            num_features=in_features,
+            num_heads=num_heads,
+            dropout=attention_dropout
+        )
+
+    def get_architecture_config(self) -> dict[str, Any]:
+        """Returns the full configuration of the model."""
+        config = super().get_architecture_config()
+        config['num_heads'] = self.num_heads
+        config['attention_dropout'] = self.attention_dropout
+        return config
+    
+    def __repr__(self) -> str:
+        """Returns the developer-friendly string representation of the model."""
+        mlp_part = " -> ".join(
+            [str(self.in_features)] + 
+            [str(h) for h in self.hidden_layers] + 
+            [str(self.out_targets)]
+        )
+        arch_str = f"{self.in_features} -> [MultiHead(h={self.num_heads})] -> {mlp_part}"
+        
+        return f"DragonMultiHeadAttentionNet(arch: {arch_str})"
+

ml_tools/{_core → ML_models}/_models_advanced_helpers.py

@@ -5,7 +5,7 @@ from torch.autograd import Function
 import math
 import random
 import numpy as np
-from typing import Optional, Tuple, List, Callable
+from typing import Optional, Callable
 
 # HELPER module for ML_models_advanced
 # SOURCE CODE: Adapted and modified from:
@@ -183,7 +183,7 @@ def entmoid15(input):
 def sparsemoid(input):
     return (0.5 * input + 0.5).clamp_(0, 1)
 
-def t_softmax(input: torch.Tensor, t: torch.Tensor = None, dim: int = -1) -> torch.Tensor:
+def t_softmax(input: torch.Tensor, t: torch.Tensor = None, dim: int = -1) -> torch.Tensor: # type: ignore
     if t is None:
         t = torch.tensor(0.5, device=input.device)
     maxes = torch.max(input, dim=dim, keepdim=True).values
@@ -250,7 +250,7 @@ class Embedding1dLayer(nn.Module):
     def __init__(
         self,
         continuous_dim: int,
-        categorical_embedding_dims: List[Tuple[int, int]],
+        categorical_embedding_dims: list[tuple[int, int]],
         embedding_dropout: float = 0.0,
         batch_norm_continuous_input: bool = False,
         virtual_batch_size: Optional[int] = None,
@@ -297,7 +297,7 @@ class Embedding1dLayer(nn.Module):
                 
         if self.embd_dropout is not None:
             embed = self.embd_dropout(embed)
-        return embed
+        return embed # type: ignore
 
 class GatedFeatureLearningUnit(nn.Module):
     def __init__(
@@ -408,14 +408,14 @@ class NeuralDecisionStump(nn.Module):
             feature_mask = self.feature_mask_function(self.feature_mask)
             
         # Repeat W for each batch size using broadcasting
-        W = torch.ones(x.size(0), 1, 1, device=x.device) * self.W
+        W = torch.ones(x.size(0), 1, 1, device=x.device) * self.W # type: ignore
         # Binning features
         # x: (B, Features) -> (B, Features, 1) -> bmm -> (B, Features, Cutpoints+1)
         x_out = torch.bmm(x.unsqueeze(-1), W) - self.cut_points.unsqueeze(0)
         x_out = self.binning_activation(x_out)
         x_out = x_out * self.leaf_responses.unsqueeze(0)
         x_out = (x_out * feature_mask.reshape(1, -1, 1)).sum(dim=1)
-        return x_out, feature_mask
+        return x_out, feature_mask # type: ignore
 
 class NeuralDecisionTree(nn.Module):
     def __init__(
@@ -455,14 +455,14 @@ class NeuralDecisionTree(nn.Module):
             layer_nodes_list = []
             layer_feature_masks = []
             for n in range(max(2 ** (d), 1)):
-                leaf_nodes, feature_mask = self._modules[f"decision_stump_{d}_{n}"](tree_input)
+                leaf_nodes, feature_mask = self._modules[f"decision_stump_{d}_{n}"](tree_input) # type: ignore
                 layer_nodes_list.append(leaf_nodes)
                 layer_feature_masks.append(feature_mask)
             layer_nodes = torch.cat(layer_nodes_list, dim=1)
             tree_input = torch.cat([x, layer_nodes], dim=1)
             feature_masks.append(layer_feature_masks)
         
-        return self.dropout(layer_nodes), feature_masks
+        return self.dropout(layer_nodes), feature_masks # type: ignore
 
 class ODST(nn.Module):
     def __init__(
@@ -569,7 +569,7 @@ class SimpleLinearHead(nn.Module):
         super().__init__()
         layers = [nn.Linear(in_units, output_dim)]
         if dropout > 0:
-            layers.append(nn.Dropout(dropout))
+            layers.append(nn.Dropout(dropout)) # type: ignore
         self.main = nn.Sequential(*layers)
         
     def forward(self, x):
@@ -649,7 +649,7 @@ class DenseODSTBlock(nn.Sequential):
             
             # Initialize the specific ODST layer
             if hasattr(layer, 'initialize'):
-                layer.initialize(layer_inp)
+                layer.initialize(layer_inp) # type: ignore
             
             # Compute output to feed the next layer in the dense block
             h = layer(layer_inp)
@@ -696,7 +696,7 @@ class Embedding2dLayer(nn.Module):
     def __init__(
         self,
         continuous_dim: int,
-        categorical_cardinality: List[int],
+        categorical_cardinality: list[int],
         embedding_dim: int,
         shared_embedding_strategy: Optional[str] = None,
         frac_shared_embed: float = 0.25,
@@ -813,7 +813,7 @@ class Embedding2dLayer(nn.Module):
         if self.embd_dropout is not None:
             embed = self.embd_dropout(embed)
             
-        return embed
+        return embed # type: ignore
     
 
 
@@ -1007,7 +1007,7 @@ class _GateHead(nn.Module):
         else:
             # Apply separate linear layers
             y_hat = torch.cat(
-                [h(backbone_features[:, :, i]).unsqueeze(1) for i, h in enumerate(self.head)],
+                [h(backbone_features[:, :, i]).unsqueeze(1) for i, h in enumerate(self.head)], # type: ignore
                 dim=1,
             )
 

ml_tools/ML_models_sequence/__init__.py

@@ -0,0 +1,10 @@
+from ._sequence_models import (
+    DragonSequenceLSTM
+)
+
+from ._imprimir import info
+
+
+__all__ = [
+    "DragonSequenceLSTM",
+]

ml_tools/ML_models_sequence/_imprimir.py

@@ -0,0 +1,8 @@
+from .._core import _imprimir_disponibles
+
+_GRUPOS = [
+    "DragonSequenceLSTM",
+]
+
+def info():
+    _imprimir_disponibles(_GRUPOS)

ml_tools/{_core/_ML_sequence_models.py → ML_models_sequence/_sequence_models.py}

@@ -2,13 +2,13 @@ import torch
 from torch import nn
 from typing import Literal
 
-from ._logger import get_logger
-from ._script_info import _script_info
-from ._keys import MLTaskKeys
-from ._ML_models import _ArchitectureHandlerMixin
+from .._core import get_logger
+from ..keys._keys import MLTaskKeys
 
+from ..ML_models._base_save_load import _ArchitectureHandlerMixin
 
-_LOGGER = get_logger("DragonModel")
+
+_LOGGER = get_logger("DragonSequenceLSTM")
 
 
 __all__ = [
@@ -137,6 +137,3 @@ class DragonSequenceLSTM(nn.Module, _ArchitectureHandlerMixin):
             f"recurrent_layers={self.lstm.num_layers}), "
             f"mode='{self.prediction_mode}')")
 
-
-def info():
-    _script_info(__all__)

ml_tools/ML_models_vision/__init__.py

@@ -0,0 +1,29 @@
+from ._image_classification import (
+    DragonResNet,
+    DragonEfficientNet,
+    DragonVGG,
+)
+
+from ._image_segmentation import (
+    DragonFCN,
+    DragonDeepLabv3
+)
+
+from ._object_detection import (
+    DragonFastRCNN,
+)
+
+from ._imprimir import info
+
+
+__all__ = [
+    # Image Classification
+    "DragonResNet",
+    "DragonEfficientNet",
+    "DragonVGG",
+    # Image Segmentation
+    "DragonFCN",
+    "DragonDeepLabv3",
+    # Object Detection
+    "DragonFastRCNN",
+]

ml_tools/ML_models_vision/_base_wrapper.py

@@ -0,0 +1,254 @@
+import torch
+from torch import nn
+import torchvision.models as vision_models
+from typing import Any, Optional
+from abc import ABC, abstractmethod
+
+from ..ML_models._base_save_load import _ArchitectureHandlerMixin
+
+from .._core import get_logger
+
+
+_LOGGER = get_logger("DragonVisionModel")
+
+
+__all__ = [
+    "_BaseVisionWrapper",
+    "_BaseSegmentationWrapper",
+]
+
+
+class _BaseVisionWrapper(nn.Module, _ArchitectureHandlerMixin, ABC):
+    """
+    Abstract base class for torchvision model wrappers.
+    
+    Handles common logic for:
+    - Model instantiation (with/without pretrained weights)
+    - Input layer modification (for custom in_channels)
+    - Output layer modification (for custom num_classes)
+    - Architecture saving/loading and representation
+    """
+    def __init__(self,
+                 num_classes: int,
+                 in_channels: int,
+                 model_name: str,
+                 init_with_pretrained: bool,
+                 weights_enum_name: Optional[str] = None):
+        super().__init__()
+        
+        # --- 1. Validation and Configuration ---
+        if not hasattr(vision_models, model_name):
+            _LOGGER.error(f"'{model_name}' is not a valid model name in torchvision.models.")
+            raise ValueError()
+
+        self.num_classes = num_classes
+        self.in_channels = in_channels
+        self.model_name = model_name
+        self._pretrained_default_transforms = None
+
+        # --- 2. Instantiate the base model ---
+        if init_with_pretrained:
+            weights_enum = getattr(vision_models, weights_enum_name, None) if weights_enum_name else None
+            weights = weights_enum.IMAGENET1K_V1 if weights_enum else None
+            
+            # Save transformations for pretrained models
+            if weights:
+                self._pretrained_default_transforms = weights.transforms()
+            
+            if weights is None and init_with_pretrained:
+                 _LOGGER.warning(f"Could not find modern weights for {model_name}. Using 'pretrained=True' legacy fallback.")
+                 self.model = getattr(vision_models, model_name)(pretrained=True)
+            else:
+                 self.model = getattr(vision_models, model_name)(weights=weights)
+        else:
+            self.model = getattr(vision_models, model_name)(weights=None)
+
+        # --- 3. Modify the input layer (using abstract method) ---
+        if in_channels != 3:
+            original_conv1 = self._get_input_layer()
+            
+            new_conv1 = nn.Conv2d(
+                in_channels,
+                original_conv1.out_channels,
+                kernel_size=original_conv1.kernel_size, # type: ignore
+                stride=original_conv1.stride, # type: ignore
+                padding=original_conv1.padding, # type: ignore
+                bias=(original_conv1.bias is not None)
+            )
+            
+            # (Optional) Average original weights if starting from pretrained
+            if init_with_pretrained and original_conv1.in_channels == 3:
+                with torch.no_grad():
+                    avg_weights = torch.mean(original_conv1.weight, dim=1, keepdim=True)
+                    new_conv1.weight[:] = avg_weights.repeat(1, in_channels, 1, 1)
+
+            self._set_input_layer(new_conv1)
+
+        # --- 4. Modify the output layer (using abstract method) ---
+        original_fc = self._get_output_layer()
+        if original_fc is None: # Handle case where layer isn't found
+             _LOGGER.error(f"Model '{model_name}' has an unexpected classifier structure. Cannot replace final layer.")
+             raise AttributeError("Could not find final classifier layer.")
+
+        num_filters = original_fc.in_features
+        self._set_output_layer(nn.Linear(num_filters, num_classes))
+
+    @abstractmethod
+    def _get_input_layer(self) -> nn.Conv2d:
+        """Returns the first convolutional layer of the model."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _set_input_layer(self, layer: nn.Conv2d):
+        """Sets the first convolutional layer of the model."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _get_output_layer(self) -> Optional[nn.Linear]:
+        """Returns the final fully-connected layer of the model."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _set_output_layer(self, layer: nn.Linear):
+        """Sets the final fully-connected layer of the model."""
+        raise NotImplementedError
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Defines the forward pass of the model."""
+        return self.model(x)
+
+    def get_architecture_config(self) -> dict[str, Any]:
+        """
+        Returns the structural configuration of the model.
+        The 'init_with_pretrained' flag is intentionally omitted,
+        as .load() should restore the architecture, not the weights.
+        """
+        return {
+            'num_classes': self.num_classes,
+            'in_channels': self.in_channels,
+            'model_name': self.model_name
+        }
+
+    def __repr__(self) -> str:
+        """Returns the developer-friendly string representation of the model."""
+        return (
+            f"{self.__class__.__name__}(model='{self.model_name}', "
+            f"in_channels={self.in_channels}, "
+            f"num_classes={self.num_classes})"
+        )
+
+
+# Image segmentation
+class _BaseSegmentationWrapper(nn.Module, _ArchitectureHandlerMixin, ABC):
+    """
+    Abstract base class for torchvision segmentation model wrappers.
+    
+    Handles common logic for:
+    - Model instantiation (with/without pretrained weights and custom num_classes)
+    - Input layer modification (for custom in_channels)
+    - Forward pass dictionary unpacking (returns 'out' tensor)
+    - Architecture saving/loading and representation
+    """
+    def __init__(self,
+                 num_classes: int,
+                 in_channels: int,
+                 model_name: str,
+                 init_with_pretrained: bool,
+                 weights_enum_name: Optional[str] = None):
+        super().__init__()
+        
+        # --- 1. Validation and Configuration ---
+        if not hasattr(vision_models.segmentation, model_name):
+            _LOGGER.error(f"'{model_name}' is not a valid model name in torchvision.models.segmentation.")
+            raise ValueError()
+
+        self.num_classes = num_classes
+        self.in_channels = in_channels
+        self.model_name = model_name
+        self._pretrained_default_transforms = None
+
+        # --- 2. Instantiate the base model ---
+        model_kwargs = {
+            'num_classes': num_classes,
+            'weights': None
+        }
+        model_constructor = getattr(vision_models.segmentation, model_name)
+
+        if init_with_pretrained:
+            weights_enum = getattr(vision_models.segmentation, weights_enum_name, None) if weights_enum_name else None
+            weights = weights_enum.DEFAULT if weights_enum else None
+            
+            # save pretrained model transformations
+            if weights:
+                self._pretrained_default_transforms = weights.transforms()
+            
+            if weights is None:
+                 _LOGGER.warning(f"Could not find modern weights for {model_name}. Using 'pretrained=True' legacy fallback.")
+                 # Legacy models used 'pretrained=True' and num_classes was separate
+                 self.model = model_constructor(pretrained=True, **model_kwargs)
+            else:
+                 # Modern way: weights object implies pretraining
+                 model_kwargs['weights'] = weights
+                 self.model = model_constructor(**model_kwargs)
+        else:
+            self.model = model_constructor(**model_kwargs)
+
+        # --- 3. Modify the input layer (using abstract method) ---
+        if in_channels != 3:
+            original_conv1 = self._get_input_layer()
+            
+            new_conv1 = nn.Conv2d(
+                in_channels,
+                original_conv1.out_channels,
+                kernel_size=original_conv1.kernel_size, # type: ignore
+                stride=original_conv1.stride, # type: ignore
+                padding=original_conv1.padding, # type: ignore
+                bias=(original_conv1.bias is not None)
+            )
+            
+            # (Optional) Average original weights if starting from pretrained
+            if init_with_pretrained and original_conv1.in_channels == 3:
+                with torch.no_grad():
+                    avg_weights = torch.mean(original_conv1.weight, dim=1, keepdim=True)
+                    new_conv1.weight[:] = avg_weights.repeat(1, in_channels, 1, 1)
+
+            self._set_input_layer(new_conv1)
+
+    @abstractmethod
+    def _get_input_layer(self) -> nn.Conv2d:
+        """Returns the first convolutional layer of the model (in the backbone)."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _set_input_layer(self, layer: nn.Conv2d):
+        """Sets the first convolutional layer of the model (in the backbone)."""
+        raise NotImplementedError
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Defines the forward pass.
+        Returns the 'out' tensor from the segmentation model's output dict.
+        """
+        output_dict = self.model(x)
+        return output_dict['out'] # Key for standard torchvision seg models
+
+    def get_architecture_config(self) -> dict[str, Any]:
+        """
+        Returns the structural configuration of the model.
+        The 'init_with_pretrained' flag is intentionally omitted,
+        as .load() should restore the architecture, not the weights.
+        """
+        return {
+            'num_classes': self.num_classes,
+            'in_channels': self.in_channels,
+            'model_name': self.model_name
+        }
+
+    def __repr__(self) -> str:
+        """Returns the developer-friendly string representation of the model."""
+        return (
+            f"{self.__class__.__name__}(model='{self.model_name}', "
+            f"in_channels={self.in_channels}, "
+            f"num_classes={self.num_classes})"
+        )
+