multabench 0.0.1__py3-none-any.whl

multabench/__init__.py

@@ -0,0 +1,2 @@
+from multabench.benchmark.load import load_embeddings, load_split
+from multabench.benchmark.splits import Size

multabench/baselines/abstract_model.py

@@ -0,0 +1,240 @@
+from typing import Tuple, Dict, Optional, Set, List, Any
+
+import numpy as np
+import torch
+from pandas import DataFrame, Series
+from sklearn.preprocessing import LabelEncoder
+from skrub import DatetimeEncoder
+
+from multabench.baselines.preprocessing.dates import fit_date_encoders, transform_date_features
+from multabench.baselines.preprocessing.feature_types import detect_feature_types, transform_feature_types
+from multabench.baselines.preprocessing.nulls import raise_if_null_target
+from multabench.baselines.preprocessing.validation import split_to_val
+from multabench.baselines.preprocessing.image_embeddings import (
+    transform_image_features,
+    fit_image_encoders,
+)
+from multabench.baselines.training.metrics import calculate_metric
+from multabench.datasets.all_datasets import MulTaBenchDatasetID
+from multabench.baselines.preprocessing.categorical import fit_categorical_encoders, transform_categorical_features
+from multabench.baselines.preprocessing.numerical import fit_numerical_median, transform_numerical_features
+from multabench.baselines.preprocessing.target import transform_preprocess_y, fit_preprocess_y
+from multabench.dino.constants import DINOV3_SMALL
+from multabench.e5.constants import E5_SMALL_V2
+from multabench.baselines.preprocessing.text_embeddings import E5ColumnEncoder, fit_text_encoders, transform_text_features
+from multabench.datasets.objects import SupervisedTask
+from multabench.utils.warnings import silence_baselines_prints
+
+
+class TabularModel:
+
+    MODEL_NAME: str
+    SHORT_NAME: str
+    USE_VAL_SPLIT: bool
+    USE_MEDIAN_FILLING: bool
+    USE_CATEGORICAL_ENCODING: bool
+    USE_TEXT_EMBEDDINGS: bool
+    USE_TARGET_ENCODER: bool
+    NEEDS_GPU = True
+
+    def __init__(self, problem_type: SupervisedTask, device: torch.device,
+                 dataset: MulTaBenchDatasetID | None = None, verbose: bool = False, image_folder: str | None = None,
+                 image_column: str | None = None,
+                 tune_dino: bool = False,
+                 dino_train_kwargs: Optional[Dict[str, Any]] = None,
+                 dino_model_name: str = DINOV3_SMALL,
+                 tune_e5: bool = False,
+                 e5_train_kwargs: Optional[Dict[str, Any]] = None,
+                 e5_model_name: str = E5_SMALL_V2,
+                 **kwargs):
+        assert problem_type in {SupervisedTask.REGRESSION, SupervisedTask.BINARY, SupervisedTask.MULTICLASS}
+        self.problem_type = problem_type
+        self.dataset = dataset
+        self.is_cls = bool(problem_type in {SupervisedTask.BINARY, SupervisedTask.MULTICLASS})
+        self.device = device
+        self.best_val_loss: Optional[None] = None
+        self.verbose = verbose
+        self.d_output: int = 0
+        self.image_folder = image_folder
+        self.image_column = image_column
+        self.tune_dino = tune_dino
+        self.dino_train_kwargs = dino_train_kwargs or {}
+        self.dino_model_name = dino_model_name
+        self._tuned_dino_model = None
+        self._tuned_dino_processor = None
+        self.tune_e5 = tune_e5
+        self.e5_train_kwargs = e5_train_kwargs or {}
+        self.e5_model_name = e5_model_name
+        self.model_ = self.initialize_model()
+        self.target_transformer: Optional[LabelEncoder] = None
+        self.date_transformers: Dict[str, DatetimeEncoder] = {}
+        self.image_transformers: Dict[str, Any] = {}  # PCA per image column
+        self.numerical_features: Set[str] = set()
+        self.numerical_medians: Dict[str, float] = {}
+        self.categorical_features: Set[str] = set()
+        self.categorical_indices: List[int] = []
+        self.categorical_encoders: Dict[str, LabelEncoder] = {}
+        self.text_features: Set[str] = set()
+        self.text_transformers: Dict[str, E5ColumnEncoder] = {}
+        self._train_embeddings: Dict[str, np.ndarray] = {}
+
+    def initialize_model(self):
+        raise NotImplementedError("Initialize model method not implemented yet")
+
+    def fit(self, x: DataFrame, y: Series):
+        x_train, y_train = x.copy(), y.copy()
+        if self.USE_VAL_SPLIT:
+            x_train, x_val, y_train, y_val = split_to_val(x=x, y=y, is_cls=self.is_cls)
+        else:
+            x_val, y_val = None, None
+        self.fit_preprocessor(x_train=x_train, y_train=y_train)
+        x_train, y_train = self.transform_preprocessor(x=x_train, y=y_train, train_embeddings=self._train_embeddings)
+        self._train_embeddings = {}
+        if x_val is not None and y_val is not None:
+            x_val, y_val = self.transform_preprocessor(x=x_val, y=y_val)
+        self._print_feature_summary(x_train)
+        with silence_baselines_prints():
+            self.fit_model(x_train=x_train, y_train=y_train, x_val=x_val, y_val=y_val)
+
+    def fit_preprocessor(self, x_train: DataFrame, y_train: Series):
+        x_train, y_train = self.do_model_agnostic_preprocessing(x=x_train, y=y_train)
+        if self.USE_MEDIAN_FILLING:
+            self.numerical_medians = fit_numerical_median(x=x_train, numerical_features=self.numerical_features)
+        if self.USE_CATEGORICAL_ENCODING:
+            self.categorical_encoders = fit_categorical_encoders(x=x_train, categorical_features=self.categorical_features)
+        if self.USE_TEXT_EMBEDDINGS:
+            self.text_transformers, text_embeddings = fit_text_encoders(
+                x=x_train,
+                text_features=self.text_features,
+                device=self.device,
+                y=y_train if self.tune_e5 else None,
+                tune_e5=self.tune_e5,
+                e5_train_kwargs=self.e5_train_kwargs if self.tune_e5 else None,
+                is_cls=self.is_cls,
+                d_output=self.d_output,
+                e5_model_name=self.e5_model_name,
+                dataset=self.dataset,
+            )
+            self._train_embeddings.update(text_embeddings)
+            self.vprint(f"📝 Detected {len(self.text_transformers)} text features: {sorted(self.text_transformers)}")
+        self.fit_internal_preprocessor(x=x_train, y=y_train)
+
+    def transform_preprocessor(self, x: DataFrame, y: Optional[Series],
+                               train_embeddings: Optional[Dict[str, np.ndarray]] = None) -> Tuple[DataFrame, Optional[Series]]:
+        x = transform_date_features(x=x, date_transformers=self.date_transformers)
+        x = transform_image_features(
+            x=x,
+            image_transformers=self.image_transformers,
+            device=self.device,
+            image_folder=self.image_folder,
+            dino_model=getattr(self, "_tuned_dino_model", None),
+            dino_processor=getattr(self, "_tuned_dino_processor", None),
+            dino_model_name=self.dino_model_name,
+            train_embeddings=train_embeddings,
+        )
+        image_features = [f"{col}_img_pca_{i}" for col in self.image_transformers.keys() for i in range(self.image_transformers[col].n_components)]
+        x = transform_feature_types(x=x, numerical_features=self.numerical_features, image_features=image_features)
+        if y is not None:
+            raise_if_null_target(y)
+            if self.USE_TARGET_ENCODER:
+                y = transform_preprocess_y(y=y, scaler=self.target_transformer)
+        if self.USE_MEDIAN_FILLING:
+            x = transform_numerical_features(x=x, numerical_medians=self.numerical_medians)
+        if self.USE_CATEGORICAL_ENCODING:
+            x = transform_categorical_features(x=x, categorical_encoders=self.categorical_encoders)
+        if self.USE_TEXT_EMBEDDINGS:
+            x = transform_text_features(
+                x=x,
+                text_encoders=self.text_transformers,
+                device=self.device,
+                train_embeddings=train_embeddings,
+            )
+        return self.transform_internal_preprocessor(x=x, y=y)
+
+    def fit_internal_preprocessor(self, x: DataFrame, y: Series) -> Tuple[DataFrame, Series]:
+        pass
+
+    def transform_internal_preprocessor(self, x: DataFrame, y: Optional[Series]) -> Tuple[DataFrame, Optional[Series]]:
+        return x, y
+
+    def fit_model(self, x_train: DataFrame, y_train: Series, x_val: Optional[DataFrame], y_val: Optional[Series]):
+        raise NotImplementedError("Fit model method not implemented yet")
+
+    def do_model_agnostic_preprocessing(self, x: DataFrame, y: Series) -> Tuple[DataFrame, Series]:
+        raise_if_null_target(y)
+        feature_types = detect_feature_types(x, image_column=self.image_column)
+        self.date_transformers = fit_date_encoders(x=x, date_features=feature_types.date_features)
+        self.vprint(f"📅 Detected {len(self.date_transformers)} date features: {sorted(self.date_transformers)}")
+        # TODO: ConTextTab supports dates natively, perhaps we need to change this to "USE_DATE_TRANSFORMATION"
+        x = transform_date_features(x=x, date_transformers=self.date_transformers)
+        if self.is_cls:
+            self.d_output = len(set(y))
+        else:
+            self.d_output = 1
+        self.image_transformers, self._tuned_dino_model, self._tuned_dino_processor, self._train_embeddings = fit_image_encoders(
+            x=x,
+            image_features=feature_types.image_features,
+            device=self.device,
+            image_folder=self.image_folder,
+            y=y if self.tune_dino else None,
+            is_cls=self.is_cls,
+            d_output=self.d_output,
+            tune_dino=self.tune_dino,
+            dino_train_kwargs=self.dino_train_kwargs if self.tune_dino else None,
+            dino_model_name=self.dino_model_name,
+            dataset=self.dataset,
+        )
+        self.vprint(f"📷 Detected {len(self.image_transformers)} image features: {sorted(self.image_transformers)}")
+        non_image_columns = [col for col in x.columns if col not in self.image_transformers]
+        x = x[non_image_columns]
+        date_parts = {part for encoder in self.date_transformers.values() for part in encoder.get_feature_names_out()}
+        self.numerical_features = feature_types.numerical_features | date_parts
+        self.text_features = feature_types.text_features
+        self.categorical_features = feature_types.categorical_features
+        self.vprint(f"🔢 Detected {len(self.numerical_features)} numerical features: {sorted(self.numerical_features)}")
+        image_features = [f"{col}_img_pca_{i}" for col in self.image_transformers.keys() for i in range(self.image_transformers[col].n_components)]
+        x = transform_feature_types(x=x, numerical_features=self.numerical_features, image_features=image_features)
+        # Indices here are for the current x (no image/text PCA yet). After transform_preprocessor, column
+        # order changes (image cols → PCA, text cols → PCA), so indices no longer match. Models that need
+        # categorical columns at fit time (e.g. LightGBM) must resolve by name: [c for c in x.columns if c in self.categorical_features].
+        self.categorical_indices = [i for i, c in enumerate(x.columns) if c in self.categorical_features]
+        self.target_transformer = fit_preprocess_y(y=y, is_cls=self.is_cls)
+        # d_output already set above for fit_image_encoders
+        # TODO: drop constant columns, where constants means all values are the same (and no nulls)
+        return x, y
+
+    def predict(self, x: DataFrame) -> np.ndarray:
+        x, _ = self.transform_preprocessor(x=x, y=None)
+        return self.predict_from_processed(x=x)
+
+    def predict_from_processed(self, x: DataFrame) -> np.ndarray:
+        if not self.is_cls:
+            return self.model_.predict(x)
+        probs = self.model_.predict_proba(x)
+        if self.d_output == 2:
+            probs = probs[:, 1]
+        return probs
+
+    def score_all_metrics(self, X, y) -> Dict:
+        x = X.copy()
+        y = y.copy()
+        y_true = transform_preprocess_y(y=y, scaler=self.target_transformer)
+        y_pred = self.predict(x)
+        metrics = calculate_metric(y_true=y_true, y_pred=y_pred, d_output=self.d_output)
+        return metrics
+
+    def _print_feature_summary(self, x: DataFrame) -> None:
+        cols = list(x.columns)
+        n_img = sum(1 for c in cols if "_img_pca_" in c)
+        n_txt = sum(1 for c in cols if "_txt_pca_" in c)
+        n_tab = len(cols) - n_img - n_txt
+        parts = [f"tabular: {n_tab}"]
+        if n_img:
+            parts.append(f"img_emb: {n_img}")
+        if n_txt:
+            parts.append(f"txt_emb: {n_txt}")
+        print(f"🔭 {len(cols)} features → {self.SHORT_NAME}  [{' | '.join(parts)}]")
+
+    def vprint(self, s: str):
+        if self.verbose:
+            print(s)

multabench/baselines/autogluon_mm.py

@@ -0,0 +1,116 @@
+import hashlib
+import os
+import shutil
+import tempfile
+import time
+
+import numpy as np
+from pandas import DataFrame, Series, notna
+
+
+from multabench.baselines.preprocessing.target import fit_preprocess_y, transform_preprocess_y
+from multabench.baselines.abstract_model import TabularModel
+from multabench.baselines.preprocessing.feature_types import detect_feature_types
+from multabench.datasets.objects import SupervisedTask
+from multabench.e5.constants import E5_SMALL_V2
+
+# DINOv2 (timm format) — DINOv3 is not supported by timm/AutoGluon as of Apr 2026
+AGMM_IMAGE_MODEL = "vit_small_patch14_dinov2.lvd142m"
+AGMM_LABEL = "__target__"
+
+TASK2AG = {
+    SupervisedTask.BINARY: "binary",
+    SupervisedTask.MULTICLASS: "multiclass",
+    SupervisedTask.REGRESSION: "regression",
+}
+
+# Same metrics the benchmark scores with (baselines/training/metrics.py).
+TASK2METRIC = {
+    SupervisedTask.BINARY: "roc_auc",
+    SupervisedTask.MULTICLASS: "log_loss",
+    SupervisedTask.REGRESSION: "root_mean_squared_error",
+}
+
+
+def unique_run_dir(dataset_name: str, x: DataFrame) -> str:
+    ts = int(time.time())
+    data_hash = hashlib.md5(x.head(5).to_csv().encode()).hexdigest()[:8]
+    return tempfile.mkdtemp(prefix=f"agmm_{dataset_name}_{ts}_{data_hash}_")
+
+
+class AutoGluonMM(TabularModel):
+
+    MODEL_NAME = "AutoGluon-MM"
+    SHORT_NAME = "agmm"
+    USE_VAL_SPLIT = False
+    USE_MEDIAN_FILLING = False
+    USE_CATEGORICAL_ENCODING = False
+    USE_TEXT_EMBEDDINGS = False
+    USE_TARGET_ENCODER = True
+
+    def initialize_model(self):
+        return None
+
+    def fit_model(self, x_train: DataFrame, y_train: Series, x_val: DataFrame, y_val: Series):
+        raise NotImplementedError("AutoGluonMM uses a custom fit() — fit_model() is never called.")
+
+    def fit(self, x: DataFrame, y: Series):
+        from autogluon.multimodal import MultiModalPredictor
+        # Fit target encoder for consistent scoring with other baselines
+        self.target_transformer = fit_preprocess_y(y=y, is_cls=self.is_cls)
+        y_enc = transform_preprocess_y(y=y, scaler=self.target_transformer)
+        if self.is_cls:
+            self.d_output = len(set(y_enc))
+        else:
+            self.d_output = 1
+
+        # Detect modalities
+        feature_types = detect_feature_types(x, image_column=self.image_column)
+        image_cols = feature_types.image_features
+        text_cols = feature_types.text_features
+
+        # Expand image file paths
+        x_ag = x.copy()
+        if image_cols and self.image_folder:
+            for col in image_cols:
+                x_ag[col] = x_ag[col].apply(lambda f: os.path.join(self.image_folder, str(f)) if notna(f) else f)
+
+        # Build base model config based on detected modalities
+        has_image = bool(image_cols and self.image_folder)
+        has_text = bool(text_cols)
+        hparams = {}
+        if has_image:
+            hparams["model.timm_image.checkpoint_name"] = AGMM_IMAGE_MODEL
+        if has_text:
+            hparams["model.hf_text.checkpoint_name"] = E5_SMALL_V2
+
+        df_train = x_ag.copy()
+        df_train[AGMM_LABEL] = y_enc.values
+
+        output_dir = unique_run_dir(dataset_name=getattr(self.dataset, "name", "unknown"), x=x)
+        self.predictor_ = MultiModalPredictor(
+            label=AGMM_LABEL,
+            problem_type=TASK2AG[self.problem_type],
+            eval_metric=TASK2METRIC[self.problem_type],
+            path=output_dir,
+            pretrained=True,
+        )
+        self.predictor_.fit(df_train, hyperparameters=hparams, clean_ckpts=True)
+
+    def _prepare_x(self, x: DataFrame) -> DataFrame:
+        x_ag = x.copy()
+        if self.image_column and self.image_folder:
+            col = self.image_column
+            x_ag[col] = x_ag[col].apply(lambda f: os.path.join(self.image_folder, str(f)) if notna(f) else f)
+        return x_ag
+
+    def predict(self, x: DataFrame) -> np.ndarray:
+        x_ag = self._prepare_x(x)
+        if not self.is_cls:
+            result = self.predictor_.predict(x_ag).to_numpy()
+        else:
+            proba = self.predictor_.predict_proba(x_ag).to_numpy()
+            result = proba[:, 1] if self.d_output == 2 else proba
+        shutil.rmtree(self.predictor_.path, ignore_errors=True)
+        self.predictor_ = None
+        return result

multabench/baselines/benchmarks/evaluate.py

@@ -0,0 +1,109 @@
+import time
+from typing import Type, Dict
+
+import torch
+
+from multabench.datasets.all_datasets import MulTaBenchDatasetID
+from multabench.datasets.objects import MultimodalDataset
+from multabench.baselines.abstract_model import TabularModel
+from multabench.benchmark.load import load_multabench_dataset, split_dataset
+from multabench.benchmark.splits import Size
+from multabench.utils.hardware import get_hardware_dict
+from multabench.dino.constants import DINOV3_SMALL
+from multabench.e5.constants import E5_SMALL_V2
+from multabench.utils.logging import get_current_commit_hash
+from multabench.result_keys import ResultKey
+from multabench.utils.profiling import PeakMemoryTracker
+from multabench.utils.timing import pop_embedding_seconds
+
+def evaluate_on_loaded_dataset(model_cls: Type[TabularModel],
+                                dataset: MultimodalDataset,
+                                fold: int,
+                                device: torch.device,
+                                size: Size = Size.TEN_K,
+                                verbose: bool = False,
+                                tune_dino: bool = False,
+                                dino_train_kwargs: dict | None = None,
+                                dino_model_name: str = DINOV3_SMALL,
+                                tune_e5: bool = False,
+                                e5_train_kwargs: dict | None = None,
+                                e5_model_name: str = E5_SMALL_V2) -> Dict:
+    start_time = time.time()
+    dataset_id = dataset.dataset_id
+    split = split_dataset(dataset, fold=fold, size=size)
+    x_train, y_train = split.x_train, split.y_train
+    x_test, y_test = split.x_test, split.y_test
+    kwargs = dict(problem_type=dataset.task_type, device=device, verbose=verbose, dataset=dataset_id,
+                  image_folder=dataset.image_folder, image_column=dataset.image_column, tune_dino=tune_dino, dino_train_kwargs=dino_train_kwargs,
+                  dino_model_name=dino_model_name,
+                  tune_e5=tune_e5, e5_train_kwargs=e5_train_kwargs, e5_model_name=e5_model_name)
+    model = model_cls(**kwargs)
+    pop_embedding_seconds()
+    with PeakMemoryTracker(phase='train', device=device) as train_tracker:
+        model.fit(x_train, y_train)
+    train_embedding_s = pop_embedding_seconds()
+    with PeakMemoryTracker(phase='inference', device=device) as test_tracker:
+        metrics = model.score_all_metrics(X=x_test, y=y_test)
+    inference_embedding_s = pop_embedding_seconds()
+    runtime = time.time() - start_time
+    d_summary = {
+        ResultKey.TIMESTAMP: time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
+        ResultKey.GIT: get_current_commit_hash(),
+        ResultKey.MODEL: model_cls.MODEL_NAME,
+        ResultKey.DATASET: dataset_id.name,
+        ResultKey.TASK_TYPE: str(dataset_id.name)[:3],
+        ResultKey.TRAIN_TYPE: "benchmark",
+        ResultKey.FOLD: fold,
+        ResultKey.SIZE: size,
+        **metrics,
+        ResultKey.RUNTIME: runtime,
+        ResultKey.N_TRAIN: len(y_train),
+        ResultKey.N_TEST: len(y_test),
+        ResultKey.M_FEATURES: x_train.shape[1],
+        ResultKey.TUNE_DINO: tune_dino,
+        ResultKey.TUNE_E5: tune_e5,
+        ResultKey.E5_MODEL_NAME: e5_model_name,
+        ResultKey.BEST_VAL_LOSS: getattr(model, "best_val_loss", None),
+        **train_tracker.summary(),
+        **test_tracker.summary(),
+        ResultKey.TRAIN_TIME_PER_1K: (train_tracker.wall_time_s - train_embedding_s) / len(y_train) * 1000,
+        ResultKey.INFERENCE_TIME_PER_1K: (test_tracker.wall_time_s - inference_embedding_s) / len(y_test) * 1000,
+        ResultKey.TRAIN_EMBEDDING_TIME_PER_1K: train_embedding_s / len(y_train) * 1000,
+        ResultKey.INFERENCE_EMBEDDING_TIME_PER_1K: inference_embedding_s / len(y_test) * 1000,
+        **get_hardware_dict(device),
+        **(dino_train_kwargs or {}),
+        **(e5_train_kwargs or {}),
+    }
+    metric, error = metrics[ResultKey.METRIC], metrics[ResultKey.TEST_ERROR]
+    print(f"{dataset_id.name} fold {fold}: {metric} error {error:.4f} ({runtime:.0f}s)")
+    return d_summary
+
+
+def evaluate_on_dataset(model_cls: Type[TabularModel],
+                        dataset_id: MulTaBenchDatasetID,
+                        fold: int,
+                        device: torch.device,
+                        size: Size = Size.TEN_K,
+                        verbose: bool = False,
+                        tune_dino: bool = False,
+                        dino_train_kwargs: dict | None = None,
+                        dino_model_name: str | None = None,
+                        tune_e5: bool = False,
+                        e5_train_kwargs: dict | None = None,
+                        e5_model_name: str = E5_SMALL_V2) -> Dict:
+    print(f"Running model {model_cls.MODEL_NAME} over dataset {dataset_id} with fold {fold}")
+    dataset = load_multabench_dataset(dataset_id)
+    return evaluate_on_loaded_dataset(
+        model_cls=model_cls,
+        dataset=dataset,
+        fold=fold,
+        device=device,
+        size=size,
+        verbose=verbose,
+        tune_dino=tune_dino,
+        dino_train_kwargs=dino_train_kwargs,
+        dino_model_name=dino_model_name,
+        tune_e5=tune_e5,
+        e5_train_kwargs=e5_train_kwargs,
+        e5_model_name=e5_model_name,
+    )

multabench/baselines/catboost.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+
+from pandas import DataFrame, Series
+
+from multabench.constants import SEED
+from multabench.utils.devices import CPU_CORES
+from multabench.baselines.abstract_model import TabularModel
+
+
+@dataclass
+class CatBoostDefaultHyperparams:
+    # Using the "default" hyperparameters of the FT-Transformer paper: https://arxiv.org/pdf/2106.11959
+    early_stopping_rounds: int = 50
+    iterations: int = 2000
+    od_pval: float = 0.001
+    random_state: int = SEED
+    thread_count = CPU_CORES
+
+class CatBoost(TabularModel):
+
+    MODEL_NAME = "CatBoost"
+    SHORT_NAME = "cat"
+    USE_VAL_SPLIT = True
+    USE_MEDIAN_FILLING = False
+    USE_CATEGORICAL_ENCODING = False
+    USE_TEXT_EMBEDDINGS = True
+    USE_TARGET_ENCODER = True
+    NEEDS_GPU = False
+
+    def initialize_model(self):
+        from catboost import CatBoostRegressor, CatBoostClassifier
+        model_cls = CatBoostClassifier if self.is_cls else CatBoostRegressor
+        params = vars(CatBoostDefaultHyperparams())
+        return model_cls(**params)
+
+    def fit_model(self, x_train: DataFrame, y_train: Series, x_val: DataFrame, y_val: Series):
+        cat_features = [i for i, c in enumerate(x_train.columns) if c in self.categorical_features]
+        logging_level = "Verbose" if self.verbose else "Silent"
+        self.model_.fit(x_train, y_train, eval_set=(x_val, y_val), use_best_model=True, cat_features=cat_features,
+                        logging_level=logging_level)

multabench/baselines/contexttab.py

@@ -0,0 +1,30 @@
+import os
+
+from pandas import DataFrame, Series
+
+from multabench.baselines.abstract_model import TabularModel
+from multabench.constants import HF_TOKEN
+
+
+class ConTextTab(TabularModel):
+    MODEL_NAME = "ConTextTab"
+    SHORT_NAME = "ctx"
+    USE_VAL_SPLIT = False
+    USE_MEDIAN_FILLING = False
+    USE_CATEGORICAL_ENCODING = False
+    # TODO: Should we use text embeddings for ConTextTab?
+    USE_TEXT_EMBEDDINGS = False
+    USE_TARGET_ENCODER = False
+
+    def initialize_model(self):
+        from sap_rpt_oss import SAP_RPT_OSS_Classifier, SAP_RPT_OSS_Regressor
+        os.environ["HF_TOKEN"] = HF_TOKEN or ""
+        model_cls = SAP_RPT_OSS_Classifier if self.is_cls else SAP_RPT_OSS_Regressor
+        # Initialize the regressor, 8k context and 8-fold bagging gives best performance, reduce if running out of memory
+        model = model_cls(max_context_size=8192, bagging=8)
+        return model
+
+    def fit_model(self, x_train: DataFrame, y_train: Series, x_val: DataFrame, y_val: Series):
+        # TODO: ConTextTab natively supports dates, perhaps we need to process them differently
+        self.model_.fit(x_train, y_train)
+

multabench/baselines/kumo_tabular.py

@@ -0,0 +1,103 @@
+import numpy as np
+import torch
+from pandas import DataFrame, Series
+
+from multabench.baselines.abstract_model import TabularModel
+from multabench.constants import SEED
+
+TARGET = "__target__"
+
+
+class KumoEstimator:
+    """Kumo Tabular on pandas, as TabArena and NVIDIA's adapter run it: fit caches the context, predict reads each
+    query row from it under float16 autocast."""
+
+    def __init__(self, task: str, size: str, num_estimators: int, n_classes: int, device: torch.device):
+        self.task = task
+        self.size = size
+        self.num_estimators = num_estimators
+        self.n_classes = n_classes
+        self.device = device
+
+    def _inference(self):
+        return torch.amp.autocast(self.device.type, dtype=torch.float16, enabled=self.device.type == "cuda")
+
+    def fit(self, x: DataFrame, y: Series):
+        import sdm
+        self.model = sdm.models.KumoTabular(task=self.task, size=self.size, device=self.device)
+        self.stypes = sdm.infer_stypes(x, _low_cardinality="infer")
+        x_context = sdm.TableTensor.from_pandas(df=x, stypes=self.stypes, device=self.device)
+        target_stype = "numerical" if self.task == "regression" else "categorical"
+        y_context = sdm.TableTensor.from_pandas(df=y.rename(TARGET).to_frame(), stypes={TARGET: target_stype},
+                                                device=self.device)
+        rows, cols = x_context.shape[-2:]
+        # NVIDIA's adapter fits one ensemble member at a time above these sizes, to bound GPU memory.
+        batch_size = 1 if rows > 2_000 or rows * cols >= 50_000 else self.num_estimators
+        generator = torch.Generator(self.device).manual_seed(SEED)
+        with torch.inference_mode(), self._inference():
+            self.model.fit(x=x_context, y=y_context, num_estimators=self.num_estimators,
+                           estimator_batch_size=batch_size, generator=generator)
+        return self
+
+    def _forward(self, x: DataFrame):
+        import sdm
+        x_query = sdm.TableTensor.from_pandas(df=x, stypes=self.stypes, device=self.device)
+        with torch.inference_mode(), self._inference():
+            return self.model.predict(x_query)
+
+    def predict(self, x: DataFrame) -> np.ndarray:
+        """The mean of the predicted quantiles."""
+        return self._forward(x).numerical.float().mean(dim=-1).cpu().numpy()
+
+    def predict_proba(self, x: DataFrame) -> np.ndarray:
+        import sdm
+        out = self._forward(x)
+        probs = np.zeros((len(x), self.n_classes), dtype=np.float32)
+        probs[:, [int(label) for label in out.columns[sdm.Stype.numerical]]] = out.numerical.float().cpu().numpy()
+        return probs
+
+
+class Kumo(TabularModel):
+    """NVIDIA's Kumo Tabular; its sizes differ only in the checkpoint (SIZE) and TabArena's NUM_ESTIMATORS."""
+
+    USE_VAL_SPLIT = False
+    USE_MEDIAN_FILLING = False
+    USE_CATEGORICAL_ENCODING = False
+    USE_TEXT_EMBEDDINGS = True
+    USE_TARGET_ENCODER = True
+
+    SIZE: str
+    NUM_ESTIMATORS: int
+
+    def initialize_model(self) -> None:
+        return None
+
+    def fit_model(self, x_train: DataFrame, y_train: Series, x_val: DataFrame, y_val: Series):
+        task = "classification" if self.is_cls else "regression"
+        self.model_ = KumoEstimator(task=task, size=self.SIZE, num_estimators=self.NUM_ESTIMATORS,
+                                    n_classes=self.d_output, device=self.device)
+        self.model_.fit(x_train, y_train)
+
+
+class KumoTabular(Kumo):
+
+    MODEL_NAME = "Kumo-Tabular"
+    SHORT_NAME = "kumo"
+    SIZE = "large"
+    NUM_ESTIMATORS = 16
+
+
+class KumoTabularMedium(Kumo):
+
+    MODEL_NAME = "Kumo-Tabular-Medium"
+    SHORT_NAME = "kumomedium"
+    SIZE = "medium"
+    NUM_ESTIMATORS = 8
+
+
+class KumoTabularSmall(Kumo):
+
+    MODEL_NAME = "Kumo-Tabular-Small"
+    SHORT_NAME = "kumosmall"
+    SIZE = "small"
+    NUM_ESTIMATORS = 8

multabench/baselines/lgbm.py

@@ -0,0 +1,28 @@
+from pandas import DataFrame, Series
+
+from multabench.constants import SEED
+from multabench.baselines.abstract_model import TabularModel
+from multabench.utils.devices import CPU_CORES
+
+
+class LightGBM(TabularModel):
+
+    MODEL_NAME = "LightGBM"
+    SHORT_NAME = "light"
+    USE_VAL_SPLIT = True
+    USE_MEDIAN_FILLING = False
+    USE_CATEGORICAL_ENCODING = True
+    USE_TEXT_EMBEDDINGS = True
+    USE_TARGET_ENCODER = True
+    NEEDS_GPU = False
+
+    def initialize_model(self):
+        from lightgbm import LGBMClassifier, LGBMRegressor
+        model_cls = LGBMClassifier if self.is_cls else LGBMRegressor
+        params = {"verbose": -1, "random_state": SEED, "n_jobs": CPU_CORES}
+        return model_cls(**params)
+
+    def fit_model(self, x_train: DataFrame, y_train: Series, x_val: DataFrame, y_val: Series):
+        # Use column names so we only mark truly categorical columns; indices can point to PCA cols after transforms.
+        cat_cols = [c for c in x_train.columns if c in self.categorical_features]
+        self.model_.fit(x_train, y_train, eval_set=[(x_val, y_val)], categorical_feature=cat_cols)