dragon-ml-toolbox 13.0.0__py3-none-any.whl → 14.7.0__py3-none-any.whl

ml_tools/ML_models.py

@@ -8,6 +8,7 @@ from ._logger import _LOGGER
 from .path_manager import make_fullpath
 from ._script_info import _script_info
 from .keys import PytorchModelArchitectureKeys
+from ._schema import FeatureSchema
 
 
 __all__ = [
@@ -298,76 +299,73 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
     """
     A Transformer-based model for tabular data tasks.
     
-    This model uses a Feature Tokenizer to convert all input features into a sequence of embeddings, prepends a [CLS] token, and processes the
+    This model uses a Feature Tokenizer to convert all input features into a
+    sequence of embeddings, prepends a [CLS] token, and processes the
     sequence with a standard Transformer Encoder.
     """
     def __init__(self, *,
-                 in_features: int,
+                 schema: FeatureSchema,
                  out_targets: int,
-                 categorical_index_map: Dict[int, int],
-                 embedding_dim: int = 32,
+                 embedding_dim: int = 256,
                  num_heads: int = 8,
                  num_layers: int = 6,
-                 dropout: float = 0.1):
+                 dropout: float = 0.2):
         """
         Args:
-            in_features (int): The total number of columns in the input data (features).
-            out_targets (int): Number of output targets (1 for regression).
-            categorical_index_map (Dict[int, int]): Maps categorical column index to its cardinality (number of unique categories).
-            embedding_dim (int): The dimension for all feature embeddings. Must be divisible by num_heads.
-            num_heads (int): The number of heads in the multi-head attention mechanism.
-            num_layers (int): The number of sub-encoder-layers in the transformer encoder.
-            dropout (float): The dropout value.
-            
-        Note: 
-        - All arguments are keyword-only to promote clarity.
-        - Column indices start at 0.
+            schema (FeatureSchema): 
+                The definitive schema object created by `data_exploration.finalize_feature_schema()`.
+            out_targets (int): 
+                Number of output targets (1 for regression).
+            embedding_dim (int): 
+                The dimension for all feature embeddings. Must be divisible by num_heads. Common values: (64, 128, 192, 256, etc.)
+            num_heads (int): 
+                The number of heads in the multi-head attention mechanism. Common values: (4, 8, 16)
+            num_layers (int): 
+                The number of sub-encoder-layers in the transformer encoder. Common values: (4, 8, 12)
+            dropout (float): 
+                The dropout value.
+                
+        ## Note:
         
-        ### Data Preparation
-        The model requires a specific input format. All columns in the input DataFrame must be numerical, but they are treated differently based on the 
-        provided index lists.
-
-        **Nominal Categorical Features** (e.g., 'City', 'Color'): Should **NOT** be one-hot encoded. 
-        Instead, convert them to integer codes (label encoding). You must then provide a dictionary mapping their column indices to 
-        their cardinality (the number of unique categories) via the `categorical_map` parameter.
-
-        **Ordinal & Binary Features** (e.g., 'Low/Medium/High', 'True/False'): Should be treated as **numerical**. Map them to numbers that 
-        represent their state (e.g., `{'Low': 0, 'Medium': 1}` or `{False: 0, True: 1}`). Their column indices should **NOT** be included in the 
-        `categorical_map` parameter.
+        **Embedding Dimension:** "Width" of the model. It's the N-dimension vector that will be used to represent each one of the features.
+            - Each continuous feature gets its own learnable N-dimension vector.
+            - Each categorical feature gets an embedding table that maps every category (e.g., "color=red", "color=blue") to a unique N-dimension vector.
+            
+        **Attention Heads:** Controls the "Multi-Head Attention" mechanism. Instead of looking at all the feature interactions at once, the model splits its attention into N parallel heads.
+            - Embedding Dimensions get divided by the number of Attention Heads, resulting in the dimensions assigned per head.
 
-        **Standard Numerical and Continuous Features** (e.g., 'Age', 'Price'): It is highly recommended to scale them before training.
+        **Number of Layers:** "Depth" of the model. Number of identical `TransformerEncoderLayer` blocks that are stacked on top of each other.
+            - Layer 1: The attention heads find simple, direct interactions between the features.
+            - Layer 2: Takes the output of Layer 1 and finds interactions between those interactions and so on.
+            - Trade-off: More layers are more powerful but are slower to train and more prone to overfitting. If the training loss goes down but the validation loss goes up, you might have too many layers (or need more dropout).
+            
         """
         super().__init__()
         
+         # --- Get info from schema ---
+        in_features = len(schema.feature_names)
+        categorical_index_map = schema.categorical_index_map
+
          # --- Validation ---
-        if categorical_index_map and max(categorical_index_map.keys()) >= in_features:
+        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()
         
-        # --- Derive numerical indices ---
-        all_indices = set(range(in_features))
-        categorical_indices_set = set(categorical_index_map.keys())
-        numerical_indices = sorted(list(all_indices - categorical_indices_set))
-        
         # --- Save configuration ---
-        self.in_features = in_features
+        self.schema = schema # <-- Save the whole schema
         self.out_targets = out_targets
-        self.numerical_indices = numerical_indices
-        self.categorical_map = categorical_index_map
         self.embedding_dim = embedding_dim
         self.num_heads = num_heads
         self.num_layers = num_layers
         self.dropout = dropout
 
-        # --- 1. Feature Tokenizer ---
+        # --- 1. Feature Tokenizer (now takes the schema) ---
         self.tokenizer = _FeatureTokenizer(
-            numerical_indices=numerical_indices,
-            categorical_map=categorical_index_map,
+            schema=schema,
             embedding_dim=embedding_dim
         )
         
         # --- 2. CLS Token ---
-        # A learnable token that will be prepended to the sequence.
         self.cls_token = nn.Parameter(torch.randn(1, 1, embedding_dim))
         
         # --- 3. Transformer Encoder ---
@@ -416,21 +414,87 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
     
     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 {
-            'in_features': self.in_features,
+            'schema_dict': schema_dict,
             'out_targets': self.out_targets,
-            'categorical_map': self.categorical_map,
             'embedding_dim': self.embedding_dim,
             'num_heads': self.num_heads,
             'num_layers': self.num_layers,
             'dropout': self.dropout
         }
+    
+    @classmethod
+    def load(cls: type, file_or_dir: Union[str, Path], verbose: bool = True) -> nn.Module:
+        """Loads a model architecture from a JSON file."""
+        user_path = make_fullpath(file_or_dir)
+        
+        if user_path.is_dir():
+            json_filename = PytorchModelArchitectureKeys.SAVENAME + ".json"
+            target_path = make_fullpath(user_path / json_filename, enforce="file")
+        elif user_path.is_file():
+            target_path = user_path
+        else:
+            _LOGGER.error(f"Invalid path: '{file_or_dir}'")
+            raise IOError()
+
+        with open(target_path, 'r') as f:
+            saved_data = json.load(f)
+
+        saved_class_name = saved_data[PytorchModelArchitectureKeys.MODEL]
+        config = saved_data[PytorchModelArchitectureKeys.CONFIG]
+
+        if saved_class_name != cls.__name__:
+            _LOGGER.error(f"Model class mismatch. File specifies '{saved_class_name}', but '{cls.__name__}' was expected.")
+            raise ValueError()
+
+        # --- RECONSTRUCTION LOGIC ---
+        if 'schema_dict' not in config:
+            _LOGGER.error("Invalid architecture file: missing 'schema_dict'. This file may be from an older version.")
+            raise ValueError("Missing 'schema_dict' in config.")
+            
+        schema_data = config.pop('schema_dict')
+        
+        # Re-hydrate the categorical_index_map
+        # JSON saves all dict keys as strings, so we must convert them back to int.
+        raw_index_map = schema_data['categorical_index_map']
+        if raw_index_map is not None:
+            rehydrated_index_map = {int(k): v for k, v in raw_index_map.items()}
+        else:
+            rehydrated_index_map = None
+
+        # Re-hydrate the FeatureSchema object
+        # JSON deserializes tuples as lists, so we must convert them back.
+        schema = FeatureSchema(
+            feature_names=tuple(schema_data['feature_names']),
+            continuous_feature_names=tuple(schema_data['continuous_feature_names']),
+            categorical_feature_names=tuple(schema_data['categorical_feature_names']),
+            categorical_index_map=rehydrated_index_map,
+            categorical_mappings=schema_data['categorical_mappings']
+        )
+        
+        config['schema'] = schema
+        # --- End Reconstruction ---
+
+        model = cls(**config)
+        if verbose:
+            _LOGGER.info(f"Successfully loaded architecture for '{saved_class_name}'")
+        return model
         
     def __repr__(self) -> str:
         """Returns the developer-friendly string representation of the model."""
         # Build the architecture string part-by-part
         parts = [
-            f"Tokenizer(features={self.in_features}, dim={self.embedding_dim})",
+            f"Tokenizer(features={len(self.schema.feature_names)}, dim={self.embedding_dim})",
             "[CLS]",
             f"TransformerEncoder(layers={self.num_layers}, heads={self.num_heads})",
             f"PredictionHead(outputs={self.out_targets})"
@@ -443,29 +507,41 @@ class TabularTransformer(nn.Module, _ArchitectureHandlerMixin):
 
 class _FeatureTokenizer(nn.Module):
     """
-    Transforms raw numerical and categorical features from any column order into a sequence of embeddings.
+    Transforms raw numerical and categorical features from any column order 
+    into a sequence of embeddings.
     """
     def __init__(self,
-                 numerical_indices: List[int],
-                 categorical_map: Dict[int, int],
+                 schema: FeatureSchema,
                  embedding_dim: int):
         """
         Args:
-            numerical_indices (List[int]): A list of column indices for the numerical features.
-            categorical_map (Dict[int, int]): A dictionary mapping each categorical column index to its cardinality (number of unique categories).
-            embedding_dim (int): The dimension for all feature embeddings.
+            schema (FeatureSchema): 
+                The definitive schema object from data_exploration.
+            embedding_dim (int): 
+                The dimension for all feature embeddings.
         """
         super().__init__()
         
-        # Unpack the dictionary into separate lists for indices and cardinalities
-        self.categorical_indices = list(categorical_map.keys())
-        cardinalities = list(categorical_map.values())
+        # --- Get info from schema ---
+        categorical_map = schema.categorical_index_map
+        
+        if categorical_map:
+            # Unpack the dictionary into separate lists
+            self.categorical_indices = list(categorical_map.keys())
+            cardinalities = list(categorical_map.values())
+        else:
+            self.categorical_indices = []
+            cardinalities = []
+        
+        # Derive numerical indices by finding what's not categorical
+        all_indices = set(range(len(schema.feature_names)))
+        categorical_indices_set = set(self.categorical_indices)
+        self.numerical_indices = sorted(list(all_indices - categorical_indices_set))
         
-        self.numerical_indices = numerical_indices
         self.embedding_dim = embedding_dim
         
         # A learnable embedding for each numerical feature
-        self.numerical_embeddings = nn.Parameter(torch.randn(len(numerical_indices), embedding_dim))
+        self.numerical_embeddings = nn.Parameter(torch.randn(len(self.numerical_indices), embedding_dim))
         
         # A standard embedding layer for each categorical feature
         self.categorical_embeddings = nn.ModuleList(
@@ -487,6 +563,8 @@ class _FeatureTokenizer(nn.Module):
         # Process categorical features
         categorical_tokens = []
         for i, embed_layer in enumerate(self.categorical_embeddings):
+            # x_categorical[:, i] selects the i-th categorical column
+            # (e.g., all values for the 'color' feature)
             token = embed_layer(x_categorical[:, i]).unsqueeze(1)
             categorical_tokens.append(token)
         
@@ -670,5 +748,7 @@ class SequencePredictorLSTM(nn.Module, _ArchitectureHandlerMixin):
         )
 
 
+# ---- PyTorch models ---
+
 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__)