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

autogluon/tabular/models/mitra/_internal/core/prediction_metrics.py

@@ -11,23 +11,20 @@ from ..._internal.data.preprocessor import Preprocessor
 
 
 @dataclass
-class PredictionMetrics():
+class PredictionMetrics:
     task: Task
     loss: float
     score: float
     metrics: dict[MetricName, float]
 
-
     @classmethod
     def from_prediction(cls, y_pred: np.ndarray, y_true: np.ndarray, task: Task) -> "PredictionMetrics":
-
         loss, score, metrics = compute_metrics(y_pred, y_true, task)
 
         return cls(task=task, loss=loss, score=score, metrics=metrics)
 
 
 def compute_metrics(y_pred: np.ndarray, y_true: np.ndarray, task: Task) -> tuple[float, float, dict]:
-
     if task == Task.CLASSIFICATION:
         return compute_classification_metrics(y_pred, y_true)
     elif task == Task.REGRESSION:
@@ -39,14 +36,20 @@ def compute_classification_metrics(y_pred: np.ndarray, y_true: np.ndarray) -> tu
 
     y_pred_class = np.argmax(y_pred, axis=1)
     y_pred_proba = scipy.special.softmax(y_pred, axis=1)
-    y_pred_proba = y_pred_proba / y_pred_proba.sum(axis=1, keepdims=True)   # softmax not completely numerically stable, so a small correction is needed
+    y_pred_proba = y_pred_proba / y_pred_proba.sum(
+        axis=1, keepdims=True
+    )  # softmax not completely numerically stable, so a small correction is needed
     labels = np.arange(y_pred_proba.shape[1])
 
     metrics = {
         MetricName.ACCURACY: (y_true == y_pred_class).mean(),
         MetricName.F1: f1_score(y_true, y_pred_class, average="weighted"),
-        MetricName.AUC: roc_auc_score_multiclass(y_true, y_pred_proba, multi_class='ovo', average='macro', labels=labels),
-        MetricName.LOG_LOSS: torch.nn.functional.cross_entropy(torch.from_numpy(y_pred), torch.from_numpy(y_true)).item()
+        MetricName.AUC: roc_auc_score_multiclass(
+            y_true, y_pred_proba, multi_class="ovo", average="macro", labels=labels
+        ),
+        MetricName.LOG_LOSS: torch.nn.functional.cross_entropy(
+            torch.from_numpy(y_pred), torch.from_numpy(y_true)
+        ).item(),
     }
 
     loss = metrics[MetricName.LOG_LOSS]
@@ -55,14 +58,14 @@ def compute_classification_metrics(y_pred: np.ndarray, y_true: np.ndarray) -> tu
     return loss, score, metrics
 
 
-def roc_auc_score_multiclass(y_true, y_pred_proba, multi_class='ovo', average='macro', labels=None) -> float:
-    """ 
+def roc_auc_score_multiclass(y_true, y_pred_proba, multi_class="ovo", average="macro", labels=None) -> float:
+    """
     The roc_auc_score multi_class is not supported for binary classification
     """
 
     if np.unique(y_true).shape[0] == 1:
         # AUC is not defined if there is only one class
-        return float('nan')
+        return float("nan")
 
     try:
         if y_pred_proba.shape[1] == 2:
@@ -76,12 +79,11 @@ def roc_auc_score_multiclass(y_true, y_pred_proba, multi_class='ovo', average='m
 
 
 def compute_regression_metrics(y_pred: np.ndarray, y_true: np.ndarray) -> tuple[float, float, dict]:
-
     metrics = {
         MetricName.RMSE: root_mean_squared_error(y_true, y_pred),
         MetricName.MSE: mean_squared_error(y_true, y_pred),
         MetricName.MAE: np.abs(y_true - y_pred).mean(),
-        MetricName.R2: r2_score(y_true, y_pred)
+        MetricName.R2: r2_score(y_true, y_pred),
     }
 
     loss = metrics[MetricName.MSE]
@@ -90,27 +92,22 @@ def compute_regression_metrics(y_pred: np.ndarray, y_true: np.ndarray) -> tuple[
     return loss, score, metrics
 
 
-class PredictionMetricsTracker():
+class PredictionMetricsTracker:
     """
     Prediction metrics tracker that accumulates predictions and true values to compute metrics at the end.
     Uses torch.Tensor for predictions and true values.
     """
 
     def __init__(self, task: Task, preprocessor: Preprocessor) -> None:
-
         self.task = task
         self.preprocessor = preprocessor
         self.reset()
 
-
     def reset(self) -> None:
-
         self.ys_pred: list[np.ndarray] = []
         self.ys_true: list[np.ndarray] = []
 
-
     def update(self, y_pred: torch.Tensor, y_true: torch.Tensor, train: bool) -> None:
-
         y_pred_np = y_pred.detach().cpu().numpy()[0]
         y_pred_ori = self.preprocessor.inverse_transform_y(y_pred_np)
 
@@ -123,10 +120,8 @@ class PredictionMetricsTracker():
         self.ys_pred.append(y_pred_ori)
         self.ys_true.append(y_true_ori)
 
-
     def get_metrics(self) -> PredictionMetrics:
-
         y_pred = np.concatenate(self.ys_pred, axis=0)
         y_true = np.concatenate(self.ys_true, axis=0)
 
-        return PredictionMetrics.from_prediction(y_pred, y_true, self.task)
+        return PredictionMetrics.from_prediction(y_pred, y_true, self.task)

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

@@ -18,17 +18,15 @@ from ..._internal.data.preprocessor import Preprocessor
 
 
 class TrainerFinetune(BaseEstimator):
-
     def __init__(
-            self,
-            cfg: ConfigRun,
-            model: torch.nn.Module,
-            n_classes: int,
-            device: str,
-            rng: np.random.RandomState = None,
-            verbose: bool = True,
+        self,
+        cfg: ConfigRun,
+        model: torch.nn.Module,
+        n_classes: int,
+        device: str,
+        rng: np.random.RandomState = None,
+        verbose: bool = True,
     ):
-
         self.cfg = cfg
         if rng is None:
             rng = np.random.RandomState(self.cfg.seed)
@@ -42,36 +40,36 @@ class TrainerFinetune(BaseEstimator):
         self.optimizer = get_optimizer(self.cfg.hyperparams, self.model)
         self.scheduler_warmup, self.scheduler_reduce_on_plateau = get_scheduler(self.cfg.hyperparams, self.optimizer)
         self.scaler = GradScaler(
-            enabled=self.cfg.hyperparams['grad_scaler_enabled'],
-            scale_init=self.cfg.hyperparams['grad_scaler_scale_init'],
-            scale_min=self.cfg.hyperparams['grad_scaler_scale_min'],
-            growth_interval=self.cfg.hyperparams['grad_scaler_growth_interval'],
-            device=self.device
+            enabled=self.cfg.hyperparams["grad_scaler_enabled"],
+            scale_init=self.cfg.hyperparams["grad_scaler_scale_init"],
+            scale_min=self.cfg.hyperparams["grad_scaler_scale_min"],
+            growth_interval=self.cfg.hyperparams["grad_scaler_growth_interval"],
+            device=self.device,
         )
 
-        self.early_stopping = EarlyStopping(patience=self.cfg.hyperparams['early_stopping_patience'])
+        self.early_stopping = EarlyStopping(patience=self.cfg.hyperparams["early_stopping_patience"])
         self.checkpoint = Checkpoint()
         self.preprocessor = Preprocessor(
-            dim_embedding=self.cfg.hyperparams['dim_embedding'],
+            dim_embedding=self.cfg.hyperparams["dim_embedding"],
             n_classes=self.n_classes,
-            dim_output=self.cfg.hyperparams['dim_output'],
-            use_quantile_transformer=self.cfg.hyperparams['use_quantile_transformer'],
-            use_feature_count_scaling=self.cfg.hyperparams['use_feature_count_scaling'],
-            use_random_transforms=self.cfg.hyperparams['use_random_transforms'],
-            shuffle_classes=self.cfg.hyperparams['shuffle_classes'],
-            shuffle_features=self.cfg.hyperparams['shuffle_features'],
-            random_mirror_x=self.cfg.hyperparams['random_mirror_x'],
-            random_mirror_regression=self.cfg.hyperparams['random_mirror_regression'],
-            task=self.cfg.task
+            dim_output=self.cfg.hyperparams["dim_output"],
+            use_quantile_transformer=self.cfg.hyperparams["use_quantile_transformer"],
+            use_feature_count_scaling=self.cfg.hyperparams["use_feature_count_scaling"],
+            use_random_transforms=self.cfg.hyperparams["use_random_transforms"],
+            shuffle_classes=self.cfg.hyperparams["shuffle_classes"],
+            shuffle_features=self.cfg.hyperparams["shuffle_features"],
+            random_mirror_x=self.cfg.hyperparams["random_mirror_x"],
+            random_mirror_regression=self.cfg.hyperparams["random_mirror_regression"],
+            task=self.cfg.task,
         )
 
         self.checkpoint.reset(self.model)
 
-        if self.cfg.task == Task.REGRESSION and self.cfg.hyperparams['regression_loss'] == LossName.CROSS_ENTROPY:
-            self.bins = torch.linspace(-0.5, 1.5, self.cfg.hyperparams['dim_output']+1, device=cfg.device)
+        if self.cfg.task == Task.REGRESSION and self.cfg.hyperparams["regression_loss"] == LossName.CROSS_ENTROPY:
+            self.bins = torch.linspace(-0.5, 1.5, self.cfg.hyperparams["dim_output"] + 1, device=cfg.device)
             self.bin_width = self.bins[1] - self.bins[0]
 
-        self.metric = self.cfg.hyperparams['metric']
+        self.metric = self.cfg.hyperparams["metric"]
 
     def set_device(self, device: str):
         self.device = device
@@ -89,7 +87,6 @@ class TrainerFinetune(BaseEstimator):
         self.metric = None
 
     def train(self, x_train: np.ndarray, y_train: np.ndarray, x_val: np.ndarray, y_val: np.ndarray):
-
         self.preprocessor.fit(x_train, y_train)
 
         x_train_transformed = self.preprocessor.transform_X(x_train)
@@ -97,11 +94,11 @@ class TrainerFinetune(BaseEstimator):
 
         dataset_train_generator = DatasetFinetuneGenerator(
             self.cfg,
-            x = x_train_transformed,
-            y = y_train_transformed,
-            task = self.cfg.task,
-            max_samples_support = self.cfg.hyperparams['max_samples_support'],
-            max_samples_query = self.cfg.hyperparams['max_samples_query'],
+            x=x_train_transformed,
+            y=y_train_transformed,
+            task=self.cfg.task,
+            max_samples_support=self.cfg.hyperparams["max_samples_support"],
+            max_samples_query=self.cfg.hyperparams["max_samples_query"],
             rng=self.rng,
         )
 
@@ -114,8 +111,7 @@ class TrainerFinetune(BaseEstimator):
 
         start_time = time.time()
 
-        for epoch in range(1, self.cfg.hyperparams['max_epochs']+1):
-
+        for epoch in range(1, self.cfg.hyperparams["max_epochs"] + 1):
             dataset_train = next(dataset_train_generator)
             loader_train = self.make_loader(dataset_train, training=True)
             self.model.train()
@@ -123,31 +119,39 @@ class TrainerFinetune(BaseEstimator):
             prediction_metrics_tracker = PredictionMetricsTracker(task=self.cfg.task, preprocessor=self.preprocessor)
 
             for batch in loader_train:
-
-                with torch.autocast(device_type=self.device, dtype=getattr(torch, self.cfg.hyperparams['precision'])):
-
-                    x_support = batch['x_support'].to(self.device, non_blocking=True)
-                    y_support = batch['y_support'].to(self.device, non_blocking=True)
-                    x_query = batch['x_query'].to(self.device, non_blocking=True)
-                    y_query = batch['y_query'].to(self.device, non_blocking=True)
-                    padding_features = batch['padding_features'].to(self.device, non_blocking=True)
-                    padding_obs_support = batch['padding_obs_support'].to(self.device, non_blocking=True)
-                    padding_obs_query = batch['padding_obs_query'].to(self.device, non_blocking=True)
+                with torch.autocast(device_type=self.device, dtype=getattr(torch, self.cfg.hyperparams["precision"])):
+                    x_support = batch["x_support"].to(self.device, non_blocking=True)
+                    y_support = batch["y_support"].to(self.device, non_blocking=True)
+                    x_query = batch["x_query"].to(self.device, non_blocking=True)
+                    y_query = batch["y_query"].to(self.device, non_blocking=True)
+                    padding_features = batch["padding_features"].to(self.device, non_blocking=True)
+                    padding_obs_support = batch["padding_obs_support"].to(self.device, non_blocking=True)
+                    padding_obs_query = batch["padding_obs_query"].to(self.device, non_blocking=True)
 
                     # Convert numerical y_support to bin ids
-                    if self.cfg.task == Task.REGRESSION and self.cfg.hyperparams['regression_loss'] == LossName.CROSS_ENTROPY:
+                    if (
+                        self.cfg.task == Task.REGRESSION
+                        and self.cfg.hyperparams["regression_loss"] == LossName.CROSS_ENTROPY
+                    ):
                         y_support = torch.bucketize(y_support, self.bins) - 1
-                        y_support = torch.clamp(y_support, 0, self.cfg.hyperparams['dim_output']-1).to(torch.int64)
+                        y_support = torch.clamp(y_support, 0, self.cfg.hyperparams["dim_output"] - 1).to(torch.int64)
                         y_query_bin_ids = torch.bucketize(y_query, self.bins) - 1
-                        y_query_bin_ids = torch.clamp(y_query_bin_ids, 0, self.cfg.hyperparams['dim_output']-1).to(torch.int64)
+                        y_query_bin_ids = torch.clamp(y_query_bin_ids, 0, self.cfg.hyperparams["dim_output"] - 1).to(
+                            torch.int64
+                        )
 
                     if self.cfg.model_name == ModelName.TABPFN:
                         y_hat = self.model(x_support, y_support, x_query, task=self.cfg.task).squeeze(-1)
                     elif self.cfg.model_name in [ModelName.TAB2D, ModelName.TAB2D_COL_ROW, ModelName.TAB2D_SDPA]:
-                        y_hat = self.model(x_support, y_support, x_query, padding_features, padding_obs_support, padding_obs_query)
+                        y_hat = self.model(
+                            x_support, y_support, x_query, padding_features, padding_obs_support, padding_obs_query
+                        )
 
                     # Convert numerical y_query to bin ids
-                    if self.cfg.task == Task.REGRESSION and self.cfg.hyperparams['regression_loss'] == LossName.CROSS_ENTROPY:
+                    if (
+                        self.cfg.task == Task.REGRESSION
+                        and self.cfg.hyperparams["regression_loss"] == LossName.CROSS_ENTROPY
+                    ):
                         loss = self.loss(y_hat, y_query_bin_ids)
                     elif self.cfg.task == Task.CLASSIFICATION:
                         # for b in range(y_support.shape[0]):
@@ -163,7 +167,10 @@ class TrainerFinetune(BaseEstimator):
                 self.scaler.update()
 
                 # Convert bin id predictions to numerical values
-                if self.cfg.task == Task.REGRESSION and self.cfg.hyperparams['regression_loss'] == LossName.CROSS_ENTROPY:
+                if (
+                    self.cfg.task == Task.REGRESSION
+                    and self.cfg.hyperparams["regression_loss"] == LossName.CROSS_ENTROPY
+                ):
                     y_hat = torch.argmax(y_hat, dim=-1)
                     y_hat = self.bins[y_hat] + self.bin_width / 2
 
@@ -187,19 +194,24 @@ class TrainerFinetune(BaseEstimator):
                     logger.info("Early stopping")
                 break
 
-            if self.cfg.hyperparams["budget"] is not None and self.cfg.hyperparams["budget"] > 0 and time.time() - start_time > self.cfg.hyperparams["budget"]:
+            if (
+                self.cfg.hyperparams["budget"] is not None
+                and self.cfg.hyperparams["budget"] > 0
+                and time.time() - start_time > self.cfg.hyperparams["budget"]
+            ):
                 logger.info("Time limit reached")
                 break
 
-            if epoch < self.cfg.hyperparams['warmup_steps']:
+            if epoch < self.cfg.hyperparams["warmup_steps"]:
                 self.scheduler_warmup.step()
             else:
                 self.scheduler_reduce_on_plateau.step(metrics_valid.loss)
 
         self.checkpoint.set_to_best(self.model)
 
-    def evaluate(self, x_support: np.ndarray, y_support: np.ndarray, x_query: np.ndarray, y_query: np.ndarray) -> PredictionMetrics:
-
+    def evaluate(
+        self, x_support: np.ndarray, y_support: np.ndarray, x_query: np.ndarray, y_query: np.ndarray
+    ) -> PredictionMetrics:
         self.model.eval()
 
         x_support_transformed = self.preprocessor.transform_X(x_support)
@@ -209,12 +221,12 @@ class TrainerFinetune(BaseEstimator):
 
         dataset = DatasetFinetune(
             self.cfg,
-            x_support = x_support_transformed,
-            y_support = y_support_transformed,
-            x_query = x_query_transformed,
-            y_query = y_query,
-            max_samples_support = self.cfg.hyperparams['max_samples_support'],
-            max_samples_query = self.cfg.hyperparams['max_samples_query'],
+            x_support=x_support_transformed,
+            y_support=y_support_transformed,
+            x_query=x_query_transformed,
+            y_query=y_query,
+            max_samples_support=self.cfg.hyperparams["max_samples_support"],
+            max_samples_query=self.cfg.hyperparams["max_samples_query"],
             rng=self.rng,
         )
 
@@ -223,21 +235,22 @@ class TrainerFinetune(BaseEstimator):
 
         with torch.no_grad():
             for batch in loader:
-
-                with torch.autocast(device_type=self.device, dtype=getattr(torch, self.cfg.hyperparams['precision'])):
-
-                    x_s = batch['x_support'].to(self.device, non_blocking=True)
-                    y_s = batch['y_support'].to(self.device, non_blocking=True)
-                    x_q = batch['x_query'].to(self.device, non_blocking=True)
-                    y_q = batch['y_query'].to(self.device, non_blocking=True)
-                    padding_features = batch['padding_features'].to(self.device, non_blocking=True)
-                    padding_obs_support = batch['padding_obs_support'].to(self.device, non_blocking=True)
-                    padding_obs_query = batch['padding_obs_query'].to(self.device, non_blocking=True)
+                with torch.autocast(device_type=self.device, dtype=getattr(torch, self.cfg.hyperparams["precision"])):
+                    x_s = batch["x_support"].to(self.device, non_blocking=True)
+                    y_s = batch["y_support"].to(self.device, non_blocking=True)
+                    x_q = batch["x_query"].to(self.device, non_blocking=True)
+                    y_q = batch["y_query"].to(self.device, non_blocking=True)
+                    padding_features = batch["padding_features"].to(self.device, non_blocking=True)
+                    padding_obs_support = batch["padding_obs_support"].to(self.device, non_blocking=True)
+                    padding_obs_query = batch["padding_obs_query"].to(self.device, non_blocking=True)
 
                     # Convert numerical y_support to bin ids
-                    if self.cfg.task == Task.REGRESSION and self.cfg.hyperparams['regression_loss'] == LossName.CROSS_ENTROPY:
+                    if (
+                        self.cfg.task == Task.REGRESSION
+                        and self.cfg.hyperparams["regression_loss"] == LossName.CROSS_ENTROPY
+                    ):
                         y_s = torch.bucketize(y_s, self.bins) - 1
-                        y_s = torch.clamp(y_s, 0, self.cfg.hyperparams['dim_output']-1).to(torch.int64)
+                        y_s = torch.clamp(y_s, 0, self.cfg.hyperparams["dim_output"] - 1).to(torch.int64)
 
                     if self.cfg.model_name == ModelName.TABPFN:
                         y_hat = self.model(x_s, y_s, x_q, task=self.cfg.task).squeeze(-1)
@@ -245,7 +258,10 @@ class TrainerFinetune(BaseEstimator):
                         y_hat = self.model(x_s, y_s, x_q, padding_features, padding_obs_support, padding_obs_query)
 
                 # Convert bin id predictions to numerical values
-                if self.cfg.task == Task.REGRESSION and self.cfg.hyperparams['regression_loss'] == LossName.CROSS_ENTROPY:
+                if (
+                    self.cfg.task == Task.REGRESSION
+                    and self.cfg.hyperparams["regression_loss"] == LossName.CROSS_ENTROPY
+                ):
                     y_hat = torch.argmax(y_hat, dim=-1)
                     y_hat = self.bins[y_hat] + self.bin_width / 2
 
@@ -255,21 +271,19 @@ class TrainerFinetune(BaseEstimator):
         metrics_eval = prediction_metrics_tracker.get_metrics()
         return metrics_eval
 
-
     def predict(self, x_support: np.ndarray, y_support: np.ndarray, x_query: np.ndarray) -> np.ndarray:
-
         x_support_transformed = self.preprocessor.transform_X(x_support)
         x_query_transformed = self.preprocessor.transform_X(x_query)
         y_support_transformed = self.preprocessor.transform_y(y_support)
 
         dataset = DatasetFinetune(
             self.cfg,
-            x_support = x_support_transformed,
-            y_support = y_support_transformed,
-            x_query = x_query_transformed,
-            y_query = None,
-            max_samples_support = self.cfg.hyperparams['max_samples_support'],
-            max_samples_query = self.cfg.hyperparams['max_samples_query'],
+            x_support=x_support_transformed,
+            y_support=y_support_transformed,
+            x_query=x_query_transformed,
+            y_query=None,
+            max_samples_support=self.cfg.hyperparams["max_samples_support"],
+            max_samples_query=self.cfg.hyperparams["max_samples_query"],
             rng=self.rng,
         )
 
@@ -280,20 +294,21 @@ class TrainerFinetune(BaseEstimator):
 
         with torch.no_grad():
             for batch in loader:
-
-                with torch.autocast(device_type=self.device, dtype=getattr(torch, self.cfg.hyperparams['precision'])):
-
-                    x_s = batch['x_support'].to(self.device, non_blocking=True)
-                    y_s = batch['y_support'].to(self.device, non_blocking=True)
-                    x_q = batch['x_query'].to(self.device, non_blocking=True)
-                    padding_features = batch['padding_features'].to(self.device, non_blocking=True)
-                    padding_obs_support = batch['padding_obs_support'].to(self.device, non_blocking=True)
-                    padding_obs_query = batch['padding_obs_query'].to(self.device, non_blocking=True)
+                with torch.autocast(device_type=self.device, dtype=getattr(torch, self.cfg.hyperparams["precision"])):
+                    x_s = batch["x_support"].to(self.device, non_blocking=True)
+                    y_s = batch["y_support"].to(self.device, non_blocking=True)
+                    x_q = batch["x_query"].to(self.device, non_blocking=True)
+                    padding_features = batch["padding_features"].to(self.device, non_blocking=True)
+                    padding_obs_support = batch["padding_obs_support"].to(self.device, non_blocking=True)
+                    padding_obs_query = batch["padding_obs_query"].to(self.device, non_blocking=True)
 
                     # Convert numerical y_support to bin ids
-                    if self.cfg.task == Task.REGRESSION and self.cfg.hyperparams['regression_loss'] == LossName.CROSS_ENTROPY:
+                    if (
+                        self.cfg.task == Task.REGRESSION
+                        and self.cfg.hyperparams["regression_loss"] == LossName.CROSS_ENTROPY
+                    ):
                         y_s = torch.bucketize(y_s, self.bins) - 1
-                        y_s = torch.clamp(y_s, 0, self.cfg.hyperparams['dim_output']-1).to(torch.int64)
+                        y_s = torch.clamp(y_s, 0, self.cfg.hyperparams["dim_output"] - 1).to(torch.int64)
 
                     if self.cfg.model_name == ModelName.TABPFN:
                         y_hat = self.model(x_s, y_s, x_q, task=self.cfg.task).squeeze(-1)
@@ -303,7 +318,10 @@ class TrainerFinetune(BaseEstimator):
                 y_hat = y_hat[0].float().cpu().numpy()
 
                 # Convert bin id predictions to numerical values
-                if self.cfg.task == Task.REGRESSION and self.cfg.hyperparams['regression_loss'] == LossName.CROSS_ENTROPY:
+                if (
+                    self.cfg.task == Task.REGRESSION
+                    and self.cfg.hyperparams["regression_loss"] == LossName.CROSS_ENTROPY
+                ):
                     y_hat = np.argmax(y_hat, axis=-1)
                     y_hat = (self.bins[y_hat] + self.bin_width / 2).cpu().numpy()
 
@@ -314,13 +332,10 @@ class TrainerFinetune(BaseEstimator):
 
         return y_pred
 
-
     def load_params(self, path):
         self.model.load_state_dict(torch.load(path))
 
-
     def make_loader(self, dataset: torch.utils.data.Dataset, training: bool) -> torch.utils.data.DataLoader:
-
         if self.cfg.model_name == ModelName.TABPFN:
             pad_to_max_features = True
         elif self.cfg.model_name in [ModelName.TAB2D, ModelName.TAB2D_COL_ROW, ModelName.TAB2D_SDPA]:
@@ -336,16 +351,14 @@ class TrainerFinetune(BaseEstimator):
             num_workers=0,
             drop_last=False,
             collate_fn=CollatorWithPadding(
-                max_features=self.cfg.hyperparams['dim_embedding'],
-                pad_to_max_features=pad_to_max_features
+                max_features=self.cfg.hyperparams["dim_embedding"], pad_to_max_features=pad_to_max_features
             ),
         )
 
-
     def log_start_metrics(self, metrics_valid: PredictionMetrics):
-
         if self.cfg.task == Task.REGRESSION:
-            logger.info((
+            logger.info(
+                (
                     f"Epoch 000 "
                     f"| Train MSE: -.---- "
                     f"| Train MAE: -.---- "
@@ -353,21 +366,24 @@ class TrainerFinetune(BaseEstimator):
                     f"| Val MSE: {metrics_valid.metrics[MetricName.MSE]:.4f} "
                     f"| Val MAE: {metrics_valid.metrics[MetricName.MAE]:.4f} "
                     f"| Val r2: {metrics_valid.metrics[MetricName.R2]:.4f}"
-                ))
+                )
+            )
 
         elif self.cfg.task == Task.CLASSIFICATION:
-            logger.info((
+            logger.info(
+                (
                     f"Epoch 000 "
                     f"| Train CE: -.---- "
                     f"| Train acc: -.---- "
                     f"| Val CE: {metrics_valid.metrics[MetricName.LOG_LOSS]:.4f} "
                     f"| Val acc: {metrics_valid.metrics[MetricName.ACCURACY]:.4f}"
-                ))
+                )
+            )
 
     def log_metrics(self, epoch: int, metrics_train: PredictionMetrics, metrics_valid: PredictionMetrics):
-
         if self.cfg.task == Task.REGRESSION:
-            logger.info((
+            logger.info(
+                (
                     f"Epoch {epoch:03d} "
                     f"| Train MSE: {metrics_train.metrics[MetricName.MSE]:.4f} "
                     f"| Train MAE: {metrics_train.metrics[MetricName.MAE]:.4f} "
@@ -375,12 +391,15 @@ class TrainerFinetune(BaseEstimator):
                     f"| Val MSE: {metrics_valid.metrics[MetricName.MSE]:.4f} "
                     f"| Val MAE: {metrics_valid.metrics[MetricName.MAE]:.4f} "
                     f"| Val r2: {metrics_valid.metrics[MetricName.R2]:.4f}"
-                ))
+                )
+            )
         elif self.cfg.task == Task.CLASSIFICATION:
-            logger.info((
+            logger.info(
+                (
                     f"Epoch {epoch:03d} "
                     f"| Train CE: {metrics_train.metrics[MetricName.LOG_LOSS]:.4f} "
                     f"| Train acc: {metrics_train.metrics[MetricName.ACCURACY]:.4f} "
                     f"| Val CE: {metrics_valid.metrics[MetricName.LOG_LOSS]:.4f} "
                     f"| Val acc: {metrics_valid.metrics[MetricName.ACCURACY]:.4f}"
-                ))
+                )
+            )

autogluon/tabular/models/mitra/_internal/data/__init__.py

@@ -1 +1 @@
-# Data processing modules for MitraModel 
+# Data processing modules for MitraModel

autogluon/tabular/models/mitra/_internal/data/collator.py

@@ -1,23 +1,19 @@
 import torch
 
 
-class CollatorWithPadding():
-
+class CollatorWithPadding:
     def __init__(
-            self, 
-            max_features: int,
-            pad_to_max_features: bool,
-        ) -> None:
-        
+        self,
+        max_features: int,
+        pad_to_max_features: bool,
+    ) -> None:
         self.max_features = max_features
         self.pad_to_max_features = pad_to_max_features
 
-
     def __call__(self, batch: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
-
-        max_support_samples = max(dataset['x_support'].shape[0] for dataset in batch)
-        max_query_samples = max(dataset['x_query'].shape[0] for dataset in batch)
-        max_features = max(dataset['x_support'].shape[1] for dataset in batch)
+        max_support_samples = max(dataset["x_support"].shape[0] for dataset in batch)
+        max_query_samples = max(dataset["x_query"].shape[0] for dataset in batch)
+        max_features = max(dataset["x_support"].shape[1] for dataset in batch)
 
         if self.pad_to_max_features:
             max_features = self.max_features
@@ -25,22 +21,30 @@ class CollatorWithPadding():
         batch_size = len(batch)
 
         tensor_dict = {
-            'x_support': torch.zeros((batch_size, max_support_samples, max_features), dtype=batch[0]['x_support'].dtype),
-            'y_support': torch.full((batch_size, max_support_samples), fill_value=-100, dtype=batch[0]['y_support'].dtype),
-            'x_query': torch.zeros((batch_size, max_query_samples, max_features), dtype=batch[0]['x_query'].dtype),
-            'y_query': torch.full((batch_size, max_query_samples), fill_value=-100, dtype=batch[0]['y_query'].dtype),
-            'padding_features': torch.ones((batch_size, max_features), dtype=torch.bool),
-            'padding_obs_support': torch.ones((batch_size, max_support_samples), dtype=torch.bool),
-            'padding_obs_query': torch.ones((batch_size, max_query_samples), dtype=torch.bool),
+            "x_support": torch.zeros(
+                (batch_size, max_support_samples, max_features), dtype=batch[0]["x_support"].dtype
+            ),
+            "y_support": torch.full(
+                (batch_size, max_support_samples), fill_value=-100, dtype=batch[0]["y_support"].dtype
+            ),
+            "x_query": torch.zeros((batch_size, max_query_samples, max_features), dtype=batch[0]["x_query"].dtype),
+            "y_query": torch.full((batch_size, max_query_samples), fill_value=-100, dtype=batch[0]["y_query"].dtype),
+            "padding_features": torch.ones((batch_size, max_features), dtype=torch.bool),
+            "padding_obs_support": torch.ones((batch_size, max_support_samples), dtype=torch.bool),
+            "padding_obs_query": torch.ones((batch_size, max_query_samples), dtype=torch.bool),
         }
 
         for i, dataset in enumerate(batch):
-            tensor_dict['x_support'][i, :dataset['x_support'].shape[0], :dataset['x_support'].shape[1]] = dataset['x_support']
-            tensor_dict['y_support'][i, :dataset['y_support'].shape[0]] = dataset['y_support']
-            tensor_dict['x_query'][i, :dataset['x_query'].shape[0], :dataset['x_support'].shape[1]] = dataset['x_query']
-            tensor_dict['y_query'][i, :dataset['y_query'].shape[0]] = dataset['y_query']
-            tensor_dict['padding_features'][i, :dataset['x_support'].shape[1]] = False
-            tensor_dict['padding_obs_support'][i, :dataset['x_support'].shape[0]] = False
-            tensor_dict['padding_obs_query'][i, :dataset['x_query'].shape[0]] = False
+            tensor_dict["x_support"][i, : dataset["x_support"].shape[0], : dataset["x_support"].shape[1]] = dataset[
+                "x_support"
+            ]
+            tensor_dict["y_support"][i, : dataset["y_support"].shape[0]] = dataset["y_support"]
+            tensor_dict["x_query"][i, : dataset["x_query"].shape[0], : dataset["x_support"].shape[1]] = dataset[
+                "x_query"
+            ]
+            tensor_dict["y_query"][i, : dataset["y_query"].shape[0]] = dataset["y_query"]
+            tensor_dict["padding_features"][i, : dataset["x_support"].shape[1]] = False
+            tensor_dict["padding_obs_support"][i, : dataset["x_support"].shape[0]] = False
+            tensor_dict["padding_obs_query"][i, : dataset["x_query"].shape[0]] = False
 
         return tensor_dict