autogluon.tabular 1.5.1b20260105__py3-none-any.whl → 1.5.1b20260116__py3-none-any.whl

autogluon/tabular/models/tabpfnmix/_internal/core/trainer_finetune.py

@@ -6,12 +6,15 @@ import time
 import einops
 import numpy as np
 import torch
-from sklearn.base import BaseEstimator
-
 from numpy.random import Generator
+from sklearn.base import BaseEstimator
 
 from autogluon.core.metrics import Scorer
 
+from ..config.config_run import ConfigRun
+from ..data.dataset_finetune import DatasetFinetune, DatasetFinetuneGenerator
+from ..data.preprocessor import Preprocessor
+from ..results.prediction_metrics import PredictionMetrics
 from .callbacks import Checkpoint, EarlyStopping, TrackOutput
 from .collator import CollatorWithPadding
 from .enums import Task
@@ -19,16 +22,11 @@ from .get_loss import get_loss
 from .get_optimizer import get_optimizer
 from .get_scheduler import get_scheduler
 from .y_transformer import create_y_transformer
-from ..config.config_run import ConfigRun
-from ..data.dataset_finetune import DatasetFinetune, DatasetFinetuneGenerator
-from ..data.preprocessor import Preprocessor
-from ..results.prediction_metrics import PredictionMetrics
 
 logger = logging.getLogger(__name__)
 
 
 class TrainerFinetune(BaseEstimator):
-
     def __init__(
         self,
         cfg: ConfigRun,
@@ -42,18 +40,18 @@ class TrainerFinetune(BaseEstimator):
         self.model = model
         self.model.to(self.cfg.device)
         self.n_classes = n_classes
-        
+
         self.loss = get_loss(self.cfg.task)
         self.optimizer = get_optimizer(self.cfg.hyperparams, self.model)
         self.scheduler = get_scheduler(self.cfg.hyperparams, self.optimizer)
         self.use_best_epoch = use_best_epoch
         self.compute_train_metrics = compute_train_metrics
 
-        self.early_stopping = EarlyStopping(patience=self.cfg.hyperparams['early_stopping_patience'])
-        self.preprocessor = Preprocessor( 
-            use_quantile_transformer=self.cfg.hyperparams['use_quantile_transformer'],
-            use_feature_count_scaling=self.cfg.hyperparams['use_feature_count_scaling'],
-            max_features=self.cfg.hyperparams['n_features'],
+        self.early_stopping = EarlyStopping(patience=self.cfg.hyperparams["early_stopping_patience"])
+        self.preprocessor = Preprocessor(
+            use_quantile_transformer=self.cfg.hyperparams["use_quantile_transformer"],
+            use_feature_count_scaling=self.cfg.hyperparams["use_feature_count_scaling"],
+            max_features=self.cfg.hyperparams["n_features"],
             task=self.cfg.task,
         )
 
@@ -70,7 +68,14 @@ class TrainerFinetune(BaseEstimator):
         self.optimizer = get_optimizer(self.cfg.hyperparams, self.model)
         self.scheduler = get_scheduler(self.cfg.hyperparams, self.optimizer)
 
-    def train(self, x_train: np.ndarray, y_train: np.ndarray, x_val: np.ndarray = None, y_val: np.ndarray = None, time_limit: float = None):
+    def train(
+        self,
+        x_train: np.ndarray,
+        y_train: np.ndarray,
+        x_val: np.ndarray = None,
+        y_val: np.ndarray = None,
+        time_limit: float = None,
+    ):
         time_start = time.time()
         if self.optimizer is None:
             self.reset_optimizer()
@@ -86,27 +91,27 @@ class TrainerFinetune(BaseEstimator):
         if use_val:
             x_val = self.preprocessor.transform(x_val)
         self.y_transformer = create_y_transformer(y_train, self.cfg.task)
-        
+
         dataset_train_generator = DatasetFinetuneGenerator(
             self.cfg,
-            x = x_train,
-            y = self.y_transformer.transform(y_train),
-            task = self.cfg.task,
-            max_samples_support = self.cfg.hyperparams['max_samples_support'],
-            max_samples_query = self.cfg.hyperparams['max_samples_query'],
-            split = 0.8,
+            x=x_train,
+            y=self.y_transformer.transform(y_train),
+            task=self.cfg.task,
+            max_samples_support=self.cfg.hyperparams["max_samples_support"],
+            max_samples_query=self.cfg.hyperparams["max_samples_query"],
+            split=0.8,
             random_state=rng,
         )
 
         if use_val:
             dataset_valid = DatasetFinetune(
                 self.cfg,
-                x_support = x_train,
-                y_support = self.y_transformer.transform(y_train),
-                x_query = x_val,
-                y_query = y_val,
-                max_samples_support = self.cfg.hyperparams['max_samples_support'],
-                max_samples_query = self.cfg.hyperparams['max_samples_query'],
+                x_support=x_train,
+                y_support=self.y_transformer.transform(y_train),
+                x_query=x_val,
+                y_query=y_val,
+                max_samples_support=self.cfg.hyperparams["max_samples_support"],
+                max_samples_query=self.cfg.hyperparams["max_samples_query"],
             )
             loader_valid = self.make_loader(dataset_valid, training=False)
         else:
@@ -115,7 +120,7 @@ class TrainerFinetune(BaseEstimator):
         if use_val and self.use_best_epoch:
             checkpoint.reset()
 
-        max_epochs = self.cfg.hyperparams['max_epochs']
+        max_epochs = self.cfg.hyperparams["max_epochs"]
 
         epoch = 0
         if max_epochs != 0 and use_val:
@@ -135,15 +140,15 @@ class TrainerFinetune(BaseEstimator):
                 time_cur = time.time()
                 time_elapsed = time_cur - time_start
                 time_left = time_limit - time_elapsed
-                if time_left < (time_elapsed*3+3):
+                if time_left < (time_elapsed * 3 + 3):
                     # Fine-tuning an epoch will take longer than this, so triple the time required
                     logger.log(15, "Early stopping due to running out of time...")
                     max_epochs = 0
 
-        for epoch in range(1, max_epochs+1):
-            dataset_train = next(dataset_train_generator)            
+        for epoch in range(1, max_epochs + 1):
+            dataset_train = next(dataset_train_generator)
             loader_train = self.make_loader(dataset_train, training=True)
-            
+
             metrics_train = self.train_epoch(loader_train, return_metrics=self.compute_train_metrics)
             if use_val:
                 metrics_valid = self.test_epoch(loader_valid, y_val)
@@ -165,7 +170,9 @@ class TrainerFinetune(BaseEstimator):
                 if self.early_stopping.we_should_stop():
                     logger.info("Early stopping")
                     break
-                self.scheduler.step(metrics_valid.loss)  # TODO: Make scheduler work properly during refit with no val data, to mimic scheduler in OG fit
+                self.scheduler.step(
+                    metrics_valid.loss
+                )  # TODO: Make scheduler work properly during refit with no val data, to mimic scheduler in OG fit
 
             if time_limit is not None:
                 time_cur = time.time()
@@ -173,7 +180,7 @@ class TrainerFinetune(BaseEstimator):
 
                 time_per_epoch = time_elapsed / epoch
                 time_left = time_limit - time_elapsed
-                if time_left < (time_per_epoch+3):
+                if time_left < (time_per_epoch + 3):
                     logger.log(15, "Early stopping due to running out of time...")
                     break
 
@@ -189,7 +196,9 @@ class TrainerFinetune(BaseEstimator):
         self.optimizer = None
         self.scheduler = None
 
-    def train_epoch(self, dataloader: torch.utils.data.DataLoader, return_metrics: bool = False) -> PredictionMetrics | None:
+    def train_epoch(
+        self, dataloader: torch.utils.data.DataLoader, return_metrics: bool = False
+    ) -> PredictionMetrics | None:
         """
 
         Parameters
@@ -214,20 +223,25 @@ class TrainerFinetune(BaseEstimator):
         for batch in dataloader:
             self.optimizer.zero_grad()
 
-            x_support = batch['x_support'].to(self.cfg.device)
-            y_support = batch['y_support'].to(self.cfg.device)
-            x_query = batch['x_query'].to(self.cfg.device)
-            y_query = batch['y_query'].to(self.cfg.device)
+            x_support = batch["x_support"].to(self.cfg.device)
+            y_support = batch["y_support"].to(self.cfg.device)
+            x_query = batch["x_query"].to(self.cfg.device)
+            y_query = batch["y_query"].to(self.cfg.device)
 
             if self.cfg.task == Task.REGRESSION:
-                x_support, y_support, x_query, y_query = x_support.float(), y_support.float(), x_query.float(), y_query.float()
+                x_support, y_support, x_query, y_query = (
+                    x_support.float(),
+                    y_support.float(),
+                    x_query.float(),
+                    y_query.float(),
+                )
 
             y_hat = self.model(x_support, y_support, x_query)
 
             if self.cfg.task == Task.REGRESSION:
                 y_hat = y_hat[0, :, 0]
             else:
-                y_hat = y_hat[0, :, :self.n_classes]
+                y_hat = y_hat[0, :, : self.n_classes]
 
             y_query = y_query[0, :]
 
@@ -241,7 +255,9 @@ class TrainerFinetune(BaseEstimator):
         if return_metrics:
             y_true, y_pred = output_tracker.get()
             y_pred = self.y_transformer.inverse_transform(y_pred)
-            prediction_metrics = PredictionMetrics.from_prediction(y_pred, y_true, self.cfg.task, metric=self.stopping_metric)
+            prediction_metrics = PredictionMetrics.from_prediction(
+                y_pred, y_true, self.cfg.task, metric=self.stopping_metric
+            )
             return prediction_metrics
         else:
             return None
@@ -251,13 +267,16 @@ class TrainerFinetune(BaseEstimator):
         y_hat = self.predict_epoch(dataloader)
         y_hat_finish = self.y_transformer.inverse_transform(y_hat)
 
-        prediction_metrics = PredictionMetrics.from_prediction(y_hat_finish, y_test, self.cfg.task, metric=self.stopping_metric)
+        prediction_metrics = PredictionMetrics.from_prediction(
+            y_hat_finish, y_test, self.cfg.task, metric=self.stopping_metric
+        )
         return prediction_metrics
 
     def _get_memory_size(self) -> int:
         import gc
-        import sys
         import pickle
+        import sys
+
         gc.collect()  # Try to avoid OOM error
         return sys.getsizeof(pickle.dumps(self, protocol=4))
 
@@ -270,24 +289,24 @@ class TrainerFinetune(BaseEstimator):
         x_query = self.preprocessor.transform(x_query)
 
         dataset = DatasetFinetune(
-            self.cfg, 
-            x_support = x_support, 
-            y_support = self.y_transformer.transform(y_support),
-            x_query = x_query,
-            y_query = None,
-            max_samples_support = self.cfg.hyperparams['max_samples_support'],
-            max_samples_query = self.cfg.hyperparams['max_samples_query'],
+            self.cfg,
+            x_support=x_support,
+            y_support=self.y_transformer.transform(y_support),
+            x_query=x_query,
+            y_query=None,
+            max_samples_support=self.cfg.hyperparams["max_samples_support"],
+            max_samples_query=self.cfg.hyperparams["max_samples_query"],
         )
 
         loader = self.make_loader(dataset, training=False)
 
         y_hat_ensembles = []
 
-        for _ in range(self.cfg.hyperparams['n_ensembles']):
+        for _ in range(self.cfg.hyperparams["n_ensembles"]):
             y_hat = self.predict_epoch(loader)
             y_hat_ensembles.append(y_hat)
 
-        y_hat_ensembled = sum(y_hat_ensembles) / self.cfg.hyperparams['n_ensembles']
+        y_hat_ensembled = sum(y_hat_ensembles) / self.cfg.hyperparams["n_ensembles"]
         y_hat_finish = self.y_transformer.inverse_transform(y_hat_ensembled)
 
         return y_hat_finish
@@ -304,20 +323,19 @@ class TrainerFinetune(BaseEstimator):
 
         with torch.no_grad():
             for batch in dataloader:
-                
-                x_support = batch['x_support'].to(self.cfg.device)
-                y_support = batch['y_support'].to(self.cfg.device)
-                x_query = batch['x_query'].to(self.cfg.device)
+                x_support = batch["x_support"].to(self.cfg.device)
+                y_support = batch["y_support"].to(self.cfg.device)
+                x_query = batch["x_query"].to(self.cfg.device)
 
                 if self.cfg.task == Task.REGRESSION:
                     y_support = y_support.float()
 
                 y_hat = self.model(x_support, y_support, x_query)
-               
+
                 if self.cfg.task == Task.REGRESSION:
                     y_hat = y_hat[0, :, 0]
                 else:
-                    y_hat = y_hat[0, :, :self.n_classes]
+                    y_hat = y_hat[0, :, : self.n_classes]
 
                 y_hat_list.append(einops.asnumpy(y_hat))
 
@@ -325,7 +343,6 @@ class TrainerFinetune(BaseEstimator):
         return y_hat
 
     def make_loader(self, dataset, training):
-
         return torch.utils.data.DataLoader(
             dataset,
             batch_size=1,
@@ -333,7 +350,5 @@ class TrainerFinetune(BaseEstimator):
             pin_memory=True,
             num_workers=0,
             drop_last=False,
-            collate_fn=CollatorWithPadding(
-                pad_to_n_support_samples=None
-            )
-        )
+            collate_fn=CollatorWithPadding(pad_to_n_support_samples=None),
+        )

autogluon/tabular/models/tabpfnmix/_internal/core/y_transformer.py

@@ -1,5 +1,4 @@
 import numpy as np
-
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.pipeline import FunctionTransformer
 from sklearn.preprocessing import QuantileTransformer
@@ -22,16 +21,15 @@ def create_y_transformer(y_train: np.ndarray, task: Task) -> TransformerMixin:
 
 
 class QuantileTransformer1D(BaseEstimator, TransformerMixin):
-
     def __init__(self, output_distribution="normal") -> None:
         self.quantile_transformer = QuantileTransformer(output_distribution=output_distribution)
 
     def fit(self, x: np.ndarray):
         self.quantile_transformer.fit(x[:, None])
         return self
-    
+
     def transform(self, x: np.ndarray):
         return self.quantile_transformer.transform(x[:, None])[:, 0]
-    
+
     def inverse_transform(self, x: np.ndarray):
-        return self.quantile_transformer.inverse_transform(x[:, None])[:, 0]
+        return self.quantile_transformer.inverse_transform(x[:, None])[:, 0]

autogluon/tabular/models/tabpfnmix/_internal/data/dataset_finetune.py

@@ -19,24 +19,24 @@ class DatasetFinetune(torch.utils.data.Dataset):
     """
 
     def __init__(
-        self, 
+        self,
         cfg: ConfigRun,
-        x_support: np.ndarray, 
-        y_support: np.ndarray, 
-        x_query: np.ndarray, 
+        x_support: np.ndarray,
+        y_support: np.ndarray,
+        x_query: np.ndarray,
         y_query: Optional[np.ndarray],
         max_samples_support: int,
-        max_samples_query: int
+        max_samples_query: int,
     ):
         """
         :param: max_features: number of features the tab pfn model has been trained on
         """
 
         self.cfg = cfg
-        
+
         self.x_support = x_support
         self.y_support = y_support
-        self.x_query = x_query        
+        self.x_query = x_query
         self.y_query = y_query
 
         if self.y_query is None:
@@ -53,17 +53,11 @@ class DatasetFinetune(torch.utils.data.Dataset):
         # We push the whole training data through the model, unless it is too large
         self.support_size = min(self.max_samples_support, self.n_samples_support)
 
-
     def __len__(self):
         return len(self.x_queries)
 
     def __getitem__(self, idx):
-
-        support_indices = np.random.choice(
-            self.n_samples_support, 
-            size=self.support_size, 
-            replace=False
-        )
+        support_indices = np.random.choice(self.n_samples_support, size=self.support_size, replace=False)
 
         x_support = self.x_support[support_indices]
         y_support = self.y_support[support_indices]
@@ -74,13 +68,11 @@ class DatasetFinetune(torch.utils.data.Dataset):
         y_query_tensor = torch.as_tensor(self.y_queries[idx])
 
         return {
-            'x_support': x_support_tensor,
-            'y_support': y_support_tensor,
-            'x_query': x_query_tensor,
-            'y_query': y_query_tensor,
+            "x_support": x_support_tensor,
+            "y_support": y_support_tensor,
+            "x_query": x_query_tensor,
+            "y_query": y_query_tensor,
         }
-    
-
 
     def split_in_chunks(self, x: np.ndarray, batch_size: int) -> list[np.ndarray]:
         """
@@ -91,19 +83,15 @@ class DatasetFinetune(torch.utils.data.Dataset):
         x_chunks = []
 
         for i in range(n_chunks):
-            x_chunks.append(x[i * batch_size: (i + 1) * batch_size])
+            x_chunks.append(x[i * batch_size : (i + 1) * batch_size])
 
         return x_chunks
-    
-
-
-
 
 
 def DatasetFinetuneGenerator(
     cfg: ConfigRun,
-    x: np.ndarray, 
-    y: np.ndarray, 
+    x: np.ndarray,
+    y: np.ndarray,
     task: Task,
     max_samples_support: int,
     max_samples_query: int,
@@ -116,9 +104,8 @@ def DatasetFinetuneGenerator(
     Every single iteration, the generator yields a different support and query set split.
     The dataset made always has exactly one batch.
     """
-        
-    while True:
 
+    while True:
         x_support, x_query, y_support, y_query = make_dataset_split(x=x, y=y, task=task, random_state=random_state)
         n_samples_support = x_support.shape[0]
         n_samples_query = x_query.shape[0]
@@ -136,4 +123,4 @@ def DatasetFinetuneGenerator(
             max_samples_query=max_samples_query,
         )
 
-        yield dataset_finetune
+        yield dataset_finetune

autogluon/tabular/models/tabpfnmix/_internal/data/preprocessor.py

@@ -14,26 +14,24 @@ class Preprocessor(TransformerMixin, BaseEstimator):
     """
     This class is used to preprocess the data before it is pushed through the model.
     The preprocessor assures that the data has the right shape and is normalized,
-    This way the model always gets the same input distribution, 
+    This way the model always gets the same input distribution,
     no matter whether the input data is synthetic or real.
 
     """
 
     def __init__(
-            self, 
-            max_features: int,
-            use_quantile_transformer: bool,
-            use_feature_count_scaling: bool,
-            task: Task,
-        ):
-
+        self,
+        max_features: int,
+        use_quantile_transformer: bool,
+        use_feature_count_scaling: bool,
+        task: Task,
+    ):
         self.max_features = max_features
         self.use_quantile_transformer = use_quantile_transformer
         self.use_feature_count_scaling = use_feature_count_scaling
         self.task = task
-    
-    def fit(self, X: np.ndarray, y: np.ndarray):
 
+    def fit(self, X: np.ndarray, y: np.ndarray):
         self.compute_pre_nan_mean(X)
         X = self.impute_nan_features_with_mean(X)
 
@@ -46,26 +44,24 @@ class Preprocessor(TransformerMixin, BaseEstimator):
         if self.use_quantile_transformer:
             n_obs, n_features = X.shape
             n_quantiles = min(n_obs, 1000)
-            self.quantile_transformer = QuantileTransformer(n_quantiles=n_quantiles, output_distribution='normal')
+            self.quantile_transformer = QuantileTransformer(n_quantiles=n_quantiles, output_distribution="normal")
             X = self.quantile_transformer.fit_transform(X)
-        
+
         self.mean, self.std = self.calc_mean_std(X)
         X = self.normalize_by_mean_std(X, self.mean, self.std)
 
         assert np.isnan(X).sum() == 0, "There are NaNs in the data after preprocessing"
 
         return self
-    
 
     def transform(self, X: np.ndarray):
-
         X = self.cutoff_singular_features(X, self.singular_features)
         X = self.impute_nan_features_with_mean(X)
         X = self.select_features(X)
 
         if self.use_quantile_transformer:
             X = self.quantile_transformer.transform(X)
-        
+
         X = self.normalize_by_mean_std(X, self.mean, self.std)
 
         if self.use_feature_count_scaling:
@@ -76,16 +72,15 @@ class Preprocessor(TransformerMixin, BaseEstimator):
         assert np.isnan(X).sum() == 0, "There are NaNs in the data after preprocessing"
 
         return X
-        
 
     def determine_which_features_are_singular(self, x: np.ndarray) -> None:
-
-        self.singular_features = np.array([ len(np.unique(x_col)) for x_col in x.T ]) == 1
+        self.singular_features = np.array([len(np.unique(x_col)) for x_col in x.T]) == 1
 
     def determine_which_features_to_select(self, x: np.ndarray, y: np.ndarray) -> None:
-
         if x.shape[1] > self.max_features:
-            logger.info(f"A maximum of {self.max_features} features are allowed, but the dataset has {x.shape[1]} features. A subset of {self.max_features} are selected using SelectKBest")
+            logger.info(
+                f"A maximum of {self.max_features} features are allowed, but the dataset has {x.shape[1]} features. A subset of {self.max_features} are selected using SelectKBest"
+            )
 
             if self.task == Task.CLASSIFICATION:
                 self.select_k_best = SelectKBest(k=self.max_features, score_func=f_classif)
@@ -99,30 +94,23 @@ class Preprocessor(TransformerMixin, BaseEstimator):
         """
         self.pre_nan_mean = np.nanmean(x, axis=0)
 
-
     def impute_nan_features_with_mean(self, x: np.ndarray) -> np.ndarray:
-
         inds = np.where(np.isnan(x))
         x[inds] = np.take(self.pre_nan_mean, inds[1])
         return x
 
-    
     def select_features(self, x: np.ndarray) -> np.ndarray:
-
         if x.shape[1] > self.max_features:
             x = self.select_k_best.transform(x)
 
         return x
-    
 
     def cutoff_singular_features(self, x: np.ndarray, singular_features: np.ndarray) -> np.ndarray:
-
         if singular_features.any():
             x = x[:, ~singular_features]
 
         return x
 
-
     def calc_mean_std(self, x: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
         """
         Calculates the mean and std of the training data
@@ -130,7 +118,6 @@ class Preprocessor(TransformerMixin, BaseEstimator):
         mean = x.mean(axis=0)
         std = x.std(axis=0)
         return mean, std
-    
 
     def normalize_by_mean_std(self, x: np.ndarray, mean: np.ndarray, std: np.ndarray) -> np.ndarray:
         """
@@ -140,7 +127,6 @@ class Preprocessor(TransformerMixin, BaseEstimator):
         x = (x - mean) / std
         return x
 
-
     def normalize_by_feature_count(self, x: np.ndarray, max_features) -> np.ndarray:
         """
         An interesting way of normalization by the tabPFN paper
@@ -149,8 +135,6 @@ class Preprocessor(TransformerMixin, BaseEstimator):
         x = x * max_features / x.shape[1]
         return x
 
-
-
     def extend_feature_dim_to_max_features(self, x: np.ndarray, max_features) -> np.ndarray:
         """
         Increases the number of features to the number of features the model has been trained on
@@ -158,7 +142,3 @@ class Preprocessor(TransformerMixin, BaseEstimator):
         added_zeros = np.zeros((x.shape[0], max_features - x.shape[1]), dtype=np.float32)
         x = np.concatenate([x, added_zeros], axis=1)
         return x
-    
-
-
-