dragon-ml-toolbox 14.1.0__tar.gz → 14.2.0__tar.gz

{dragon_ml_toolbox-14.1.0 → dragon_ml_toolbox-14.2.0}/LICENSE-THIRD-PARTY.md

@@ -27,3 +27,13 @@ This project depends on the following third-party packages. Each is governed by
 - [plotnine](https://github.com/has2k1/plotnine/blob/main/LICENSE)
 - [tqdm](https://github.com/tqdm/tqdm/blob/master/LICENSE)
 - [pyarrow](https://github.com/apache/arrow/blob/main/LICENSE.txt)
+- [colorlog](https://github.com/borntyping/python-colorlog/blob/main/LICENSE)
+- [evotorch](https://github.com/nnaisense/evotorch/blob/master/LICENSE)
+- [FreeSimpleGUI](https://github.com/spyoungtech/FreeSimpleGUI/blob/main/license.txt)
+- [nuitka](https://github.com/Nuitka/Nuitka/blob/main/LICENSE.txt)
+- [omegaconf](https://github.com/omry/omegaconf/blob/master/LICENSE)
+- [ordered-set](https://github.com/rspeer/ordered-set/blob/master/MIT-LICENSE)
+- [pyinstaller](https://github.com/pyinstaller/pyinstaller/blob/develop/COPYING.txt)
+- [pytorch_tabular](https://github.com/manujosephv/pytorch_tabular/blob/main/LICENSE)
+- [torchmetrics](https://github.com/Lightning-AI/torchmetrics/blob/master/LICENSE)
+- [zstandard](https://github.com/indygreg/python-zstandard/blob/main/LICENSE)

{dragon_ml_toolbox-14.1.0/dragon_ml_toolbox.egg-info → dragon_ml_toolbox-14.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dragon-ml-toolbox
-Version: 14.1.0
+Version: 14.2.0
 Summary: A collection of tools for data science and machine learning projects.
 Author-email: "Karl L. Loza Vidaurre" <luigiloza@gmail.com>
 License-Expression: MIT
@@ -35,6 +35,9 @@ Requires-Dist: evotorch; extra == "ml"
 Requires-Dist: pyarrow; extra == "ml"
 Requires-Dist: colorlog; extra == "ml"
 Requires-Dist: torchmetrics; extra == "ml"
+Provides-Extra: py-tab
+Requires-Dist: pytorch_tabular; extra == "py-tab"
+Requires-Dist: omegaconf; extra == "py-tab"
 Provides-Extra: mice
 Requires-Dist: numpy<2.0; extra == "mice"
 Requires-Dist: pandas; extra == "mice"
@@ -143,10 +146,16 @@ ML_evaluation_multi
 ML_evaluation
 ML_inference
 ML_models
+ML_models_advanced # Requires the extra flag [py-tab]
 ML_optimization
 ML_scaler
 ML_trainer
 ML_utilities
+ML_vision_datasetmaster
+ML_vision_evaluation
+ML_vision_inference
+ML_vision_models
+ML_vision_transformers
 optimization_tools
 path_manager
 PSO_optimization
@@ -192,7 +201,6 @@ pip install "dragon-ml-toolbox[excel]"
 #### Modules:
 
 ```Bash
-constants
 custom_logger
 handle_excel
 path_manager

{dragon_ml_toolbox-14.1.0 → dragon_ml_toolbox-14.2.0}/README.md

@@ -67,10 +67,16 @@ ML_evaluation_multi
 ML_evaluation
 ML_inference
 ML_models
+ML_models_advanced # Requires the extra flag [py-tab]
 ML_optimization
 ML_scaler
 ML_trainer
 ML_utilities
+ML_vision_datasetmaster
+ML_vision_evaluation
+ML_vision_inference
+ML_vision_models
+ML_vision_transformers
 optimization_tools
 path_manager
 PSO_optimization
@@ -116,7 +122,6 @@ pip install "dragon-ml-toolbox[excel]"
 #### Modules:
 
 ```Bash
-constants
 custom_logger
 handle_excel
 path_manager

{dragon_ml_toolbox-14.1.0 → dragon_ml_toolbox-14.2.0/dragon_ml_toolbox.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dragon-ml-toolbox
-Version: 14.1.0
+Version: 14.2.0
 Summary: A collection of tools for data science and machine learning projects.
 Author-email: "Karl L. Loza Vidaurre" <luigiloza@gmail.com>
 License-Expression: MIT
@@ -35,6 +35,9 @@ Requires-Dist: evotorch; extra == "ml"
 Requires-Dist: pyarrow; extra == "ml"
 Requires-Dist: colorlog; extra == "ml"
 Requires-Dist: torchmetrics; extra == "ml"
+Provides-Extra: py-tab
+Requires-Dist: pytorch_tabular; extra == "py-tab"
+Requires-Dist: omegaconf; extra == "py-tab"
 Provides-Extra: mice
 Requires-Dist: numpy<2.0; extra == "mice"
 Requires-Dist: pandas; extra == "mice"
@@ -143,10 +146,16 @@ ML_evaluation_multi
 ML_evaluation
 ML_inference
 ML_models
+ML_models_advanced # Requires the extra flag [py-tab]
 ML_optimization
 ML_scaler
 ML_trainer
 ML_utilities
+ML_vision_datasetmaster
+ML_vision_evaluation
+ML_vision_inference
+ML_vision_models
+ML_vision_transformers
 optimization_tools
 path_manager
 PSO_optimization
@@ -192,7 +201,6 @@ pip install "dragon-ml-toolbox[excel]"
 #### Modules:
 
 ```Bash
-constants
 custom_logger
 handle_excel
 path_manager

{dragon_ml_toolbox-14.1.0 → dragon_ml_toolbox-14.2.0}/dragon_ml_toolbox.egg-info/SOURCES.txt

@@ -17,6 +17,7 @@ ml_tools/ML_evaluation.py
 ml_tools/ML_evaluation_multi.py
 ml_tools/ML_inference.py
 ml_tools/ML_models.py
+ml_tools/ML_models_advanced.py
 ml_tools/ML_optimization.py
 ml_tools/ML_scaler.py
 ml_tools/ML_trainer.py
@@ -30,7 +31,6 @@ ml_tools/PSO_optimization.py
 ml_tools/RNN_forecast.py
 ml_tools/SQL.py
 ml_tools/VIF_factor.py
-ml_tools/_ML_pytorch_tabular.py
 ml_tools/_ML_vision_recipe.py
 ml_tools/__init__.py
 ml_tools/_logger.py

{dragon_ml_toolbox-14.1.0 → dragon_ml_toolbox-14.2.0}/dragon_ml_toolbox.egg-info/requires.txt

@@ -63,5 +63,9 @@ nuitka
 zstandard
 ordered-set
 
+[py-tab]
+pytorch_tabular
+omegaconf
+
 [pyinstaller]
 pyinstaller

{dragon_ml_toolbox-14.1.0 → dragon_ml_toolbox-14.2.0}/ml_tools/ML_models.py

@@ -1,6 +1,6 @@
 import torch
 from torch import nn
-from typing import List, Union, Tuple, Dict, Any, Literal, Optional
+from typing import List, Union, Tuple, Dict, Any
 from pathlib import Path
 import json
 

dragon_ml_toolbox-14.2.0/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__)

{dragon_ml_toolbox-14.1.0 → dragon_ml_toolbox-14.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dragon-ml-toolbox"
-version = "14.1.0"
+version = "14.2.0"
 description = "A collection of tools for data science and machine learning projects."
 authors = [
     { name = "Karl L. Loza Vidaurre", email = "luigiloza@gmail.com" }
@@ -45,14 +45,11 @@ ML = [
     "torchmetrics",
 ]
 
-# pytorch-tabular API. Additionally Requires PyTorch with CUDA / MPS support
-# py-tab = [
-#     "pytorch_tabular",
-#     "pytorch-lightning",
-#     "wandb",
-#     "plotly",
-#     "captum",
-# ]
+# pytorch-tabular API for advanced models
+py-tab = [
+    "pytorch_tabular",
+    "omegaconf"
+]
 
 # MICE and VIF - Requires a new virtual-env due to dependency version conflicts
 mice = [

dragon_ml_toolbox-14.1.0/ml_tools/_ML_pytorch_tabular.py

@@ -1,543 +0,0 @@
-import torch
-import pandas as pd
-import numpy as np
-import matplotlib.pyplot as plt
-from typing import List, Literal, Union, Optional, Dict, Any
-from pathlib import Path
-import warnings
-
-# --- Third-party imports ---
-try:
-    from pytorch_tabular.models.common.heads import LinearHeadConfig
-    from pytorch_tabular.config import (
-        DataConfig,
-        ModelConfig,
-        OptimizerConfig,
-        TrainerConfig,
-        ExperimentConfig,
-    )
-    from pytorch_tabular.models import (
-        CategoryEmbeddingModelConfig,
-        TabNetModelConfig,
-        TabTransformerConfig,
-        FTTransformerConfig,
-        AutoIntConfig,
-        NodeConfig,
-        GANDALFConfig
-    )
-    from pytorch_tabular.tabular_model import TabularModel
-except ImportError:
-    print("----------------------------------------------------------------")
-    print("ERROR: `pytorch-tabular` is not installed.")
-    print("Please install it to use the models in this script:")
-    print('\npip install "dragon-ml-toolbox[py-tab]"')
-    print("----------------------------------------------------------------")
-    raise
-
-# --- Local ML-Tools imports ---
-from ._logger import _LOGGER
-from ._script_info import _script_info
-from ._schema import FeatureSchema
-from .path_manager import make_fullpath, sanitize_filename
-from .keys import SHAPKeys
-from .ML_datasetmaster import _PytorchDataset
-from .ML_evaluation import (
-    classification_metrics, 
-    regression_metrics
-)
-from .ML_evaluation_multi import (
-    multi_target_regression_metrics, 
-    multi_label_classification_metrics
-)
-
-
-__all__ = [
-    "PyTabularTrainer"
-]
-
-
-# --- Model Configuration Mapping ---
-# Maps a simple string name to the required ModelConfig class
-SUPPORTED_MODELS: Dict[str, Any] = {
-    "TabNet": TabNetModelConfig,
-    "TabTransformer": TabTransformerConfig,
-    "FTTransformer": FTTransformerConfig,
-    "AutoInt": AutoIntConfig,
-    "NODE": NodeConfig,
-    "GATE": GANDALFConfig, # Gated Additive Tree Ensemble
-    "CategoryEmbedding": CategoryEmbeddingModelConfig, # A basic MLP
-}
-
-
-class PyTabularTrainer:
-    """
-    A wrapper for models from the `pytorch-tabular` library, designed to be
-    compatible with the `dragon-ml-toolbox` ecosystem.
-    
-    This class acts as a high-level trainer that adapts the `ML_datasetmaster`
-    datasets into the format required by `pytorch-tabular` and routes
-    evaluation results to the standard `ML_evaluation` functions.
-    
-    It handles:
-    - Automatic `DataConfig` creation from a `FeatureSchema`.
-    - Model and Trainer configuration.
-    - Training and evaluation.
-    - SHAP explanations.
-    """
-    
-    def __init__(self,
-                 schema: FeatureSchema,
-                 target_names: List[str],
-                 kind: Literal["regression", "classification", "multi_target_regression", "multi_label_classification"],
-                 model_name: str,
-                 model_config_params: Optional[Dict[str, Any]] = None,
-                 optimizer_config_params: Optional[Dict[str, Any]] = None,
-                 trainer_config_params: Optional[Dict[str, Any]] = None):
-        """
-        Initializes the Model, Data, and Trainer configurations.
-
-        Args:
-            schema (FeatureSchema): 
-                The definitive schema object from data_exploration.
-            target_names (List[str]): 
-                A list of target column names.
-            kind (Literal[...]): 
-                The type of ML task. This is used to set the `pytorch-tabular`
-                task and to route to the correct evaluation function.
-            model_name (str): 
-                The name of the model to use. Must be one of:
-                "TabNet", "TabTransformer", "FTTransformer", "AutoInt",
-                "NODE", "GATE", "CategoryEmbedding".
-            model_config_params (Dict, optional): 
-                Overrides for the chosen model's `ModelConfig`.
-                (e.g., `{"n_d": 16, "n_a": 16}` for TabNet).
-            optimizer_config_params (Dict, optional): 
-                Overrides for the `OptimizerConfig` (e.g., `{"lr": 0.005}`).
-            trainer_config_params (Dict, optional): 
-                Overrides for the `TrainerConfig` (e.g., `{"max_epochs": 100}`).
-        """
-        _LOGGER.info(f"Initializing PyTabularTrainer for model: {model_name}")
-        
-        # --- 1. Store key info ---
-        self.schema = schema
-        self.target_names = target_names
-        self.kind = kind
-        self.model_name = model_name
-        self._is_fitted = False
-
-        if model_name not in SUPPORTED_MODELS:
-            _LOGGER.error(f"Model '{model_name}' is not supported. Choose from: {list(SUPPORTED_MODELS.keys())}")
-            raise ValueError(f"Unsupported model: {model_name}")
-
-        # --- 2. Map ML-Tools 'kind' to pytorch-tabular 'task' ---
-        if kind == "regression":
-            self.task = "regression"
-            self._pt_target_names = target_names
-        elif kind == "classification":
-            self.task = "classification"
-            self._pt_target_names = target_names
-        elif kind == "multi_target_regression":
-            self.task = "multi-label-regression" # pytorch-tabular's name
-            self._pt_target_names = target_names
-        elif kind == "multi_label_classification":
-            self.task = "multi-label-classification"
-            self._pt_target_names = target_names
-        else:
-            _LOGGER.error(f"Unknown task 'kind': {kind}")
-            raise ValueError()
-
-        # --- 3. Create DataConfig from FeatureSchema ---
-        # Note: pytorch-tabular handles scaling internally
-        self.data_config = DataConfig(
-            target=self._pt_target_names,
-            continuous_cols=list(schema.continuous_feature_names),
-            categorical_cols=list(schema.categorical_feature_names),
-            continuous_feature_transform="quantile_normal", 
-        )
-
-        # --- 4. Create ModelConfig ---
-        model_config_class = SUPPORTED_MODELS[model_name]
-        
-        # Apply user overrides
-        if model_config_params is None:
-            model_config_params = {}
-            
-        # Set task in params
-        model_config_params["task"] = self.task
-        
-        # Handle multi-target output for regression
-        if self.task == "multi-label-regression":
-            # Must configure the model's output head
-            if "head" not in model_config_params:
-                _LOGGER.info("Configuring model head for multi-target regression.")
-                model_config_params["head"] = "LinearHead"
-                model_config_params["head_config"] = {
-                    "layers": "", # No hidden layers in the head
-                    "output_dim": len(self.target_names)
-                }
-
-        self.model_config = model_config_class(**model_config_params)
-
-        # --- 5. Create OptimizerConfig ---
-        if optimizer_config_params is None:
-            optimizer_config_params = {}
-        self.optimizer_config = OptimizerConfig(**optimizer_config_params)
-
-        # --- 6. Create TrainerConfig ---
-        if trainer_config_params is None:
-            trainer_config_params = {}
-        
-        # Default to GPU if available
-        if "accelerator" not in trainer_config_params:
-            if torch.cuda.is_available():
-                trainer_config_params["accelerator"] = "cuda"
-            elif torch.backends.mps.is_available():
-                 trainer_config_params["accelerator"] = "mps"
-            else:
-                 trainer_config_params["accelerator"] = "cpu"
-        
-        # Set other sensible defaults
-        if "checkpoints" not in trainer_config_params:
-            trainer_config_params["checkpoints"] = "val_loss"
-            trainer_config_params["load_best_at_end"] = True
-            
-        if "early_stopping" not in trainer_config_params:
-            trainer_config_params["early_stopping"] = "val_loss"
-        
-        self.trainer_config = TrainerConfig(**trainer_config_params)
-
-        # --- 7. Instantiate the TabularModel ---
-        self.tabular_model = TabularModel(
-            data_config=self.data_config,
-            model_config=self.model_config,
-            optimizer_config=self.optimizer_config,
-            trainer_config=self.trainer_config,
-        )
-
-    def _dataset_to_dataframe(self, dataset: _PytorchDataset) -> pd.DataFrame:
-        """Converts an _PytorchDataset back into a pandas DataFrame."""
-        try:
-            features_np = dataset.features.cpu().numpy()
-            labels_np = dataset.labels.cpu().numpy()
-            feature_names = dataset.feature_names
-            target_names = dataset.target_names
-        except Exception as e:
-            _LOGGER.error(f"Failed to extract data from provided dataset: {e}")
-            raise
-            
-        # Create features DataFrame
-        df = pd.DataFrame(features_np, columns=feature_names)
-        
-        # Add labels
-        if labels_np.ndim == 1:
-            df[target_names[0]] = labels_np
-        elif labels_np.ndim == 2:
-            for i, name in enumerate(target_names):
-                df[name] = labels_np[:, i]
-        
-        return df
-
-    def fit(self, 
-            train_dataset: _PytorchDataset, 
-            test_dataset: _PytorchDataset, 
-            epochs: int = 20, 
-            batch_size: int = 10):
-        """
-        Trains the model using the provided datasets.
-
-        Args:
-            train_dataset (_PytorchDataset): The training dataset.
-            test_dataset (_PytorchDataset): The validation dataset.
-            epochs (int): The number of epochs to train for.
-            batch_size (int): The batch size.
-        """
-        _LOGGER.info(f"Converting datasets to pandas DataFrame for {self.model_name}...")
-        train_df = self._dataset_to_dataframe(train_dataset)
-        test_df = self._dataset_to_dataframe(test_dataset)
-        
-        _LOGGER.info(f"Starting training for {epochs} epochs...")
-        with warnings.catch_warnings():
-            # Suppress abundant pytorch-lightning warnings
-            warnings.simplefilter("ignore")
-            self.tabular_model.fit(
-                train=train_df,
-                validation=test_df,
-                max_epochs=epochs
-            )
-            
-        self._is_fitted = True
-        _LOGGER.info("Training complete.")
-
-    def evaluate(self, 
-                 save_dir: Union[str, Path], 
-                 data: _PytorchDataset, 
-                 classification_threshold: float = 0.5):
-        """
-        Evaluates the model and saves reports using the standard ML_evaluation functions.
-        
-        Args:
-            save_dir (str | Path): Directory to save all reports and plots.
-            data (_PytorchDataset): The data to evaluate on.
-            classification_threshold (float): Threshold for multi-label tasks.
-        """
-        if not self._is_fitted:
-            _LOGGER.error("Model is not fitted. Call .fit() first.")
-            raise RuntimeError()
-            
-        print("\n--- Model Evaluation (PyTorch-Tabular) ---")
-        
-        eval_df = self._dataset_to_dataframe(data)
-        
-        # Get raw predictions from pytorch-tabular
-        raw_preds_df = self.tabular_model.predict(
-            eval_df, 
-            include_input_features=False
-        )
-        
-        # Extract y_true from the dataframe
-        y_true = eval_df[self.target_names].to_numpy()
-        
-        y_pred = None
-        y_prob = None
-        
-        # --- Route based on task kind ---
-        
-        if self.kind == "regression":
-            pred_col_name = f"{self.target_names[0]}_prediction"
-            y_pred = raw_preds_df[pred_col_name].to_numpy()
-            regression_metrics(y_true.flatten(), y_pred.flatten(), save_dir)
-            
-        elif self.kind == "classification":
-            y_pred = raw_preds_df["prediction"].to_numpy()
-            # Get class names from the model's datamodule
-            if self.tabular_model.datamodule is None:
-                _LOGGER.error("Model's datamodule is not initialized. Cannot extract class names for probabilities.")
-                raise RuntimeError("Datamodule not found. Was the model trained or loaded correctly?")
-            class_names = self.tabular_model.datamodule.data_config.target_classes[self.target_names[0]]
-            prob_cols = [f"{c}_probability" for c in class_names]
-            y_prob = raw_preds_df[prob_cols].values
-            classification_metrics(save_dir, y_true.flatten(), y_pred, y_prob)
-            
-        elif self.kind == "multi_target_regression":
-            pred_cols = [f"{name}_prediction" for name in self.target_names]
-            y_pred = raw_preds_df[pred_cols].to_numpy()
-            multi_target_regression_metrics(y_true, y_pred, self.target_names, save_dir)
-            
-        elif self.kind == "multi_label_classification":
-            prob_cols = [f"{name}_probability" for name in self.target_names]
-            y_prob = raw_preds_df[prob_cols].to_numpy()
-            # y_pred is derived from y_prob
-            multi_label_classification_metrics(y_true, y_prob, self.target_names, save_dir, classification_threshold)
-            
-    def explain(self, 
-                save_dir: Union[str, Path], 
-                explain_dataset: _PytorchDataset):
-        """
-        Generates SHAP explanations and saves plots and summary CSVs.
-        
-        This method uses pytorch-tabular's internal `.explain()` method
-        and then formats the output to match the ML_evaluation standard.
-
-        Args:
-            save_dir (str | Path): Directory to save all SHAP artifacts.
-            explain_dataset (_PytorchDataset): The dataset to explain.
-        """
-        if not self._is_fitted:
-            _LOGGER.error("Model is not fitted. Call .fit() first.")
-            raise RuntimeError()
-
-        print(f"\n--- SHAP Value Explanation ({self.model_name}) ---")
-        
-        explain_df = self._dataset_to_dataframe(explain_dataset)
-        
-        # We must use the dataframe *without* the target columns for explanation
-        feature_df: pd.DataFrame = explain_df[self.schema.feature_names] # type: ignore
-        
-        # This returns a DataFrame (single-target) or Dict[str, DataFrame]
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            shap_output = self.tabular_model.explain(feature_df)
-        
-        save_dir_path = make_fullpath(save_dir, make=True, enforce="directory")
-        plt.ioff()
-        
-        # --- 1. Handle single-target (regression/classification) ---
-        if isinstance(shap_output, pd.DataFrame):
-            # shap_output is (n_samples, n_features)
-            shap_values = shap_output.to_numpy()
-            
-            # Save Bar Plot
-            self._save_shap_plots(
-                shap_values=shap_values,
-                instances_df=feature_df,
-                save_dir=save_dir_path,
-                suffix="" # No suffix for single target
-            )
-            # Save Summary Data
-            self._save_shap_csv(
-                shap_values=shap_values,
-                feature_names=list(self.schema.feature_names),
-                save_dir=save_dir_path,
-                suffix=""
-            )
-        
-        # --- 2. Handle multi-target ---
-        elif isinstance(shap_output, dict):
-            for target_name, shap_df in shap_output.items(): # type: ignore
-                _LOGGER.info(f"  -> Generating SHAP plots for target: '{target_name}'")
-                shap_values = shap_df.values
-                sanitized_name = sanitize_filename(target_name)
-                
-                # Save Bar Plot
-                self._save_shap_plots(
-                    shap_values=shap_values,
-                    instances_df=feature_df,
-                    save_dir=save_dir_path,
-                    suffix=f"_{sanitized_name}",
-                    title_suffix=f" for '{target_name}'"
-                )
-                # Save Summary Data
-                self._save_shap_csv(
-                    shap_values=shap_values,
-                    feature_names=list(self.schema.feature_names),
-                    save_dir=save_dir_path,
-                    suffix=f"_{sanitized_name}"
-                )
-
-        plt.ion()
-        _LOGGER.info(f"All SHAP plots saved to '{save_dir_path.name}'")
-
-    def _save_shap_plots(self, shap_values: np.ndarray, 
-                         instances_df: pd.DataFrame, 
-                         save_dir: Path, 
-                         suffix: str = "",
-                         title_suffix: str = ""):
-        """Internal helper to save standard SHAP plots."""
-        try:
-            import shap
-        except ImportError:
-            _LOGGER.error("`shap` is required for plotting. Please install it: pip install shap")
-            return
-            
-        # Save Bar Plot
-        bar_path = save_dir / f"shap_bar_plot{suffix}.svg"
-        shap.summary_plot(shap_values, instances_df, plot_type="bar", show=False)
-        ax = plt.gca()
-        ax.set_xlabel("SHAP Value Impact", labelpad=10)
-        plt.title(f"SHAP Feature Importance{title_suffix}")
-        plt.tight_layout()
-        plt.savefig(bar_path)
-        plt.close()
-
-        # Save Dot Plot
-        dot_path = save_dir / f"shap_dot_plot{suffix}.svg"
-        shap.summary_plot(shap_values, instances_df, plot_type="dot", show=False)
-        ax = plt.gca()
-        ax.set_xlabel("SHAP Value Impact", labelpad=10)
-        if plt.gcf().axes and len(plt.gcf().axes) > 1:
-            cb = plt.gcf().axes[-1]
-            cb.set_ylabel("", size=1)
-        plt.title(f"SHAP Feature Importance{title_suffix}")
-        plt.tight_layout()
-        plt.savefig(dot_path)
-        plt.close()
-
-    def _save_shap_csv(self, shap_values: np.ndarray, 
-                       feature_names: List[str], 
-                       save_dir: Path, 
-                       suffix: str = ""):
-        """Internal helper to save standard SHAP summary CSV."""
-        
-        shap_summary_filename = f"{SHAPKeys.SAVENAME}{suffix}.csv"
-        summary_path = save_dir / shap_summary_filename
-        
-        # Handle multi-class (list of arrays) vs. regression (single array)
-        if isinstance(shap_values, list):
-            mean_abs_shap = np.abs(np.stack(shap_values)).mean(axis=0).mean(axis=0)
-        else:
-            mean_abs_shap = np.abs(shap_values).mean(axis=0)
-
-        mean_abs_shap = mean_abs_shap.flatten()
-            
-        summary_df = pd.DataFrame({
-            SHAPKeys.FEATURE_COLUMN: feature_names,
-            SHAPKeys.SHAP_VALUE_COLUMN: mean_abs_shap
-        }).sort_values(SHAPKeys.SHAP_VALUE_COLUMN, ascending=False)
-        
-        summary_df.to_csv(summary_path, index=False)
-
-    def save_model(self, directory: Union[str, Path]):
-        """
-        Saves the entire trained model, configuration, and datamodule
-        to a directory.
-
-        Args:
-            directory (str | Path): The directory to save the model.
-                                    The directory will be created.
-        """
-        if not self._is_fitted:
-            _LOGGER.error("Cannot save a model that has not been fitted.")
-            return
-            
-        save_path = make_fullpath(directory, make=True, enforce="directory")
-        self.tabular_model.save_model(str(save_path))
-        _LOGGER.info(f"Model saved to '{save_path.name}'")
-
-    @classmethod
-    def load_model(cls,
-                   directory: Union[str, Path],
-                   schema: FeatureSchema,
-                   target_names: List[str],
-                   kind: Literal["regression", "classification", "multi_target_regression", "multi_label_classification"]
-                   ) -> 'PyTabularTrainer':
-        """
-        Loads a saved model and reconstructs the PyTabularTrainer wrapper.
-        
-        Note: The schema, target_names, and kind must be provided again
-        as they are not serialized by pytorch-tabular.
-
-        Args:
-            directory (str | Path): The directory from which to load the model.
-            schema (FeatureSchema): The schema used during original training.
-            target_names (List[str]): The target names used during original training.
-            kind (Literal[...]): The task 'kind' used during original training.
-
-        Returns:
-            PyTabularTrainer: A new instance of the trainer with the loaded model.
-        """
-        load_path = make_fullpath(directory, enforce="directory")
-        
-        _LOGGER.info(f"Loading model from '{load_path.name}'...")
-        
-        # Load the internal pytorch-tabular model
-        loaded_tabular_model = TabularModel.load_model(str(load_path))
-        
-        if loaded_tabular_model.model is None:
-             _LOGGER.error("Loaded model's internal '.model' attribute is None. Load failed.")
-             raise RuntimeError("Loaded model is incomplete.")
-        
-        model_name = loaded_tabular_model.model._model_name
-        
-        if model_name.startswith("GANDALF"): # Handle GANDALF's dynamic name
-            model_name = "GATE"
-        
-        # Re-create the wrapper
-        wrapper = cls(
-            schema=schema,
-            target_names=target_names,
-            kind=kind,
-            model_name=model_name 
-            # Configs are already part of the loaded_tabular_model
-            # We just need to pass the minimum to the __init__
-        )
-        
-        # Overwrite the un-trained model with the loaded trained model
-        wrapper.tabular_model = loaded_tabular_model
-        wrapper._is_fitted = True
-        
-        _LOGGER.info(f"Successfully loaded '{model_name}' model.")
-        return wrapper
-
-
-def info():
-    _script_info(__all__)