autogluon.tabular 1.3.2b20250715__py3-none-any.whl → 1.3.2b20250716__py3-none-any.whl

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

@@ -0,0 +1 @@
+# Internal modules for MitraModel 

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

@@ -0,0 +1 @@
+# Configuration modules for MitraModel 

autogluon/tabular/models/mitra/_internal/config/config_run.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Self
 
 import torch
 
 from ..._internal.config.config_pretrain import ConfigSaveLoadMixin
 from ..._internal.config.enums import ModelName
 
+
 @dataclass
 class ConfigRun(ConfigSaveLoadMixin):
     device: torch.device
@@ -22,11 +22,11 @@ class ConfigRun(ConfigSaveLoadMixin):
         seed: int, 
         model_name: ModelName, 
         hyperparams: dict
-    ) -> Self:
+    ) -> "ConfigRun":
        
         return cls(
             device=device,
             seed=seed,
             model_name=model_name,
             hyperparams=hyperparams
-    )
+        )

autogluon/tabular/models/mitra/_internal/config/enums.py

@@ -1,4 +1,21 @@
-from enum import IntEnum, StrEnum
+from enum import IntEnum
+
+try:
+    from enum import StrEnum
+except ImportError:
+    # StrEnum is not available in Python < 3.11, so we create a compatible version
+    from enum import Enum
+    class StrEnum(str, Enum):
+        """
+        Enum where members are also (and must be) strings
+        """
+        def __new__(cls, value):
+            if not isinstance(value, str):
+                raise TypeError(f"{value!r} is not a string")
+            return super().__new__(cls, value)
+
+        def __str__(self):
+            return self.value
 
 
 class Task(StrEnum):
@@ -95,7 +112,7 @@ class BenchmarkName(StrEnum):
     NUMERICAL_CLASSIFICATION_LARGE = "numerical_classification_large"
     CATEGORICAL_REGRESSION_LARGE = "categorical_regression_large"
     NUMERICAL_REGRESSION_LARGE = "numerical_regression_large"
-    
+
     TABZILLA_HARD = "tabzilla_hard"
     TABZILLA_HARD_MAX_TEN_CLASSES = "tabzilla_hard_max_ten_classes"
     TABZILLA_HAS_COMPLETED_RUNS = "tabzilla_has_completed_runs"
@@ -142,4 +159,4 @@ class MetricName(StrEnum):
 class LossName(StrEnum):
     CROSS_ENTROPY = "cross_entropy"
     MSE = "mse"
-    MAE = "mae"
+    MAE = "mae"

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

@@ -0,0 +1 @@
+# Core modules for MitraModel 

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

@@ -1,10 +1,11 @@
-import torch
 import einops
+import torch
 
 from ..._internal.config.config_pretrain import ConfigPretrain
 from ..._internal.config.config_run import ConfigRun
 from ..._internal.config.enums import LossName, Task
 
+
 class CrossEntropyLossExtraBatch(torch.nn.Module):
 
     def __init__(self, label_smoothing: float):
@@ -28,28 +29,26 @@ class CrossEntropyLossExtraBatch(torch.nn.Module):
 
 def get_loss(cfg: ConfigRun):
 
-    match (cfg.task, cfg.hyperparams['regression_loss']):
-        case (Task.REGRESSION, LossName.MSE):
-            return torch.nn.MSELoss()
-        case (Task.REGRESSION, LossName.MAE):
-            return torch.nn.L1Loss()
-        case (Task.REGRESSION, LossName.CROSS_ENTROPY):
-            return CrossEntropyLossExtraBatch(cfg.hyperparams['label_smoothing'])
-        case (Task.CLASSIFICATION, _):
-            return CrossEntropyLossExtraBatch(cfg.hyperparams['label_smoothing'])
-        case (_, _):
-            raise ValueError(f"Unsupported task {cfg.task} and (regression) loss {cfg.hyperparams['regression_loss']}")
+    if cfg.task == Task.REGRESSION and cfg.hyperparams['regression_loss'] == LossName.MSE:
+        return torch.nn.MSELoss()
+    elif cfg.task == Task.REGRESSION and cfg.hyperparams['regression_loss'] == LossName.MAE:
+        return torch.nn.L1Loss()
+    elif cfg.task == Task.REGRESSION and cfg.hyperparams['regression_loss'] == LossName.CROSS_ENTROPY:
+        return CrossEntropyLossExtraBatch(cfg.hyperparams['label_smoothing'])
+    elif cfg.task == Task.CLASSIFICATION:
+        return CrossEntropyLossExtraBatch(cfg.hyperparams['label_smoothing'])
+    else:
+        raise ValueError(f"Unsupported task {cfg.task} and (regression) loss {cfg.hyperparams['regression_loss']}")
 
 def get_loss_pretrain(cfg: ConfigPretrain):
 
-    match (cfg.data.task, cfg.optim.regression_loss):
-        case (Task.REGRESSION, LossName.MSE):
-            return torch.nn.MSELoss()
-        case (Task.REGRESSION, LossName.MAE):
-            return torch.nn.L1Loss()
-        case (Task.REGRESSION, LossName.CROSS_ENTROPY):
-            return CrossEntropyLossExtraBatch(cfg.optim.label_smoothing)
-        case (Task.CLASSIFICATION, _):
-            return CrossEntropyLossExtraBatch(cfg.optim.label_smoothing)
-        case (_, _):
-            raise ValueError(f"Unsupported task {cfg.data.task} and (regression) loss {cfg.optim.regression_loss}")
+    if cfg.data.task == Task.REGRESSION and cfg.optim.regression_loss == LossName.MSE:
+        return torch.nn.MSELoss()
+    elif cfg.data.task == Task.REGRESSION and cfg.optim.regression_loss == LossName.MAE:
+        return torch.nn.L1Loss()
+    elif cfg.data.task == Task.REGRESSION and cfg.optim.regression_loss == LossName.CROSS_ENTROPY:
+        return CrossEntropyLossExtraBatch(cfg.optim.label_smoothing)
+    elif cfg.data.task == Task.CLASSIFICATION:
+        return CrossEntropyLossExtraBatch(cfg.optim.label_smoothing)
+    else:
+        raise ValueError(f"Unsupported task {cfg.data.task} and (regression) loss {cfg.optim.regression_loss}")

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

@@ -1,14 +1,13 @@
 from dataclasses import dataclass
-from typing import Self
 
 import numpy as np
-import scipy
+import scipy.special
 import torch
 from loguru import logger
 from sklearn.metrics import f1_score, mean_squared_error, r2_score, roc_auc_score, root_mean_squared_error
 
-from ..._internal.data.preprocessor import Preprocessor
 from ..._internal.config.enums import MetricName, Task
+from ..._internal.data.preprocessor import Preprocessor
 
 
 @dataclass
@@ -20,21 +19,20 @@ class PredictionMetrics():
 
 
     @classmethod
-    def from_prediction(cls, y_pred: np.ndarray, y_true: np.ndarray, task: Task) -> Self:
+    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]:
 
-    match task:
-        case Task.CLASSIFICATION:
-            return compute_classification_metrics(y_pred, y_true)
-        case Task.REGRESSION:
-            return compute_regression_metrics(y_pred, y_true)
-        
+    if task == Task.CLASSIFICATION:
+        return compute_classification_metrics(y_pred, y_true)
+    elif task == Task.REGRESSION:
+        return compute_regression_metrics(y_pred, y_true)
+
 
 def compute_classification_metrics(y_pred: np.ndarray, y_true: np.ndarray) -> tuple[float, float, dict]:
     # predictions are assumed to be log-probabilities
@@ -121,7 +119,7 @@ class PredictionMetricsTracker():
             y_true_ori = self.preprocessor.inverse_transform_y(y_true_np)
         else:
             y_true_ori = y_true_np
-    
+
         self.ys_pred.append(y_pred_ori)
         self.ys_true.append(y_true_ori)
 
@@ -131,4 +129,4 @@ class PredictionMetricsTracker():
         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

@@ -1,26 +1,26 @@
 import time
+
 import numpy as np
 import torch
 from loguru import logger
 from sklearn.base import BaseEstimator
-import torch.nn.functional as F
 
 from ..._internal.config.config_run import ConfigRun
+from ..._internal.config.enums import LossName, MetricName, ModelName, Task
 from ..._internal.core.callbacks import Checkpoint, EarlyStopping
-from ..._internal.data.collator import CollatorWithPadding
-from ..._internal.config.enums import MetricName, ModelName, Task, LossName
 from ..._internal.core.get_loss import get_loss
-from ..._internal.core.get_optimizer import get_optimizer, GradScaler
+from ..._internal.core.get_optimizer import GradScaler, get_optimizer
 from ..._internal.core.get_scheduler import get_scheduler
+from ..._internal.core.prediction_metrics import PredictionMetrics, PredictionMetricsTracker
+from ..._internal.data.collator import CollatorWithPadding
 from ..._internal.data.dataset_finetune import DatasetFinetune, DatasetFinetuneGenerator
 from ..._internal.data.preprocessor import Preprocessor
-from ..._internal.core.prediction_metrics import PredictionMetrics, PredictionMetricsTracker
 
 
 class TrainerFinetune(BaseEstimator):
 
     def __init__(
-            self, 
+            self,
             cfg: ConfigRun,
             model: torch.nn.Module,
             n_classes: int,
@@ -31,7 +31,7 @@ class TrainerFinetune(BaseEstimator):
         self.device = device
         self.model = model.to(self.device, non_blocking=True)
         self.n_classes = n_classes
-        
+
         self.loss = get_loss(self.cfg)
         self.optimizer = get_optimizer(self.cfg.hyperparams, self.model)
         self.scheduler_warmup, self.scheduler_reduce_on_plateau = get_scheduler(self.cfg.hyperparams, self.optimizer)
@@ -51,7 +51,7 @@ class TrainerFinetune(BaseEstimator):
             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'], 
+            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'],
@@ -62,19 +62,19 @@ class TrainerFinetune(BaseEstimator):
         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)  
+            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']
 
 
     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)       
+        self.preprocessor.fit(x_train, y_train)
 
-        x_train_transformed = self.preprocessor.transform_X(x_train) 
+        x_train_transformed = self.preprocessor.transform_X(x_train)
         y_train_transformed = self.preprocessor.transform_y(y_train)
-        
+
         dataset_train_generator = DatasetFinetuneGenerator(
             self.cfg,
             x = x_train_transformed,
@@ -94,16 +94,16 @@ class TrainerFinetune(BaseEstimator):
 
         for epoch in range(1, self.cfg.hyperparams['max_epochs']+1):
 
-            dataset_train = next(dataset_train_generator)            
+            dataset_train = next(dataset_train_generator)
             loader_train = self.make_loader(dataset_train, training=True)
             self.model.train()
-            
+
             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)
@@ -118,13 +118,12 @@ class TrainerFinetune(BaseEstimator):
                         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)
-                    
-                    match self.cfg.model_name:
-                        case ModelName.TABPFN:
-                            y_hat = self.model(x_support, y_support, x_query, task=self.cfg.task).squeeze(-1)
-                        case 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)
-                    
+
+                    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)
+
                     # Convert numerical y_query to bin ids
                     if self.cfg.task == Task.REGRESSION and self.cfg.hyperparams['regression_loss'] == LossName.CROSS_ENTROPY:
                         loss = self.loss(y_hat, y_query_bin_ids)
@@ -143,7 +142,7 @@ class TrainerFinetune(BaseEstimator):
 
                 # Convert bin id predictions to numerical values
                 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 = torch.argmax(y_hat, dim=-1)
                     y_hat = self.bins[y_hat] + self.bin_width / 2
 
                 y_hat = y_hat.float()
@@ -153,12 +152,12 @@ class TrainerFinetune(BaseEstimator):
                     prediction_metrics_tracker.update(y_hat, y_query, train=False)
 
             metrics_train = prediction_metrics_tracker.get_metrics()
-            metrics_valid = self.evaluate(x_train, y_train, x_val, y_val)   
+            metrics_valid = self.evaluate(x_train, y_train, x_val, y_val)
 
             self.log_metrics(epoch, metrics_train, metrics_valid)
 
             self.checkpoint(self.model, metrics_valid.loss)
-            
+
             self.early_stopping(metrics_valid.metrics[self.metric])
             if self.early_stopping.we_should_stop():
                 logger.info("Early stopping")
@@ -175,9 +174,9 @@ class TrainerFinetune(BaseEstimator):
 
         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:
-        
+
         self.model.eval()
 
         x_support_transformed = self.preprocessor.transform_X(x_support)
@@ -186,9 +185,9 @@ class TrainerFinetune(BaseEstimator):
         # y_query_transformed = self.preprocessor.transform_y(y_query)
 
         dataset = DatasetFinetune(
-            self.cfg, 
-            x_support = x_support_transformed, 
-            y_support = y_support_transformed, 
+            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'],
@@ -202,7 +201,7 @@ class TrainerFinetune(BaseEstimator):
             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)
@@ -215,16 +214,15 @@ class TrainerFinetune(BaseEstimator):
                     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)
-                    
-                    match self.cfg.model_name:
-                        case ModelName.TABPFN:
-                            y_hat = self.model(x_s, y_s, x_q, task=self.cfg.task).squeeze(-1)
-                        case ModelName.TAB2D | ModelName.TAB2D_COL_ROW | ModelName.TAB2D_SDPA:
-                            y_hat = self.model(x_s, y_s, x_q, padding_features, padding_obs_support, padding_obs_query)
+
+                    if self.cfg.model_name == ModelName.TABPFN:
+                        y_hat = self.model(x_s, y_s, x_q, 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_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:
-                    y_hat = torch.argmax(y_hat, dim=-1)  
+                    y_hat = torch.argmax(y_hat, dim=-1)
                     y_hat = self.bins[y_hat] + self.bin_width / 2
 
                 y_hat = y_hat.float()
@@ -232,7 +230,7 @@ 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:
 
@@ -241,9 +239,9 @@ class TrainerFinetune(BaseEstimator):
         y_support_transformed = self.preprocessor.transform_y(y_support)
 
         dataset = DatasetFinetune(
-            self.cfg, 
-            x_support = x_support_transformed, 
-            y_support = y_support_transformed, 
+            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'],
@@ -259,7 +257,7 @@ class TrainerFinetune(BaseEstimator):
             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)
@@ -271,18 +269,17 @@ class TrainerFinetune(BaseEstimator):
                     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)
-                    
-                    match self.cfg.model_name:
-                        case ModelName.TABPFN:
-                            y_hat = self.model(x_s, y_s, x_q, task=self.cfg.task).squeeze(-1)
-                        case ModelName.TAB2D | ModelName.TAB2D_COL_ROW | ModelName.TAB2D_SDPA:
-                            y_hat = self.model(x_s, y_s, x_q, padding_features, padding_obs_support, padding_obs_query)
-                    
+
+                    if self.cfg.model_name == ModelName.TABPFN:
+                        y_hat = self.model(x_s, y_s, x_q, 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_s, y_s, x_q, padding_features, padding_obs_support, padding_obs_query)
+
                 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:
-                    y_hat = np.argmax(y_hat, axis=-1)  
+                    y_hat = np.argmax(y_hat, axis=-1)
                     y_hat = (self.bins[y_hat] + self.bin_width / 2).cpu().numpy()
 
                 y_hat = self.preprocessor.inverse_transform_y(y_hat)
@@ -291,21 +288,20 @@ class TrainerFinetune(BaseEstimator):
         y_pred = np.concatenate(y_pred_list, axis=0)
 
         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:
 
-        match self.cfg.model_name:
-            case ModelName.TABPFN:
-                pad_to_max_features = True
-            case ModelName.TAB2D | ModelName.TAB2D_COL_ROW | ModelName.TAB2D_SDPA:
-                pad_to_max_features = False
-            case _:
-                raise NotImplementedError(f"Model {self.cfg.model_name} not implemented")
+        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]:
+            pad_to_max_features = False
+        else:
+            raise NotImplementedError(f"Model {self.cfg.model_name} not implemented")
 
         return torch.utils.data.DataLoader(
             dataset,
@@ -319,13 +315,12 @@ class TrainerFinetune(BaseEstimator):
                 pad_to_max_features=pad_to_max_features
             ),
         )
-    
+
 
     def log_start_metrics(self, metrics_valid: PredictionMetrics):
 
-        match self.cfg.task:
-            case Task.REGRESSION:
-                logger.info((
+        if self.cfg.task == Task.REGRESSION:
+            logger.info((
                     f"Epoch 000 "
                     f"| Train MSE: -.---- "
                     f"| Train MAE: -.---- "
@@ -334,21 +329,20 @@ class TrainerFinetune(BaseEstimator):
                     f"| Val MAE: {metrics_valid.metrics[MetricName.MAE]:.4f} "
                     f"| Val r2: {metrics_valid.metrics[MetricName.R2]:.4f}"
                 ))
-            case Task.CLASSIFICATION:
-                logger.info((
+
+        elif self.cfg.task == Task.CLASSIFICATION:
+            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):
 
-        match self.cfg.task:
-            case Task.REGRESSION:
-                logger.info((
+        if self.cfg.task == Task.REGRESSION:
+            logger.info((
                     f"Epoch {epoch:03d} "
                     f"| Train MSE: {metrics_train.metrics[MetricName.MSE]:.4f} "
                     f"| Train MAE: {metrics_train.metrics[MetricName.MAE]:.4f} "
@@ -357,11 +351,11 @@ class TrainerFinetune(BaseEstimator):
                     f"| Val MAE: {metrics_valid.metrics[MetricName.MAE]:.4f} "
                     f"| Val r2: {metrics_valid.metrics[MetricName.R2]:.4f}"
                 ))
-            case Task.CLASSIFICATION:
-                logger.info((
+        elif self.cfg.task == Task.CLASSIFICATION:
+            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

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