autogluon.tabular 1.5.0b20251228__py3-none-any.whl → 1.5.1b20260116__py3-none-any.whl

autogluon/tabular/models/mitra/sklearn_interface.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
+import contextlib
 import os
 import time
 from pathlib import Path
-import contextlib
 
 import numpy as np
 import pandas as pd
@@ -18,34 +18,36 @@ from ._internal.models.tab2d import Tab2D
 from ._internal.utils.set_seed import set_seed
 
 # Hyperparameter search space
-DEFAULT_FINE_TUNE = True # [True, False]
-DEFAULT_FINE_TUNE_STEPS = 50 # [50, 60, 70, 80, 90, 100]
-DEFAULT_CLS_METRIC = 'log_loss' # ['log_loss', 'accuracy', 'auc']
-DEFAULT_REG_METRIC = 'mse' # ['mse', 'mae', 'rmse', 'r2']
-SHUFFLE_CLASSES = False # [True, False]
-SHUFFLE_FEATURES = False # [True, False]
-USE_RANDOM_TRANSFORMS = False # [True, False]
-RANDOM_MIRROR_REGRESSION = True # [True, False]
-RANDOM_MIRROR_X = True # [True, False]
-LR = 0.0001 # [0.00001, 0.000025, 0.00005, 0.000075, 0.0001, 0.00025, 0.0005, 0.00075, 0.001]
-PATIENCE = 40 # [30, 35, 40, 45, 50]
-WARMUP_STEPS = 1000 # [500, 750, 1000, 1250, 1500]
-DEFAULT_CLS_MODEL = 'autogluon/mitra-classifier'
-DEFAULT_REG_MODEL = 'autogluon/mitra-regressor'
+DEFAULT_FINE_TUNE = True  # [True, False]
+DEFAULT_FINE_TUNE_STEPS = 50  # [50, 60, 70, 80, 90, 100]
+DEFAULT_CLS_METRIC = "log_loss"  # ['log_loss', 'accuracy', 'auc']
+DEFAULT_REG_METRIC = "mse"  # ['mse', 'mae', 'rmse', 'r2']
+SHUFFLE_CLASSES = False  # [True, False]
+SHUFFLE_FEATURES = False  # [True, False]
+USE_RANDOM_TRANSFORMS = False  # [True, False]
+RANDOM_MIRROR_REGRESSION = True  # [True, False]
+RANDOM_MIRROR_X = True  # [True, False]
+LR = 0.0001  # [0.00001, 0.000025, 0.00005, 0.000075, 0.0001, 0.00025, 0.0005, 0.00075, 0.001]
+PATIENCE = 40  # [30, 35, 40, 45, 50]
+WARMUP_STEPS = 1000  # [500, 750, 1000, 1250, 1500]
+DEFAULT_CLS_MODEL = "autogluon/mitra-classifier"
+DEFAULT_REG_MODEL = "autogluon/mitra-regressor"
 
 # Constants
 SEED = 0
 DEFAULT_MODEL_TYPE = "Tab2D"
 
+
 def _get_default_device():
     """Get the best available device for the current system."""
     if torch.cuda.is_available():
         return "cuda"
-    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
         return "mps"  # Apple silicon
     else:
         return "cpu"
 
+
 DEFAULT_DEVICE = _get_default_device()
 DEFAULT_ENSEMBLE = 1
 DEFAULT_DIM = 512
@@ -55,32 +57,34 @@ DEFAULT_CLASSES = 10
 DEFAULT_VALIDATION_SPLIT = 0.2
 USE_HF = True  # Use Hugging Face pretrained models if available
 
+
 class MitraBase(BaseEstimator):
     """Base class for Mitra models with common functionality."""
 
-    def __init__(self,
-            model_type=DEFAULT_MODEL_TYPE,
-            n_estimators=DEFAULT_ENSEMBLE,
-            device=DEFAULT_DEVICE,
-            fine_tune=DEFAULT_FINE_TUNE,
-            fine_tune_steps=DEFAULT_FINE_TUNE_STEPS,
-            metric=DEFAULT_CLS_METRIC,
-            state_dict=None,
-            hf_model=None,
-            patience=PATIENCE,
-            lr=LR,
-            warmup_steps=WARMUP_STEPS,
-            shuffle_classes=SHUFFLE_CLASSES,
-            shuffle_features=SHUFFLE_FEATURES,
-            use_random_transforms=USE_RANDOM_TRANSFORMS,
-            random_mirror_regression=RANDOM_MIRROR_REGRESSION,
-            random_mirror_x=RANDOM_MIRROR_X,
-            seed=SEED,
-            verbose=True,
-        ):
+    def __init__(
+        self,
+        model_type=DEFAULT_MODEL_TYPE,
+        n_estimators=DEFAULT_ENSEMBLE,
+        device=DEFAULT_DEVICE,
+        fine_tune=DEFAULT_FINE_TUNE,
+        fine_tune_steps=DEFAULT_FINE_TUNE_STEPS,
+        metric=DEFAULT_CLS_METRIC,
+        state_dict=None,
+        hf_model=None,
+        patience=PATIENCE,
+        lr=LR,
+        warmup_steps=WARMUP_STEPS,
+        shuffle_classes=SHUFFLE_CLASSES,
+        shuffle_features=SHUFFLE_FEATURES,
+        use_random_transforms=USE_RANDOM_TRANSFORMS,
+        random_mirror_regression=RANDOM_MIRROR_REGRESSION,
+        random_mirror_x=RANDOM_MIRROR_X,
+        seed=SEED,
+        verbose=True,
+    ):
         """
         Initialize the base Mitra model.
-        
+
         Parameters
         ----------
         model_type : str, default="Tab2D"
@@ -125,59 +129,62 @@ class MitraBase(BaseEstimator):
             model_name=ModelName.TAB2D,
             seed=self.seed,
             hyperparams={
-                'dim_embedding': None,
-                'early_stopping_data_split': 'VALID',
-                'early_stopping_max_samples': 2048,
-                'early_stopping_patience': self.patience,
-                'grad_scaler_enabled': False,
-                'grad_scaler_growth_interval': 1000,
-                'grad_scaler_scale_init': 65536.0,
-                'grad_scaler_scale_min': 65536.0,
-                'label_smoothing': 0.0,
-                'lr_scheduler': False,
-                'lr_scheduler_patience': 25,
-                'max_epochs': self.fine_tune_steps if self.fine_tune else 0,
-                'max_samples_query': 1024,
-                'max_samples_support': 8192,
-                'optimizer': 'adamw',
-                'lr': self.lr,
-                'weight_decay': 0.1,
-                'warmup_steps': self.warmup_steps,
-                'path_to_weights': self.state_dict,
-                'precision': 'bfloat16',
-                'random_mirror_regression': self.random_mirror_regression,
-                'random_mirror_x': self.random_mirror_x,
-                'shuffle_classes': self.shuffle_classes,
-                'shuffle_features': self.shuffle_features,
-                'use_random_transforms': self.use_random_transforms,
-                'use_feature_count_scaling': False,
-                'use_pretrained_weights': False,
-                'use_quantile_transformer': False,
-                'budget': time_limit,
-                'metric': self.metric,
+                "dim_embedding": None,
+                "early_stopping_data_split": "VALID",
+                "early_stopping_max_samples": 2048,
+                "early_stopping_patience": self.patience,
+                "grad_scaler_enabled": False,
+                "grad_scaler_growth_interval": 1000,
+                "grad_scaler_scale_init": 65536.0,
+                "grad_scaler_scale_min": 65536.0,
+                "label_smoothing": 0.0,
+                "lr_scheduler": False,
+                "lr_scheduler_patience": 25,
+                "max_epochs": self.fine_tune_steps if self.fine_tune else 0,
+                "max_samples_query": 1024,
+                "max_samples_support": 8192,
+                "optimizer": "adamw",
+                "lr": self.lr,
+                "weight_decay": 0.1,
+                "warmup_steps": self.warmup_steps,
+                "path_to_weights": self.state_dict,
+                "precision": "bfloat16",
+                "random_mirror_regression": self.random_mirror_regression,
+                "random_mirror_x": self.random_mirror_x,
+                "shuffle_classes": self.shuffle_classes,
+                "shuffle_features": self.shuffle_features,
+                "use_random_transforms": self.use_random_transforms,
+                "use_feature_count_scaling": False,
+                "use_pretrained_weights": False,
+                "use_quantile_transformer": False,
+                "budget": time_limit,
+                "metric": self.metric,
             },
         )
 
         cfg.task = task
-        cfg.hyperparams.update({
-            'n_ensembles': self.n_estimators,
-            'dim': DEFAULT_DIM,
-            'dim_output': dim_output,
-            'n_layers': DEFAULT_LAYERS,
-            'n_heads': DEFAULT_HEADS,
-            'regression_loss': 'mse',
-        })
+        cfg.hyperparams.update(
+            {
+                "n_ensembles": self.n_estimators,
+                "dim": DEFAULT_DIM,
+                "dim_output": dim_output,
+                "n_layers": DEFAULT_LAYERS,
+                "n_heads": DEFAULT_HEADS,
+                "regression_loss": "mse",
+            }
+        )
 
         return cfg, Tab2D
 
-
     def _split_data(self, X, y):
         """Split data into training and validation sets."""
-        if hasattr(self, 'task') and self.task == 'classification':
+        if hasattr(self, "task") and self.task == "classification":
             return make_stratified_dataset_split(X, y, seed=self.seed)
         else:
             # For regression, use random split
-            val_indices = np.random.choice(range(len(X)), int(DEFAULT_VALIDATION_SPLIT * len(X)), replace=False).tolist()
+            val_indices = np.random.choice(
+                range(len(X)), int(DEFAULT_VALIDATION_SPLIT * len(X)), replace=False
+            ).tolist()
             train_indices = [i for i in range(len(X)) if i not in val_indices]
             return X[train_indices], X[val_indices], y[train_indices], y[val_indices]
 
@@ -188,7 +195,9 @@ class MitraBase(BaseEstimator):
         rng = np.random.RandomState(cfg.seed)
 
         success = False
-        while not (success and cfg.hyperparams["max_samples_support"] > 0 and cfg.hyperparams["max_samples_query"] > 0):
+        while not (
+            success and cfg.hyperparams["max_samples_support"] > 0 and cfg.hyperparams["max_samples_query"] > 0
+        ):
             try:
                 self.trainers.clear()
 
@@ -199,16 +208,18 @@ class MitraBase(BaseEstimator):
                         model = Tab2D.from_pretrained(self.hf_model, device=self.device)
                     else:
                         model = Tab2D(
-                            dim=cfg.hyperparams['dim'],
+                            dim=cfg.hyperparams["dim"],
                             dim_output=dim_output,
-                            n_layers=cfg.hyperparams['n_layers'],
-                            n_heads=cfg.hyperparams['n_heads'],
+                            n_layers=cfg.hyperparams["n_layers"],
+                            n_heads=cfg.hyperparams["n_heads"],
                             task=task.upper(),
                             use_pretrained_weights=True,
                             path_to_weights=Path(self.state_dict),
                             device=self.device,
                         )
-                    trainer = TrainerFinetune(cfg, model, n_classes=n_classes, device=self.device, rng=rng, verbose=self.verbose)
+                    trainer = TrainerFinetune(
+                        cfg, model, n_classes=n_classes, device=self.device, rng=rng, verbose=self.verbose
+                    )
 
                     start_time = time.time()
                     trainer.train(X_train, y_train, X_valid, y_valid)
@@ -221,27 +232,25 @@ class MitraBase(BaseEstimator):
 
             except torch.cuda.OutOfMemoryError:
                 if cfg.hyperparams["max_samples_support"] >= 2048:
-                    cfg.hyperparams["max_samples_support"] = int(
-                        cfg.hyperparams["max_samples_support"] // 2
+                    cfg.hyperparams["max_samples_support"] = int(cfg.hyperparams["max_samples_support"] // 2)
+                    print(
+                        f"Reducing max_samples_support from {cfg.hyperparams['max_samples_support'] * 2}"
+                        f"to {cfg.hyperparams['max_samples_support']} due to OOM error."
                     )
-                    print(f"Reducing max_samples_support from {cfg.hyperparams['max_samples_support'] * 2}"
-                          f"to {cfg.hyperparams['max_samples_support']} due to OOM error.")
                 else:
-                    cfg.hyperparams["max_samples_support"] = int(
-                        cfg.hyperparams["max_samples_support"] // 2
+                    cfg.hyperparams["max_samples_support"] = int(cfg.hyperparams["max_samples_support"] // 2)
+                    print(
+                        f"Reducing max_samples_support from {cfg.hyperparams['max_samples_support'] * 2}"
+                        f"to {cfg.hyperparams['max_samples_support']} due to OOM error."
                     )
-                    print(f"Reducing max_samples_support from {cfg.hyperparams['max_samples_support'] * 2}"
-                          f"to {cfg.hyperparams['max_samples_support']} due to OOM error.")
-                    cfg.hyperparams["max_samples_query"] = int(
-                        cfg.hyperparams["max_samples_query"] // 2
+                    cfg.hyperparams["max_samples_query"] = int(cfg.hyperparams["max_samples_query"] // 2)
+                    print(
+                        f"Reducing max_samples_query from {cfg.hyperparams['max_samples_query'] * 2}"
+                        f"to {cfg.hyperparams['max_samples_query']} due to OOM error."
                     )
-                    print(f"Reducing max_samples_query from {cfg.hyperparams['max_samples_query'] * 2}"
-                          f"to {cfg.hyperparams['max_samples_query']} due to OOM error.")
 
         if not success:
-            raise RuntimeError(
-                "Failed to train Mitra model after multiple attempts due to out of memory error."
-            )
+            raise RuntimeError("Failed to train Mitra model after multiple attempts due to out of memory error.")
 
         return self
 
@@ -249,26 +258,27 @@ class MitraBase(BaseEstimator):
 class MitraClassifier(MitraBase, ClassifierMixin):
     """Classifier implementation of Mitra model."""
 
-    def __init__(self,
-            model_type=DEFAULT_MODEL_TYPE,
-            n_estimators=DEFAULT_ENSEMBLE,
-            device=DEFAULT_DEVICE,
-            fine_tune=DEFAULT_FINE_TUNE,
-            fine_tune_steps=DEFAULT_FINE_TUNE_STEPS,
-            metric=DEFAULT_CLS_METRIC,
-            state_dict=None,
-            hf_model=DEFAULT_CLS_MODEL,
-            patience=PATIENCE,
-            lr=LR,
-            warmup_steps=WARMUP_STEPS,
-            shuffle_classes=SHUFFLE_CLASSES,
-            shuffle_features=SHUFFLE_FEATURES,
-            use_random_transforms=USE_RANDOM_TRANSFORMS,
-            random_mirror_regression=RANDOM_MIRROR_REGRESSION,
-            random_mirror_x=RANDOM_MIRROR_X,
-            seed=SEED,
-            verbose=True,
-        ):
+    def __init__(
+        self,
+        model_type=DEFAULT_MODEL_TYPE,
+        n_estimators=DEFAULT_ENSEMBLE,
+        device=DEFAULT_DEVICE,
+        fine_tune=DEFAULT_FINE_TUNE,
+        fine_tune_steps=DEFAULT_FINE_TUNE_STEPS,
+        metric=DEFAULT_CLS_METRIC,
+        state_dict=None,
+        hf_model=DEFAULT_CLS_MODEL,
+        patience=PATIENCE,
+        lr=LR,
+        warmup_steps=WARMUP_STEPS,
+        shuffle_classes=SHUFFLE_CLASSES,
+        shuffle_features=SHUFFLE_FEATURES,
+        use_random_transforms=USE_RANDOM_TRANSFORMS,
+        random_mirror_regression=RANDOM_MIRROR_REGRESSION,
+        random_mirror_x=RANDOM_MIRROR_X,
+        seed=SEED,
+        verbose=True,
+    ):
         """Initialize the classifier."""
         super().__init__(
             model_type,
@@ -290,19 +300,19 @@ class MitraClassifier(MitraBase, ClassifierMixin):
             seed=seed,
             verbose=verbose,
         )
-        self.task = 'classification'
+        self.task = "classification"
 
-    def fit(self, X, y, X_val = None, y_val = None, time_limit = None):
+    def fit(self, X, y, X_val=None, y_val=None, time_limit=None):
         """
         Fit the ensemble of models.
-        
+
         Parameters
         ----------
         X : array-like of shape (n_samples, n_features)
             Training data
         y : array-like of shape (n_samples,)
             Target values
-            
+
         Returns
         -------
         self : object
@@ -310,7 +320,6 @@ class MitraClassifier(MitraBase, ClassifierMixin):
         """
 
         with mitra_deterministic_context():
-
             if isinstance(X, pd.DataFrame):
                 X = X.values
             if isinstance(y, pd.Series):
@@ -327,17 +336,26 @@ class MitraClassifier(MitraBase, ClassifierMixin):
             else:
                 X_train, X_valid, y_train, y_valid = self._split_data(X, y)
 
-            return self._train_ensemble(X_train, y_train, X_valid, y_valid, self.task, DEFAULT_CLASSES, n_classes=DEFAULT_CLASSES, time_limit=time_limit)
+            return self._train_ensemble(
+                X_train,
+                y_train,
+                X_valid,
+                y_valid,
+                self.task,
+                DEFAULT_CLASSES,
+                n_classes=DEFAULT_CLASSES,
+                time_limit=time_limit,
+            )
 
     def predict(self, X):
         """
         Predict class labels for samples in X.
-        
+
         Parameters
         ----------
         X : array-like of shape (n_samples, n_features)
             The input samples
-            
+
         Returns
         -------
         y : ndarray of shape (n_samples,)
@@ -352,12 +370,12 @@ class MitraClassifier(MitraBase, ClassifierMixin):
     def predict_proba(self, X):
         """
         Predict class probabilities for samples in X.
-        
+
         Parameters
         ----------
         X : array-like of shape (n_samples, n_features)
             The input samples
-            
+
         Returns
         -------
         p : ndarray of shape (n_samples, n_classes)
@@ -365,14 +383,13 @@ class MitraClassifier(MitraBase, ClassifierMixin):
         """
 
         with mitra_deterministic_context():
-
             if isinstance(X, pd.DataFrame):
                 X = X.values
 
             preds = []
             for trainer in self.trainers:
-                logits = trainer.predict(self.X, self.y, X)[...,:len(np.unique(self.y))] # Remove extra classes
-                preds.append(np.exp(logits) / np.exp(logits).sum(axis=1, keepdims=True)) # Softmax
+                logits = trainer.predict(self.X, self.y, X)[..., : len(np.unique(self.y))]  # Remove extra classes
+                preds.append(np.exp(logits) / np.exp(logits).sum(axis=1, keepdims=True))  # Softmax
             preds = sum(preds) / len(preds)  # Averaging ensemble predictions
 
             return preds
@@ -381,26 +398,27 @@ class MitraClassifier(MitraBase, ClassifierMixin):
 class MitraRegressor(MitraBase, RegressorMixin):
     """Regressor implementation of Mitra model."""
 
-    def __init__(self,
-            model_type=DEFAULT_MODEL_TYPE,
-            n_estimators=DEFAULT_ENSEMBLE,
-            device=DEFAULT_DEVICE,
-            fine_tune=DEFAULT_FINE_TUNE,
-            fine_tune_steps=DEFAULT_FINE_TUNE_STEPS,
-            metric=DEFAULT_REG_METRIC,
-            state_dict=None,
-            hf_model=DEFAULT_REG_MODEL,
-            patience=PATIENCE,
-            lr=LR,
-            warmup_steps=WARMUP_STEPS,
-            shuffle_classes=SHUFFLE_CLASSES,
-            shuffle_features=SHUFFLE_FEATURES,
-            use_random_transforms=USE_RANDOM_TRANSFORMS,
-            random_mirror_regression=RANDOM_MIRROR_REGRESSION,
-            random_mirror_x=RANDOM_MIRROR_X,
-            seed=SEED,
-            verbose=True,
-        ):
+    def __init__(
+        self,
+        model_type=DEFAULT_MODEL_TYPE,
+        n_estimators=DEFAULT_ENSEMBLE,
+        device=DEFAULT_DEVICE,
+        fine_tune=DEFAULT_FINE_TUNE,
+        fine_tune_steps=DEFAULT_FINE_TUNE_STEPS,
+        metric=DEFAULT_REG_METRIC,
+        state_dict=None,
+        hf_model=DEFAULT_REG_MODEL,
+        patience=PATIENCE,
+        lr=LR,
+        warmup_steps=WARMUP_STEPS,
+        shuffle_classes=SHUFFLE_CLASSES,
+        shuffle_features=SHUFFLE_FEATURES,
+        use_random_transforms=USE_RANDOM_TRANSFORMS,
+        random_mirror_regression=RANDOM_MIRROR_REGRESSION,
+        random_mirror_x=RANDOM_MIRROR_X,
+        seed=SEED,
+        verbose=True,
+    ):
         """Initialize the regressor."""
         super().__init__(
             model_type,
@@ -422,19 +440,19 @@ class MitraRegressor(MitraBase, RegressorMixin):
             seed=seed,
             verbose=verbose,
         )
-        self.task = 'regression'
+        self.task = "regression"
 
-    def fit(self, X, y, X_val = None, y_val = None, time_limit = None):
+    def fit(self, X, y, X_val=None, y_val=None, time_limit=None):
         """
         Fit the ensemble of models.
-        
+
         Parameters
         ----------
         X : array-like of shape (n_samples, n_features)
             Training data
         y : array-like of shape (n_samples,)
             Target values
-            
+
         Returns
         -------
         self : object
@@ -442,7 +460,6 @@ class MitraRegressor(MitraBase, RegressorMixin):
         """
 
         with mitra_deterministic_context():
-
             if isinstance(X, pd.DataFrame):
                 X = X.values
             if isinstance(y, pd.Series):
@@ -464,12 +481,12 @@ class MitraRegressor(MitraBase, RegressorMixin):
     def predict(self, X):
         """
         Predict regression target for samples in X.
-        
+
         Parameters
         ----------
         X : array-like of shape (n_samples, n_features)
             The input samples
-            
+
         Returns
         -------
         y : ndarray of shape (n_samples,)
@@ -477,16 +494,15 @@ class MitraRegressor(MitraBase, RegressorMixin):
         """
 
         with mitra_deterministic_context():
-
             if isinstance(X, pd.DataFrame):
                 X = X.values
-            
+
             preds = []
             for trainer in self.trainers:
                 preds.append(trainer.predict(self.X, self.y, X))
-        
+
             return sum(preds) / len(preds)  # Averaging ensemble predictions
-    
+
 
 @contextlib.contextmanager
 def mitra_deterministic_context():