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

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

@@ -0,0 +1,108 @@
+import torch
+from torch.optim import SGD, Adam, AdamW
+
+from ..._internal.config.config_pretrain import ConfigPretrain
+
+
+def get_optimizer(hyperparams: dict, model: torch.nn.Module) -> torch.optim.Optimizer:
+
+    optimizer: torch.optim.Optimizer
+
+    if hyperparams['optimizer'] == "adam":
+        optimizer = Adam(
+            model.parameters(), 
+            lr=hyperparams['lr'],
+            betas=(0.9, 0.999),
+            weight_decay=hyperparams['weight_decay']
+        )
+    elif hyperparams['optimizer'] == "adamw":
+        optimizer = AdamW(
+            model.parameters(), 
+            lr=hyperparams['lr'],
+            betas=(0.9, 0.999),
+            weight_decay=hyperparams['weight_decay']
+        )
+    elif hyperparams['optimizer'] == "sgd":
+        optimizer = SGD(
+            model.parameters(),
+            lr=hyperparams['lr'],
+            weight_decay=hyperparams['weight_decay']
+        )
+    else:
+        raise ValueError("Optimizer not recognized")
+    
+    return optimizer
+
+
+def get_optimizer_pretrain(cfg: ConfigPretrain, model: torch.nn.Module) -> torch.optim.Optimizer:
+
+    parameters = [(name, param) for name, param in model.named_parameters()]
+
+    parameters_with_weight_decay = []
+    parameters_without_weight_decay = []
+
+    for name, param in parameters:
+        if name.endswith("bias") or 'norm' in name or 'embedding' in name:
+            parameters_without_weight_decay.append(param)
+        else:
+            parameters_with_weight_decay.append(param)
+
+    optimizer_parameters = [
+        {"params": parameters_with_weight_decay, "weight_decay": cfg.optim.weight_decay},
+        {"params": parameters_without_weight_decay, "weight_decay": 0.0},
+    ]
+
+    optimizer = torch.optim.AdamW(
+        optimizer_parameters, 
+        lr=cfg.optim.lr,
+        betas=(cfg.optim.beta1, cfg.optim.beta2),
+        weight_decay=cfg.optim.weight_decay
+    )
+    
+    return optimizer
+
+
+class GradScaler(torch.amp.GradScaler):
+
+    def __init__(
+        self, 
+        enabled: bool = True,
+        scale_init: float = 2.**16,
+        scale_min: float = 1.,
+        growth_interval: int = 2000,
+        device: str = 'cuda'
+    ):
+        super().__init__(enabled=enabled, device="cpu", init_scale=scale_init, growth_interval=growth_interval) # type: ignore
+        self._enabled = enabled
+        self.scale_min = scale_min
+        self.device = device
+
+        if not self._enabled:
+            # We write scale=1 to log if the scaler is disabled
+            self._scale = torch.tensor((1,), dtype=torch.float32, device=self.device)
+
+
+    def update(self):
+
+        if not self._enabled:
+            return
+
+        super().update()
+
+        if self._scale < self.scale_min:
+            super().update(self.scale_min)
+
+
+def move_optimizer_to(optim, device):
+    for param in optim.state.values():
+        # Not sure there are any global tensors in the state dict
+        if isinstance(param, torch.Tensor):
+            param.data = param.data.to(device)
+            if param._grad is not None:
+                param._grad.data = param._grad.data.to(device)
+        elif isinstance(param, dict):
+            for subparam in param.values():
+                if isinstance(subparam, torch.Tensor):
+                    subparam.data = subparam.data.to(device)
+                    if subparam._grad is not None:
+                        subparam._grad.data = subparam._grad.data.to(device)

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

@@ -0,0 +1,67 @@
+import torch
+from torch.optim.lr_scheduler import ReduceLROnPlateau, LinearLR
+from transformers import get_constant_schedule_with_warmup
+from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup
+
+from ..._internal.config.config_pretrain import ConfigPretrain
+
+
+def get_scheduler(hyperparams: dict, optimizer: torch.optim.Optimizer) -> tuple[torch.optim.lr_scheduler.LambdaLR, ReduceLROnPlateau]:
+
+    warmup_steps = hyperparams['warmup_steps']
+
+    # if warmup_steps > 0:
+    #     scheduler_warmup = torch.optim.lr_scheduler.LambdaLR(
+    #         optimizer, lambda step: min((step + 1) / warmup_steps, 1.0)
+    #     )
+    # else:
+    #     scheduler_warmup = torch.optim.lr_scheduler.LambdaLR(
+    #         optimizer, lambda step: 1.0
+    #     )
+    
+    if warmup_steps > 0:
+        scheduler_warmup = LinearLR(
+            optimizer,
+            start_factor=1.0 / warmup_steps,
+            end_factor=1.0,
+            total_iters=warmup_steps,
+        )
+    else:
+        scheduler_warmup = torch.optim.lr_scheduler.ConstantLR(optimizer, factor=1.0, total_iters=1)
+
+    if hyperparams['lr_scheduler']:      
+        scheduler_reduce_on_plateau = ReduceLROnPlateau(
+            optimizer, 
+            patience=hyperparams['lr_scheduler_patience'], 
+            min_lr=0, 
+            factor=0.2
+        )
+    else:
+        # With ReduceLROnPlateau, the scheduler accepts a metric to monitor, so our dummy metric must also be a ReduceLRonPlateau scheduler
+        scheduler_reduce_on_plateau = ReduceLROnPlateau(
+            optimizer, 
+            patience=1000000000, 
+            min_lr=0, 
+            factor=0.2
+        )
+
+    return scheduler_warmup, scheduler_reduce_on_plateau
+
+
+def get_scheduler_pretrain(cfg: ConfigPretrain, optimizer: torch.optim.Optimizer):
+
+    
+    if cfg.optim.cosine_scheduler:
+        schedule = get_cosine_with_min_lr_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=cfg.optim.warmup_steps,
+            num_training_steps=cfg.optim.steps,
+            min_lr_rate=0.1
+        )
+    else:
+        schedule = get_constant_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=cfg.optim.warmup_steps
+        )
+
+    return schedule

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

@@ -0,0 +1,134 @@
+from dataclasses import dataclass
+from typing import Self
+
+import numpy as np
+import scipy
+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
+
+
+@dataclass
+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) -> Self:
+
+        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)
+        
+
+def compute_classification_metrics(y_pred: np.ndarray, y_true: np.ndarray) -> tuple[float, float, dict]:
+    # predictions are assumed to be log-probabilities
+
+    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
+    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()
+    }
+
+    loss = metrics[MetricName.LOG_LOSS]
+    score = metrics[MetricName.ACCURACY]
+
+    return loss, score, metrics
+
+
+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')
+
+    try:
+        if y_pred_proba.shape[1] == 2:
+            return roc_auc_score(y_true, y_pred_proba[:, 1])
+        else:
+            return roc_auc_score(y_true, y_pred_proba, multi_class=multi_class, average=average, labels=labels)
+    except ValueError as e:
+        logger.error(f"Error computing roc_auc_score: {e}")
+        logger.error(f"Returning {-1}")
+        return -1
+
+
+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)
+    }
+
+    loss = metrics[MetricName.MSE]
+    score = metrics[MetricName.R2]
+
+    return loss, score, metrics
+
+
+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)
+
+        y_true_np = y_true.detach().cpu().numpy()[0]
+        if train:
+            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)
+
+
+    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)

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

@@ -0,0 +1,367 @@
+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.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_scheduler import get_scheduler
+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, 
+            cfg: ConfigRun,
+            model: torch.nn.Module,
+            n_classes: int,
+            device: str
+        ) -> None:
+
+        self.cfg = cfg
+        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)
+        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
+        )
+
+        self.early_stopping = EarlyStopping(patience=self.cfg.hyperparams['early_stopping_patience'])
+        self.checkpoint = Checkpoint()
+        self.preprocessor = Preprocessor(
+            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
+        )
+
+        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.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)       
+
+        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,
+            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']
+        )
+
+        self.checkpoint.reset(self.model)
+
+        metrics_valid = self.evaluate(x_train, y_train, x_val, y_val)
+        self.log_start_metrics(metrics_valid)
+        self.checkpoint(self.model, metrics_valid.loss)
+
+        start_time = time.time()
+
+        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()
+            
+            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)
+
+                    # Convert numerical y_support to bin ids
+                    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_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)
+                    
+                    # 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)
+                    elif self.cfg.task == Task.CLASSIFICATION:
+                        # for b in range(y_support.shape[0]):
+                        #     unique_classes = len(torch.unique(torch.cat((y_support[b], y_query[b]))))
+                        #     y_hat[b, :, unique_classes:] = 0
+                        loss = self.loss(y_hat, y_query)
+                    else:
+                        loss = self.loss(y_hat, y_query)
+
+                self.optimizer.zero_grad()
+                self.scaler.scale(loss).backward()
+                self.scaler.step(self.optimizer)
+                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:
+                    y_hat = torch.argmax(y_hat, dim=-1)  
+                    y_hat = self.bins[y_hat] + self.bin_width / 2
+
+                y_hat = y_hat.float()
+                if self.cfg.task == Task.REGRESSION:
+                    prediction_metrics_tracker.update(y_hat, y_query, train=True)
+                else:
+                    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)   
+
+            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")
+                break
+
+            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']:
+                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:
+        
+        self.model.eval()
+
+        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)
+        # y_query_transformed = self.preprocessor.transform_y(y_query)
+
+        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'],
+        )
+
+        loader = self.make_loader(dataset, training=False)
+        prediction_metrics_tracker = PredictionMetricsTracker(task=self.cfg.task, preprocessor=self.preprocessor)
+
+        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)
+
+                    # Convert numerical y_support to bin ids
+                    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)
+
+                # 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 = self.bins[y_hat] + self.bin_width / 2
+
+                y_hat = y_hat.float()
+                prediction_metrics_tracker.update(y_hat, y_q, train=False)
+
+        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'],
+        )
+
+        loader = self.make_loader(dataset, training=False)
+        self.model.eval()
+
+        y_pred_list = []
+
+        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)
+
+                    # Convert numerical y_support to bin ids
+                    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)
+                    
+                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 = (self.bins[y_hat] + self.bin_width / 2).cpu().numpy()
+
+                y_hat = self.preprocessor.inverse_transform_y(y_hat)
+                y_pred_list.append(y_hat)
+
+        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")
+
+        return torch.utils.data.DataLoader(
+            dataset,
+            batch_size=1,
+            shuffle=training,
+            pin_memory=True,
+            num_workers=0,
+            drop_last=False,
+            collate_fn=CollatorWithPadding(
+                max_features=self.cfg.hyperparams['dim_embedding'],
+                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((
+                    f"Epoch 000 "
+                    f"| Train MSE: -.---- "
+                    f"| Train MAE: -.---- "
+                    f"| Train r2: -.---- "
+                    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}"
+                ))
+            case 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((
+                    f"Epoch {epoch:03d} "
+                    f"| Train MSE: {metrics_train.metrics[MetricName.MSE]:.4f} "
+                    f"| Train MAE: {metrics_train.metrics[MetricName.MAE]:.4f} "
+                    f"| Train r2: {metrics_train.metrics[MetricName.R2]:.4f} "
+                    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}"
+                ))
+            case 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}"
+                ))