dragon-ml-toolbox 13.3.0__py3-none-any.whl → 16.2.0__py3-none-any.whl

ml_tools/ML_models.py

@@ -7,16 +7,15 @@ import json
 from ._logger import _LOGGER
 from .path_manager import make_fullpath
 from ._script_info import _script_info
-from .keys import PytorchModelArchitectureKeys
+from ._keys import PytorchModelArchitectureKeys
 from ._schema import FeatureSchema
 
 
 __all__ = [
-    "MultilayerPerceptron",
-    "AttentionMLP",
-    "MultiHeadAttentionMLP",
-    "TabularTransformer",
-    "SequencePredictorLSTM",
+    "DragonMLP",
+    "DragonAttentionMLP",
+    "DragonMultiHeadAttentionNet",
+    "DragonTabularTransformer"
 ]
 
 
@@ -174,7 +173,7 @@ class _BaseAttention(_BaseMLP):
         return logits, attention_weights
 
 
-class MultilayerPerceptron(_BaseMLP):
+class DragonMLP(_BaseMLP):
     """
     Creates a versatile Multilayer Perceptron (MLP) for regression or classification tasks.
     """
@@ -208,10 +207,10 @@ class MultilayerPerceptron(_BaseMLP):
         # 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="MultilayerPerceptron", mlp_layers=layer_sizes)
+        return self._repr_helper(name="DragonMLP", mlp_layers=layer_sizes)
 
 
-class AttentionMLP(_BaseAttention):
+class DragonAttentionMLP(_BaseAttention):
     """
     A Multilayer Perceptron (MLP) that incorporates an Attention layer to dynamically weigh input features.
     
@@ -244,10 +243,10 @@ class AttentionMLP(_BaseAttention):
             if isinstance(layer, nn.Linear):
                 arch.append(str(layer.in_features))
         
-        return self._repr_helper(name="AttentionMLP", mlp_layers=arch)
+        return self._repr_helper(name="DragonAttentionMLP", mlp_layers=arch)
 
 
-class MultiHeadAttentionMLP(_BaseAttention):
+class DragonMultiHeadAttentionNet(_BaseAttention):
     """
     An MLP that incorporates a standard `nn.MultiheadAttention` layer to process
     the input features.
@@ -292,10 +291,10 @@ class MultiHeadAttentionMLP(_BaseAttention):
         )
         arch_str = f"{self.in_features} -> [MultiHead(h={self.num_heads})] -> {mlp_part}"
         
-        return f"MultiHeadAttentionMLP(arch: {arch_str})"
+        return f"DragonMultiHeadAttentionNet(arch: {arch_str})"
 
 
-class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
+class DragonTabularTransformer(nn.Module, _ArchitectureHandlerMixin):
     """
     A Transformer-based model for tabular data tasks.
     
@@ -306,10 +305,10 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
     def __init__(self, *,
                  schema: FeatureSchema,
                  out_targets: 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:
             schema (FeatureSchema): 
@@ -317,14 +316,28 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
             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.
+                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.
+                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.
+                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).
+            
         """
         super().__init__()
         
@@ -488,7 +501,7 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
         
         arch_str = " -> ".join(parts)
         
-        return f"TabularTransformer(arch: {arch_str})"
+        return f"DragonTabularTransformer(arch: {arch_str})"
 
 
 class _FeatureTokenizer(nn.Module):
@@ -648,91 +661,5 @@ class _MultiHeadAttentionLayer(nn.Module):
         return out, attn_weights.squeeze()
 
 
-class SequencePredictorLSTM(nn.Module, _ArchitectureHandlerMixin):
-    """
-    A simple LSTM-based network for sequence-to-sequence prediction tasks.
-
-    This model is designed for datasets where each input sequence maps to an
-    output sequence of the same length. It's suitable for forecasting problems
-    prepared by the `SequenceMaker` class.
-
-    The expected input shape is `(batch_size, sequence_length, features)`.
-
-    Args:
-        features (int): The number of features in the input sequence. Defaults to 1.
-        hidden_size (int): The number of features in the LSTM's hidden state.
-                           Defaults to 100.
-        recurrent_layers (int): The number of recurrent LSTM layers. Defaults to 1.
-        dropout (float): The dropout probability for all but the last LSTM layer.
-                         Defaults to 0.
-    """
-    def __init__(self, features: int = 1, hidden_size: int = 100,
-                 recurrent_layers: int = 1, dropout: float = 0):
-        super().__init__()
-
-        # --- Validation ---
-        if not isinstance(features, int) or features < 1:
-            raise ValueError("features must be a positive integer.")
-        if not isinstance(hidden_size, int) or hidden_size < 1:
-            raise ValueError("hidden_size must be a positive integer.")
-        if not isinstance(recurrent_layers, int) or recurrent_layers < 1:
-            raise ValueError("recurrent_layers must be a positive integer.")
-        if not (0.0 <= dropout < 1.0):
-            raise ValueError("dropout must be a float between 0.0 and 1.0.")
-        
-        # --- Save configuration ---
-        self.features = features
-        self.hidden_size = hidden_size
-        self.recurrent_layers = recurrent_layers
-        self.dropout = dropout
-        
-        # Build model
-        self.lstm = nn.LSTM(
-            input_size=features,
-            hidden_size=hidden_size,
-            num_layers=recurrent_layers,
-            dropout=dropout,
-            batch_first=True  # This is crucial for (batch, seq, feature) input
-        )
-        self.linear = nn.Linear(in_features=hidden_size, out_features=features)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Defines the forward pass.
-
-        Args:
-            x (torch.Tensor): The input tensor with shape
-                              (batch_size, sequence_length, features).
-
-        Returns:
-            torch.Tensor: The output tensor with shape
-                          (batch_size, sequence_length, features).
-        """
-        # The LSTM returns the full output sequence and the final hidden/cell states
-        lstm_out, _ = self.lstm(x)
-        
-        # Pass the LSTM's output sequence to the linear layer
-        predictions = self.linear(lstm_out)
-        
-        return predictions
-    
-    def get_architecture_config(self) -> dict:
-        """Returns the configuration of the model."""
-        return {
-            'features': self.features,
-            'hidden_size': self.hidden_size,
-            'recurrent_layers': self.recurrent_layers,
-            'dropout': self.dropout
-        }
-    
-    def __repr__(self) -> str:
-        """Returns the developer-friendly string representation of the model."""
-        return (
-            f"SequencePredictorLSTM(features={self.lstm.input_size}, "
-            f"hidden_size={self.lstm.hidden_size}, "
-            f"recurrent_layers={self.lstm.num_layers})"
-        )
-
-
 def info():
     _script_info(__all__)

ml_tools/ML_models_advanced.py

@@ -0,0 +1,323 @@
+import torch
+from torch import nn
+from typing import Union, Dict, Any
+from pathlib import Path
+import json
+
+from ._logger import _LOGGER
+from .path_manager import make_fullpath
+from ._keys import PytorchModelArchitectureKeys
+from ._schema import FeatureSchema
+from ._script_info import _script_info
+from .ML_models import _ArchitectureHandlerMixin
+
+# Imports from pytorch_tabular
+try:
+    from omegaconf import DictConfig
+    from pytorch_tabular.models import GatedAdditiveTreeEnsembleModel, NODEModel
+except ImportError:
+    _LOGGER.error(f"GATE and NODE require 'pip install pytorch_tabular omegaconf' dependencies.")
+    raise ImportError()
+
+
+__all__ = [
+    "DragonGateModel",
+    "DragonNodeModel",
+]
+
+
+class _BasePytabWrapper(nn.Module, _ArchitectureHandlerMixin):
+    """
+    Internal Base Class: Do not use directly.
+    
+    This is an adapter to make pytorch_tabular models compatible with the
+    dragon-ml-toolbox pipeline.
+    
+    It handles:
+    1.  Schema-based initialization.
+    2.  Single-tensor forward pass, which is then split into the
+        dict {'continuous': ..., 'categorical': ...} that pytorch_tabular expects.
+    3.  Saving/Loading architecture using the pipeline's _ArchitectureHandlerMixin.
+    """
+    def __init__(self, schema: FeatureSchema):
+        super().__init__()
+        
+        self.schema = schema
+        self.model_name = "Base" # To be overridden by child
+        self.internal_model: nn.Module = None # type: ignore # To be set by child
+        self.model_hparams: Dict = dict() # To be set by child
+
+        # --- Derive indices from schema ---
+        categorical_map = schema.categorical_index_map
+        
+        if categorical_map:
+            # The order of keys/values is implicitly linked and must be preserved
+            self.categorical_indices = list(categorical_map.keys())
+            self.cardinalities = list(categorical_map.values())
+        else:
+            self.categorical_indices = []
+            self.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))
+
+    def _build_pt_config(self, out_targets: int, **kwargs) -> DictConfig:
+        """Helper to create the minimal config dict for a pytorch_tabular model."""
+        # 'regression' is the most neutral for model architecture. The final output_dim is what truly matters.
+        task = "regression"
+
+        config_dict = {
+            # --- Data / Schema Params ---
+            'task': task,
+            'continuous_cols': list(self.schema.continuous_feature_names),
+            'categorical_cols': list(self.schema.categorical_feature_names),
+            'continuous_dim': len(self.numerical_indices),
+            'categorical_dim': len(self.categorical_indices),
+            'categorical_cardinality': self.cardinalities,
+            'target': ['dummy_target'], # Required, but not used
+            
+            # --- Model Params ---
+            'output_dim': out_targets,
+            **kwargs
+        }
+        
+        # Add common params that most models need
+        if 'loss' not in config_dict:
+            config_dict['loss'] = 'NotUsed'
+        if 'metrics' not in config_dict:
+            config_dict['metrics'] = []
+            
+        return DictConfig(config_dict)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Accepts a single tensor and converts it to the dict
+        that pytorch_tabular models expect.
+        """
+        # 1. Split the single tensor input
+        x_cont = x[:, self.numerical_indices].float()
+        x_cat = x[:, self.categorical_indices].long()
+
+        # 2. Create the input dict
+        input_dict = {
+            'continuous': x_cont,
+            'categorical': x_cat
+        }
+        
+        # 3. Pass to the internal pytorch_tabular model
+        #    The model returns a dict, we extract the logits
+        model_output_dict = self.internal_model(input_dict)
+        
+        # 4. Return the logits tensor
+        return model_output_dict['logits']
+
+    def get_architecture_config(self) -> Dict[str, Any]:
+        """Returns the full configuration of the model."""
+        # Deconstruct schema into a JSON-friendly dict
+        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 = {
+            'schema_dict': schema_dict,
+            'out_targets': self.out_targets,
+            **self.model_hparams
+        }
+        return config
+
+    @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')
+        
+        # JSON saves all dict keys as strings, 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
+
+        # JSON deserializes tuples as lists, 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:
+        internal_model_str = str(self.internal_model)
+        # Grab the first line of the internal model's repr
+        internal_repr = internal_model_str.split('\n')[0]
+        return f"{self.model_name}(internal_model={internal_repr})"
+
+
+class DragonGateModel(_BasePytabWrapper):
+    """
+    Adapter for the Gated Additive Tree Ensemble (GATE) model from the 'pytorch_tabular' library.
+    
+    GATE is a hybrid model that uses Gated Feature Learning Units (GFLUs) to
+    learn powerful feature representations. These learned features are then
+    fed into an additive ensemble of differentiable decision trees, combining
+    the representation learning of deep networks with the structured
+    decision-making of tree ensembles.
+    """
+    def __init__(self, *,
+                 schema: FeatureSchema,
+                 out_targets: int,
+                 embedding_dim: int = 32,
+                 gflu_stages: int = 6,
+                 num_trees: int = 20,
+                 tree_depth: int = 5,
+                 dropout: float = 0.1):
+        """
+        Args:
+            schema (FeatureSchema): 
+                The definitive schema object from data_exploration.
+            out_targets (int): 
+                Number of output targets.
+            embedding_dim (int):
+                Dimension of the categorical embeddings. (Recommended: 16 to 64)
+            gflu_stages (int):
+                Number of Gated Feature Learning Units (GFLU) stages. (Recommended: 2 to 6)
+            num_trees (int):
+                Number of trees in the ensemble. (Recommended: 10 to 50)
+            tree_depth (int):
+                Depth of each tree. (Recommended: 4 to 8)
+            dropout (float):
+                Dropout rate for the GFLU.
+        """
+        super().__init__(schema)
+        self.model_name = "DragonGateModel"
+        self.out_targets = out_targets
+        
+        # Store hparams for saving/loading
+        self.model_hparams = {
+            'embedding_dim': embedding_dim,
+            'gflu_stages': gflu_stages,
+            'num_trees': num_trees,
+            'tree_depth': tree_depth,
+            'dropout': dropout
+        }
+
+        # Build the minimal config for the GateModel
+        pt_config = self._build_pt_config(
+            out_targets=out_targets,
+            embedding_dim=embedding_dim,
+            gflu_stages=gflu_stages,
+            num_trees=num_trees,
+            tree_depth=tree_depth,
+            dropout=dropout,
+            # GATE-specific params
+            gflu_dropout=dropout,
+            chain_trees=False,
+        )
+
+        # Instantiate the internal pytorch_tabular model
+        self.internal_model = GatedAdditiveTreeEnsembleModel(config=pt_config)
+
+
+class DragonNodeModel(_BasePytabWrapper):
+    """
+    Adapter for the Neural Oblivious Decision Ensembles (NODE) model from the 'pytorch_tabular' library.
+    
+    NODE is a model based on an ensemble of differentiable 'oblivious'
+    decision trees. An oblivious tree uses the same splitting feature and
+    threshold across all nodes at the same depth. This structure, combined
+    with a differentiable formulation, allows the model to be trained
+    end-to-end with gradient descent, learning feature interactions and
+    splitting thresholds simultaneously.
+    """
+    def __init__(self, *,
+                 schema: FeatureSchema,
+                 out_targets: int,
+                 embedding_dim: int = 32,
+                 num_trees: int = 1024,
+                 tree_depth: int = 6,
+                 dropout: float = 0.1):
+        """
+        Args:
+            schema (FeatureSchema): 
+                The definitive schema object from data_exploration.
+            out_targets (int): 
+                Number of output targets.
+            embedding_dim (int):
+                Dimension of the categorical embeddings. (Recommended: 16 to 64)
+            num_trees (int):
+                Total number of trees in the ensemble. (Recommended: 256 to 2048)
+            tree_depth (int):
+                Depth of each tree. (Recommended: 4 to 8)
+            dropout (float):
+                Dropout rate.
+        """
+        super().__init__(schema)
+        self.model_name = "DragonNodeModel"
+        self.out_targets = out_targets
+
+        # Store hparams for saving/loading
+        self.model_hparams = {
+            'embedding_dim': embedding_dim,
+            'num_trees': num_trees,
+            'tree_depth': tree_depth,
+            'dropout': dropout
+        }
+
+        # Build the minimal config for the NodeModel
+        pt_config = self._build_pt_config(
+            out_targets=out_targets,
+            embedding_dim=embedding_dim,
+            num_trees=num_trees,
+            tree_depth=tree_depth,
+            # NODE-specific params
+            num_layers=1, # NODE uses num_layers=1 for a single ensemble
+            total_trees=num_trees,
+            dropout_rate=dropout,
+        )
+
+        # Instantiate the internal pytorch_tabular model
+        self.internal_model = NODEModel(config=pt_config)
+
+
+def info():
+    _script_info(__all__)

ml_tools/ML_optimization.py

@@ -14,9 +14,9 @@ from functools import partial
 from .path_manager import make_fullpath, sanitize_filename
 from ._logger import _LOGGER
 from ._script_info import _script_info
-from .ML_inference import PyTorchInferenceHandler
-from .keys import PyTorchInferenceKeys
-from .SQL import DatabaseManager
+from .ML_inference import DragonInferenceHandler
+from ._keys import PyTorchInferenceKeys
+from .SQL import DragonSQL
 from .optimization_tools import _save_result, create_optimization_bounds
 from .utilities import save_dataframe_filename
 from .math_utilities import discretize_categorical_values
@@ -24,14 +24,14 @@ from ._schema import FeatureSchema
 
 
 __all__ = [
-    "MLOptimizer",
+    "DragonOptimizer",
     "FitnessEvaluator",
     "create_pytorch_problem",
     "run_optimization"
 ]
 
 
-class MLOptimizer:
+class DragonOptimizer:
     """
     A wrapper class for setting up and running EvoTorch optimization tasks.
 
@@ -47,7 +47,7 @@ class MLOptimizer:
         >>> cont_bounds = {'feature_A': (0, 100), 'feature_B': (-10, 10)}
         >>>
         >>> # 3. Initialize the optimizer
-        >>> optimizer = MLOptimizer(
+        >>> optimizer = DragonOptimizer(
         ...     inference_handler=my_handler,
         ...     schema=schema,
         ...     continuous_bounds_map=cont_bounds,
@@ -63,7 +63,7 @@ class MLOptimizer:
         ... )
     """
     def __init__(self,
-                 inference_handler: PyTorchInferenceHandler,
+                 inference_handler: DragonInferenceHandler,
                  schema: FeatureSchema,
                  continuous_bounds_map: Dict[str, Tuple[float, float]],
                  task: Literal["min", "max"],
@@ -75,7 +75,7 @@ class MLOptimizer:
         Initializes the optimizer by creating the EvoTorch problem and searcher.
 
         Args:
-            inference_handler (PyTorchInferenceHandler): 
+            inference_handler (DragonInferenceHandler): 
                 An initialized inference handler containing the model.
             schema (FeatureSchema): 
                 The definitive schema object from data_exploration.
@@ -172,18 +172,18 @@ class FitnessEvaluator:
     A callable class that wraps the PyTorch model inference handler and performs
     on-the-fly discretization for the EvoTorch fitness function.
 
-    This class is automatically instantiated by MLOptimizer and passed to
+    This class is automatically instantiated by DragonOptimizer and passed to
     create_pytorch_problem, encapsulating the evaluation logic.
     """
     def __init__(self,
-                 inference_handler: PyTorchInferenceHandler,
+                 inference_handler: DragonInferenceHandler,
                  categorical_index_map: Optional[Dict[int, int]] = None,
                  discretize_start_at_zero: bool = True):
         """
         Initializes the fitness evaluator.
 
         Args:
-            inference_handler (PyTorchInferenceHandler): 
+            inference_handler (DragonInferenceHandler): 
                 An initialized inference handler containing the model.
             categorical_index_map (Dict[int, int] | None): 
                 Maps {column_index: cardinality} for discretization.
@@ -426,7 +426,7 @@ def run_optimization(
         _LOGGER.info(f"🏁 Starting optimal solution space analysis with {repetitions} repetitions...")
         
         first_run_logger = None # To store the logger from the first rep
-        db_context = DatabaseManager(db_path) if save_format in ['sqlite', 'both'] else nullcontext()
+        db_context = DragonSQL(db_path) if save_format in ['sqlite', 'both'] else nullcontext()
         
         with db_context as db_manager:
             # --- Setup Database Schema (if applicable) ---