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

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}"
+                ))

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

@@ -0,0 +1,46 @@
+import torch
+
+
+class CollatorWithPadding():
+
+    def __init__(
+            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)
+
+        if self.pad_to_max_features:
+            max_features = self.max_features
+
+        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),
+        }
+
+        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
+
+        return tensor_dict

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

@@ -0,0 +1,132 @@
+from typing import Optional
+
+import numpy as np
+import torch
+
+from ..._internal.config.config_run import ConfigRun
+from ..._internal.data.dataset_split import make_dataset_split
+from ..._internal.config.enums import Task
+
+
+class DatasetFinetune(torch.utils.data.Dataset):
+    """
+    The main goal of this class is to generate a dataset for fine-tuning.
+    The input data are the full (x_support, y_support, x_query, y_query)
+    But these arrays are too large to be pushed through the model at once.
+    So here we split query the data into chunks if the query data is too large.
+    If the support data is too large, we randomly sample from it.
+    Furthermore, we transition from numpy to tensors.
+    """
+
+    def __init__(
+        self, 
+        cfg: ConfigRun,
+        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
+    ):
+        """
+        :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.y_query = y_query
+
+        if self.y_query is None:
+            self.y_query = np.zeros((self.x_query.shape[0],)) - 1
+
+        self.max_samples_support = max_samples_support
+        self.max_samples_query = max_samples_query
+
+        self.x_queries = self.split_in_chunks(self.x_query, max_samples_query)
+        self.y_queries = self.split_in_chunks(self.y_query, max_samples_query)
+
+        self.n_samples_support = self.x_support.shape[0]
+
+        # 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
+        )
+
+        x_support = self.x_support[support_indices]
+        y_support = self.y_support[support_indices]
+
+        x_support_tensor = torch.as_tensor(x_support)
+        y_support_tensor = torch.as_tensor(y_support)
+        x_query_tensor = torch.as_tensor(self.x_queries[idx])
+        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,
+        }
+    
+
+
+    def split_in_chunks(self, x: np.ndarray, batch_size: int) -> list[np.ndarray]:
+        """
+        Splits the data into chunks of size batch_size
+        """
+
+        n_chunks = int(np.ceil(x.shape[0] / batch_size))
+        x_chunks = []
+
+        for i in range(n_chunks):
+            x_chunks.append(x[i * batch_size: (i + 1) * batch_size])
+
+        return x_chunks
+
+def DatasetFinetuneGenerator(
+    cfg: ConfigRun,
+    x: np.ndarray, 
+    y: np.ndarray, 
+    task: Task,
+    max_samples_support: int,
+    max_samples_query: int
+):
+    """
+    The dataset fine-tune generator is a generator that yields a dataset for fine-tuning.
+    The idea is to split the training dataset into a support and query set.
+    Every single iteration, the generator yields a different support and query set split.
+    The dataset made always has exactly one batch.
+    """
+        
+    while True:
+
+        x_support, x_query, y_support, y_query = make_dataset_split(x=x, y=y, task=task, seed=cfg.seed)
+        n_samples_support = x_support.shape[0]
+        n_samples_query = x_query.shape[0]
+
+        support_size = min(max_samples_support, n_samples_support)
+        query_size = min(max_samples_query, n_samples_query)
+
+        dataset_finetune = DatasetFinetune(
+            cfg=cfg,
+            x_support=x_support[:support_size],
+            y_support=y_support[:support_size],
+            x_query=x_query[:query_size],
+            y_query=y_query[:query_size],
+            max_samples_support=max_samples_support,
+            max_samples_query=max_samples_query,
+        )
+
+        yield dataset_finetune