dragon-ml-toolbox 20.2.0__py3-none-any.whl → 20.3.0__py3-none-any.whl

ml_tools/ML_models/_dragon_node.py

@@ -0,0 +1,268 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Any, Optional, Literal
+
+from ..schema import FeatureSchema
+from .._core import get_logger
+from ..keys._keys import SchemaKeys
+
+from ._base_save_load import _ArchitectureBuilder
+from ._models_advanced_helpers import (
+    Embedding1dLayer,
+    entmax15,
+    entmoid15,
+    sparsemax,
+    sparsemoid,
+    DenseODSTBlock,
+)
+
+
+_LOGGER = get_logger("DragonNodeModel")
+
+
+__all__ = [
+    "DragonNodeModel",
+]
+
+# SOURCE CODE: Adapted and modified from:
+# https://github.com/manujosephv/pytorch_tabular/blob/main/LICENSE
+# https://github.com/Qwicen/node/blob/master/LICENSE.md
+# https://github.com/jrzaurin/pytorch-widedeep?tab=readme-ov-file#license
+# https://github.com/rixwew/pytorch-fm/blob/master/LICENSE
+# https://arxiv.org/abs/1705.08741v2
+
+
+class DragonNodeModel(_ArchitectureBuilder):
+    """
+    Native implementation of Neural Oblivious Decision Ensembles (NODE).
+    
+    The 'Dense' architecture concatenates the outputs of previous layers to the 
+    features of subsequent layers, allowing for deep feature interaction learning.
+    """
+    ACTIVATION_MAP = {
+        "entmax": entmax15,
+        "sparsemax": sparsemax,
+        "softmax": F.softmax,
+    }
+    
+    BINARY_ACTIVATION_MAP = {
+        "entmoid": entmoid15,
+        "sparsemoid": sparsemoid,
+        "sigmoid": torch.sigmoid,
+    }
+
+    def __init__(self, *,
+                 schema: FeatureSchema,
+                 out_targets: int,
+                 embedding_dim: int = 24,
+                 num_trees: int = 1024,
+                 num_layers: int = 2,
+                 tree_depth: int = 6,
+                 additional_tree_output_dim: int = 3,
+                 max_features: Optional[int] = None,
+                 input_dropout: float = 0.0,
+                 embedding_dropout: float = 0.0,
+                 choice_function: Literal['entmax', 'sparsemax', 'softmax'] = 'entmax',
+                 bin_function: Literal['entmoid', 'sparsemoid', 'sigmoid'] = 'entmoid',
+                 batch_norm_continuous: bool = False):
+        """
+        Args:
+            schema (FeatureSchema): 
+                Schema object containing feature names and types.
+            out_targets (int): 
+                Number of output targets.
+            embedding_dim (int, optional): 
+                Embedding dimension for categorical features.
+                Suggested: 16 to 64.
+            num_trees (int, optional): 
+                Number of Oblivious Decision Trees per layer. NODE relies on a large number 
+                of trees (wider layers) compared to standard forests.
+                Suggested: 512 to 2048.
+            num_layers (int, optional): 
+                Number of DenseODST layers. Since layers are densely connected, deeper 
+                networks increase memory usage significantly.
+                Suggested: 2 to 5.
+            tree_depth (int, optional): 
+                Depth of the oblivious trees. Oblivious trees are symmetric, so 
+                parameters scale with 2^depth.
+                Suggested: 4 to 8.
+            additional_tree_output_dim (int, optional): 
+                Extra output channels per tree. These are used for internal representation 
+                in deeper layers but discarded for the final prediction.
+                Suggested: 1 to 5.
+            max_features (int, optional): 
+                Max features to keep in the dense connection to prevent explosion in 
+                feature dimension for deeper layers. If None, keeps all.
+            input_dropout (float, optional): 
+                Dropout applied to the input of the Dense Block.
+                Suggested: 0.0 to 0.2.
+            embedding_dropout (float, optional): 
+                Dropout applied specifically to embeddings.
+                Suggested: 0.0 to 0.2.
+            choice_function (str, optional): 
+                Activation for feature selection. 'entmax' allows sparse feature selection.
+                Options: 'entmax', 'sparsemax', 'softmax'. 
+            bin_function (str, optional): 
+                Activation for the soft binning steps.
+                Options: 'entmoid', 'sparsemoid', 'sigmoid'.
+            batch_norm_continuous (bool, optional): 
+                If True, applies Batch Normalization to continuous features.
+        """
+        super().__init__()
+        self.schema = schema
+        self.out_targets = out_targets
+        
+        # -- Configuration for saving --
+        self.model_hparams = {
+            'embedding_dim': embedding_dim,
+            'num_trees': num_trees,
+            'num_layers': num_layers,
+            'tree_depth': tree_depth,
+            'additional_tree_output_dim': additional_tree_output_dim,
+            'max_features': max_features,
+            'input_dropout': input_dropout,
+            'embedding_dropout': embedding_dropout,
+            'choice_function': choice_function,
+            'bin_function': bin_function,
+            'batch_norm_continuous': batch_norm_continuous
+        }
+
+        # -- 1. Setup Embeddings --
+        self.categorical_indices = []
+        self.cardinalities = []
+        if schema.categorical_index_map:
+            self.categorical_indices = list(schema.categorical_index_map.keys())
+            self.cardinalities = list(schema.categorical_index_map.values())
+        
+        all_indices = set(range(len(schema.feature_names)))
+        self.numerical_indices = sorted(list(all_indices - set(self.categorical_indices)))
+        
+        embedding_dims = [(c, embedding_dim) for c in self.cardinalities]
+        n_continuous = len(self.numerical_indices)
+        
+        self.embedding_layer = Embedding1dLayer(
+            continuous_dim=n_continuous,
+            categorical_embedding_dims=embedding_dims,
+            embedding_dropout=embedding_dropout,
+            batch_norm_continuous_input=batch_norm_continuous
+        )
+        
+        total_embedded_dim = n_continuous + sum([d for _, d in embedding_dims])
+        
+        # -- 2. Backbone (Dense ODST) --
+        # The tree output dim includes the target dim + auxiliary dims for deep learning
+        self.tree_dim = out_targets + additional_tree_output_dim
+        
+        self.backbone = DenseODSTBlock(
+            input_dim=total_embedded_dim,
+            num_trees=num_trees,
+            num_layers=num_layers,
+            tree_output_dim=self.tree_dim,
+            max_features=max_features,
+            input_dropout=input_dropout,
+            flatten_output=False, # We want (Batch, Num_Layers * Num_Trees, Tree_Dim)
+            depth=tree_depth,
+            # Activations
+            choice_function=self.ACTIVATION_MAP[choice_function],
+            bin_function=self.BINARY_ACTIVATION_MAP[bin_function],
+            # Init strategies (defaults)
+            initialize_response_=nn.init.normal_,
+            initialize_selection_logits_=nn.init.uniform_,
+        )
+        
+        # Note: NODE has a fixed Head (averaging) which is defined in forward()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Split inputs
+        x_cont = x[:, self.numerical_indices].float()
+        x_cat = x[:, self.categorical_indices].long()
+        
+        # 1. Embeddings
+        x = self.embedding_layer(x_cont, x_cat)
+        
+        # 2. Backbone
+        # Output shape: (Batch, Total_Trees, Tree_Dim)
+        x = self.backbone(x)
+        
+        # 3. Head (Averaging)
+        # We take the first 'out_targets' channels and average them across all trees
+        # subset: x[..., :out_targets]
+        # mean: .mean(dim=-2) -> average over Total_Trees dimension
+        return x[..., :self.out_targets].mean(dim=-2)
+
+    def data_aware_initialization(self, train_dataset, num_samples: int = 2000, verbose: int = 3):
+        """
+        Performs data-aware initialization for the ODST trees using a dataset.
+        Crucial for NODE convergence.
+        """
+        # 1. Prepare Data
+        if verbose >= 2:
+            _LOGGER.info(f"Performing NODE data-aware initialization on up to {num_samples} samples...")
+        device = next(self.parameters()).device
+            
+        # 2. Extract Features
+        # Fast path: If the dataset exposes the full feature tensor (like _PytorchDataset)
+        if hasattr(train_dataset, "features") and isinstance(train_dataset.features, torch.Tensor):
+             # Slice directly
+             limit = min(len(train_dataset.features), num_samples)
+             x_input = train_dataset.features[:limit]
+        else:
+            # Slow path: Iterate and stack (Generic Dataset)
+            indices = range(min(len(train_dataset), num_samples))
+            x_accum = []
+            for i in indices:
+                # Expecting (features, targets) tuple from standard datasets
+                sample = train_dataset[i]
+                if isinstance(sample, (tuple, list)):
+                    x_accum.append(sample[0])
+                elif isinstance(sample, dict) and 'features' in sample:
+                    x_accum.append(sample['features'])
+                elif isinstance(sample, dict) and 'x' in sample:
+                    x_accum.append(sample['x'])
+                else:
+                    # Fallback: assume the sample itself is the feature
+                    x_accum.append(sample)
+            
+            if not x_accum:
+                if verbose >= 1:
+                    _LOGGER.warning("Dataset empty or format unrecognized. Skipping NODE initialization.")
+                return
+                
+            x_input = torch.stack(x_accum)
+            
+        x_input = x_input.to(device).float()
+        
+        # 3. Process features (Split -> Embed)
+        x_cont = x_input[:, self.numerical_indices].float()
+        x_cat = x_input[:, self.categorical_indices].long()
+        
+        with torch.no_grad():
+            x_embedded = self.embedding_layer(x_cont, x_cat)
+            
+            # 4. Initialize Backbone
+            if hasattr(self.backbone, 'initialize'):
+                self.backbone.initialize(x_embedded)
+                if verbose >= 2:
+                    _LOGGER.info("NODE Initialization Complete. Ready to train.")
+            else:
+                if verbose >= 1:
+                    _LOGGER.warning("NODE Backbone does not have an 'initialize' method. Skipping.")
+            
+    def get_architecture_config(self) -> dict[str, Any]:
+        """Returns the full configuration of the model."""
+        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
+        }
+        
+        config = {
+            SchemaKeys.SCHEMA_DICT: schema_dict,
+            'out_targets': self.out_targets,
+            **self.model_hparams
+        }
+        return config
+

ml_tools/ML_models/_dragon_tabnet.py

@@ -0,0 +1,255 @@
+import torch
+import torch.nn as nn
+from typing import Any, Literal
+
+from ..schema import FeatureSchema
+from .._core import get_logger
+from ..keys._keys import SchemaKeys
+
+from ._base_save_load import _ArchitectureBuilder
+from ._models_advanced_helpers import (
+    FeatTransformer, 
+    AttentiveTransformer, 
+    initialize_non_glu,
+)
+
+
+_LOGGER = get_logger("DragonTabNet")
+
+
+__all__ = [
+    "DragonTabNet"
+]
+
+# SOURCE CODE: Adapted and modified from:
+# https://github.com/manujosephv/pytorch_tabular/blob/main/LICENSE
+# https://github.com/Qwicen/node/blob/master/LICENSE.md
+# https://github.com/jrzaurin/pytorch-widedeep?tab=readme-ov-file#license
+# https://github.com/rixwew/pytorch-fm/blob/master/LICENSE
+# https://arxiv.org/abs/1705.08741v2
+
+
+class DragonTabNet(_ArchitectureBuilder):
+    """
+    Native implementation of TabNet (Attentive Interpretable Tabular Learning).
+    
+    Includes the Initial Splitter, Ghost Batch Norm, and GLU scaling.
+    """
+    def __init__(self, *,
+                 schema: FeatureSchema,
+                 out_targets: int,
+                 n_d: int = 8,
+                 n_a: int = 8,
+                 n_steps: int = 3,
+                 gamma: float = 1.3,
+                 n_independent: int = 2,
+                 n_shared: int = 2,
+                 virtual_batch_size: int = 128,
+                 momentum: float = 0.02,
+                 mask_type: Literal['sparsemax', 'entmax', 'softmax'] = 'sparsemax',
+                 batch_norm_continuous: bool = False):
+        """
+        Args:
+            schema (FeatureSchema): 
+                Schema object containing feature names and types.
+            out_targets (int): 
+                Number of output targets.
+            n_d (int, optional): 
+                Dimension of the prediction layer (decision step).
+                Suggested: 8 to 64.
+            n_a (int, optional): 
+                Dimension of the attention layer (masking step).
+                Suggested: 8 to 64.
+            n_steps (int, optional): 
+                Number of sequential attention steps (architecture depth).
+                Suggested: 3 to 10.
+            gamma (float, optional): 
+                Relaxation parameter for sparsity in the mask.
+                Suggested: 1.0 to 2.0.
+            n_independent (int, optional): 
+                Number of independent Gated Linear Unit (GLU) layers in each block.
+                Suggested: 1 to 5.
+            n_shared (int, optional): 
+                Number of shared GLU layers across all blocks.
+                Suggested: 1 to 5.
+            virtual_batch_size (int, optional): 
+                Batch size for Ghost Batch Normalization.
+                Suggested: 128 to 1024.
+            momentum (float, optional): 
+                Momentum for Batch Normalization.
+                Suggested: 0.01 to 0.4.
+            mask_type (str, optional): 
+                Masking function to use. 'sparsemax' enforces sparsity.
+                Options: 'sparsemax', 'entmax', 'softmax'.
+            batch_norm_continuous (bool, optional): 
+                If True, applies Batch Normalization to continuous features before processing.
+        """
+        super().__init__()
+        self.schema = schema
+        self.out_targets = out_targets
+        
+        # Save config
+        self.model_hparams = {
+            'n_d': n_d,
+            'n_a': n_a,
+            'n_steps': n_steps,
+            'gamma': gamma,
+            'n_independent': n_independent,
+            'n_shared': n_shared,
+            'virtual_batch_size': virtual_batch_size,
+            'momentum': momentum,
+            'mask_type': mask_type,
+            'batch_norm_continuous': batch_norm_continuous
+        }
+        
+        # -- 1. Setup Input Features --
+        self.categorical_indices = []
+        self.cardinalities = []
+        if schema.categorical_index_map:
+            self.categorical_indices = list(schema.categorical_index_map.keys())
+            self.cardinalities = list(schema.categorical_index_map.values())
+        
+        all_indices = set(range(len(schema.feature_names)))
+        self.numerical_indices = sorted(list(all_indices - set(self.categorical_indices)))
+        
+        # Standard TabNet Embeddings:
+        # We use a simple embedding for each categorical feature and concat with continuous.
+        self.cat_embeddings = nn.ModuleList([
+            nn.Embedding(card, 1) for card in self.cardinalities
+        ])
+        
+        self.n_continuous = len(self.numerical_indices)
+        self.input_dim = self.n_continuous + len(self.cardinalities)
+        
+        # -- 2. TabNet Backbone Components --
+        self.n_d = n_d
+        self.n_a = n_a
+        self.n_steps = n_steps
+        self.gamma = gamma
+        self.epsilon = 1e-15
+        
+        # Initial BN
+        self.initial_bn = nn.BatchNorm1d(self.input_dim, momentum=0.01)
+
+        # Shared GLU Layers
+        if n_shared > 0:
+            self.shared_feat_transform = nn.ModuleList()
+            for i in range(n_shared):
+                if i == 0:
+                    self.shared_feat_transform.append(
+                        nn.Linear(self.input_dim, 2 * (n_d + n_a), bias=False)
+                    )
+                else:
+                    self.shared_feat_transform.append(
+                        nn.Linear(n_d + n_a, 2 * (n_d + n_a), bias=False)
+                    )
+        else:
+            self.shared_feat_transform = None
+
+        # Initial Splitter
+        # This processes the input BEFORE the first step to generate the initial attention vector 'a'
+        self.initial_splitter = FeatTransformer(
+            self.input_dim,
+            n_d + n_a,
+            self.shared_feat_transform,
+            n_glu_independent=n_independent,
+            virtual_batch_size=virtual_batch_size,
+            momentum=momentum,
+        )
+
+        # Steps
+        self.feat_transformers = nn.ModuleList()
+        self.att_transformers = nn.ModuleList()
+
+        for step in range(n_steps):
+            transformer = FeatTransformer(
+                self.input_dim,
+                n_d + n_a,
+                self.shared_feat_transform,
+                n_glu_independent=n_independent,
+                virtual_batch_size=virtual_batch_size,
+                momentum=momentum,
+            )
+            attention = AttentiveTransformer(
+                n_a,
+                self.input_dim, # We assume group_dim = input_dim (no grouping)
+                virtual_batch_size=virtual_batch_size,
+                momentum=momentum,
+                mask_type=mask_type,
+            )
+            self.feat_transformers.append(transformer)
+            self.att_transformers.append(attention)
+
+        # -- 3. Final Mapping Head --
+        self.final_mapping = nn.Linear(n_d, out_targets, bias=False)
+        initialize_non_glu(self.final_mapping, n_d, out_targets)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # -- Preprocessing --
+        x_cont = x[:, self.numerical_indices].float()
+        x_cat = x[:, self.categorical_indices].long()
+        
+        cat_list = []
+        for i, embed in enumerate(self.cat_embeddings):
+            cat_list.append(embed(x_cat[:, i])) # (B, 1)
+        
+        if cat_list:
+            x_in = torch.cat([x_cont, *cat_list], dim=1)
+        else:
+            x_in = x_cont
+            
+        
+        # -- TabNet Encoder Pass --
+        x_bn = self.initial_bn(x_in)
+        # Initial Split
+        # The splitter produces [d, a]. We only need 'a' to start the loop.
+        att = self.initial_splitter(x_bn)[:, self.n_d :]
+        priors = torch.ones(x_bn.shape, device=x.device)
+        out_accumulated = 0
+        self.regularization_loss = 0
+
+        for step in range(self.n_steps):
+            # 1. Attention
+            mask = self.att_transformers[step](priors, att)
+            # 2. Accumulate sparsity loss matching original implementation
+            loss = torch.sum(torch.mul(mask, torch.log(mask + self.epsilon)), dim=1)
+            self.regularization_loss += torch.mean(loss)
+            # 3. Update Prior
+            priors = torch.mul(self.gamma - mask, priors)
+            # 4. Masking
+            masked_x = torch.mul(mask, x_bn)
+            # 5. Feature Transformer
+            out = self.feat_transformers[step](masked_x)
+            # 6. Split Output
+            d = nn.ReLU()(out[:, :self.n_d])
+            att = out[:, self.n_d:]
+            # 7. Accumulate Decision
+            out_accumulated = out_accumulated + d
+
+        self.regularization_loss /= self.n_steps
+        return self.final_mapping(out_accumulated)
+    
+    def data_aware_initialization(self, train_dataset, num_samples: int = 2000, verbose: int = 3):
+        """  
+        TabNet does not require data-aware initialization. Method Implemented for compatibility.
+        """
+        if verbose >= 2:
+            _LOGGER.info("TabNet does not require data-aware initialization. Skipping.")
+
+    def get_architecture_config(self) -> dict[str, Any]:
+        """Returns the full configuration of the model."""
+        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
+        }
+        
+        config = {
+            SchemaKeys.SCHEMA_DICT: schema_dict,
+            'out_targets': self.out_targets,
+            **self.model_hparams
+        }
+        return config
+

ml_tools/ML_models_sequence/__init__.py

@@ -2,9 +2,13 @@ from ._sequence_models import (
     DragonSequenceLSTM
 )
 
-from ._imprimir import info
+from .._core import _imprimir_disponibles
 
 
 __all__ = [
     "DragonSequenceLSTM",
 ]
+
+
+def info():
+    _imprimir_disponibles(__all__)

ml_tools/ML_models_vision/__init__.py

@@ -13,7 +13,7 @@ from ._object_detection import (
     DragonFastRCNN,
 )
 
-from ._imprimir import info
+from .._core import _imprimir_disponibles
 
 
 __all__ = [
@@ -27,3 +27,7 @@ __all__ = [
     # Object Detection
     "DragonFastRCNN",
 ]
+
+
+def info():
+    _imprimir_disponibles(__all__)

ml_tools/ML_optimization/__init__.py

@@ -1,6 +1,10 @@
-from ._multi_dragon import DragonParetoOptimizer
+from ._multi_dragon import (
+    DragonParetoOptimizer
+)
 
-from ._single_dragon import DragonOptimizer
+from ._single_dragon import (
+    DragonOptimizer
+)
 
 from ._single_manual import (
     FitnessEvaluator,
@@ -8,7 +12,7 @@ from ._single_manual import (
     run_optimization,
 )
 
-from ._imprimir import info
+from .._core import _imprimir_disponibles
 
 
 __all__ = [
@@ -19,3 +23,7 @@ __all__ = [
     "create_pytorch_problem",
     "run_optimization",
 ]
+
+
+def info():
+    _imprimir_disponibles(__all__)

ml_tools/ML_optimization/_multi_dragon.py

@@ -378,8 +378,8 @@ class DragonParetoOptimizer:
         sanitized_filename = sanitize_filename(filename)
         csv_filename = sanitized_filename if sanitized_filename.lower().endswith(".csv") else f"{sanitized_filename}.csv"
         
-        save_dataframe_filename(df=df_to_save, save_dir=save_path, filename=csv_filename)
-        _LOGGER.info(f"💾 Pareto solutions saved to CSV: '{save_path.name}/{csv_filename}'")
+        save_dataframe_filename(df=df_to_save, save_dir=save_path, filename=csv_filename, verbose=1)
+        _LOGGER.info(f"💾 Pareto solutions saved to CSV: '{save_path.name}/{csv_filename}'. Shape: {df_to_save.shape}")
         
         # Save optimization bounds as JSON for reference (debug mode)
         if self._debug: