From 2b3007c4df8de47e4ee6c4aed9c3c50e3acd9206 Mon Sep 17 00:00:00 2001 From: Georg Grab Date: Mon, 27 Apr 2026 12:56:34 +0000 Subject: [PATCH] chore: tabpfn_client + tabpfn_plus baselines --- tabstar_paper/baselines/_sanity_folds.py | 111 ++++++++++++++++++ .../baselines/_sanity_tabpfn_client.py | 74 ++++++++++++ .../baselines/_sanity_tabpfn_plus.py | 63 ++++++++++ .../baselines/tabpfn_client_baseline.py | 70 +++++++++++ .../baselines/tabpfn_plus_baseline.py | 109 +++++++++++++++++ 5 files changed, 427 insertions(+) create mode 100644 tabstar_paper/baselines/_sanity_folds.py create mode 100644 tabstar_paper/baselines/_sanity_tabpfn_client.py create mode 100644 tabstar_paper/baselines/_sanity_tabpfn_plus.py create mode 100644 tabstar_paper/baselines/tabpfn_client_baseline.py create mode 100644 tabstar_paper/baselines/tabpfn_plus_baseline.py diff --git a/tabstar_paper/baselines/_sanity_folds.py b/tabstar_paper/baselines/_sanity_folds.py new file mode 100644 index 0000000..fe0bcbe --- /dev/null +++ b/tabstar_paper/baselines/_sanity_folds.py @@ -0,0 +1,111 @@ +"""Sanity-fit either the local make_tabpfn_plus baseline or the tabpfn-client +v3 API baseline over multiple TabSTAR folds and print a comparison table. + +Uses TabSTAR's official split logic (subsample_dataset + split_to_test) with +``DOWNSTREAM_EXAMPLES = 10_000``, RNG = ``SEED + fold``, so the numbers are +comparable across runs and to the rest of the benchmark. + +Run: + # local v3-demo-hf checkpoints + .TabSTAR-env/bin/python tabstar_paper/baselines/_sanity_folds.py local + + # tabpfn-client v3 API (prod) + TABPFN_TOKEN= \\ + .TabSTAR-env/bin/python tabstar_paper/baselines/_sanity_folds.py client + + # both + TABPFN_TOKEN= \\ + .TabSTAR-env/bin/python tabstar_paper/baselines/_sanity_folds.py both +""" +from __future__ import annotations + +import argparse +import os +import time + +import torch + +from tabstar.datasets.all_datasets import OpenMLDatasetID +from tabstar.preprocessing.splits import split_to_test +from tabstar_paper.benchmarks.evaluate import DOWNSTREAM_EXAMPLES +from tabstar_paper.datasets.downloading import download_dataset +from tabstar_paper.preprocessing.sampling import subsample_dataset + + +# Defaults applied before tabpfn_client picks them up. Overridable on the caller side. +os.environ.setdefault("TABPFN_CLIENT_API_URL", "https://api.priorlabs.ai") +os.environ.setdefault("TABPFN_AUTH_API_URL", "https://api.priorlabs.ai") +os.environ.setdefault("TABPFN_CLIENT_FORCE_REUPLOAD", "true") +os.environ.setdefault("TABPFN_CLIENT_FORCE_REFIT", "true") +os.environ.setdefault("TABPFN_CLIENT_TIMEOUT", "3600") + + +def _run_one(model_cls, dataset_id, fold, device, *, label): + """Reproduces tabstar_paper.benchmarks.evaluate.evaluate_on_dataset's split.""" + ds = download_dataset(dataset_id=dataset_id) + is_cls = ds.is_cls + x, y = subsample_dataset( + x=ds.x, y=ds.y, is_cls=is_cls, + train_examples=DOWNSTREAM_EXAMPLES, fold=fold, + ) + x_train, x_test, y_train, y_test = split_to_test( + x=x, y=y, is_cls=is_cls, + fold=fold, train_examples=DOWNSTREAM_EXAMPLES, + ) + model = model_cls(problem_type=ds.task_type, device=device, verbose=False) + t0 = time.time() + model.fit(x_train, y_train) + fit_s = time.time() - t0 + t1 = time.time() + score = model.score(x_test, y_test) + score_s = time.time() - t1 + print(f" fold {fold} [{label}] train={len(x_train):>5} test={len(x_test):>5} " + f"score={score:.4f} fit={fit_s:5.1f}s predict={score_s:5.1f}s") + return {"fold": fold, "label": label, "score": score, "fit_s": fit_s, "score_s": score_s, + "n_train": len(x_train), "n_test": len(x_test)} + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("which", choices=["local", "client", "both"]) + parser.add_argument("--dataset", default="BIN_PROFESSIONAL_FAKE_JOB_POSTING") + parser.add_argument("--folds", type=int, nargs="+", default=[0, 1, 2, 3, 4]) + args = parser.parse_args() + + dataset_id = OpenMLDatasetID[args.dataset] + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + variants: list[tuple[str, type]] = [] + if args.which in {"local", "both"}: + from tabstar_paper.baselines.tabpfn_plus_baseline import TabPFNPlus + variants.append(("local v3-demo-hf", TabPFNPlus)) + if args.which in {"client", "both"}: + if "TABPFN_TOKEN" not in os.environ: + raise SystemExit("TABPFN_TOKEN env var required for client variant") + from tabstar_paper.baselines.tabpfn_client_baseline import TabPFNClient + variants.append(("client v3 API", TabPFNClient)) + + print(f"📚 dataset={args.dataset} folds={args.folds} device={device.type}") + rows: list[dict] = [] + for label, cls in variants: + print(f"\n=== {label} ({cls.__name__}) ===") + for fold in args.folds: + rows.append(_run_one(cls, dataset_id, fold, device, label=label)) + + print("\n=== summary ===") + print(f"{'variant':<20} {'fold':>4} {'score':>8} {'fit_s':>8} {'pred_s':>8}") + for r in rows: + print(f"{r['label']:<20} {r['fold']:>4} {r['score']:>8.4f} {r['fit_s']:>8.1f} {r['score_s']:>8.1f}") + + by_label: dict[str, list[float]] = {} + for r in rows: + by_label.setdefault(r["label"], []).append(r["score"]) + print() + for label, scores in by_label.items(): + mean = sum(scores) / len(scores) + spread = max(scores) - min(scores) + print(f" {label:<20} mean={mean:.4f} min={min(scores):.4f} max={max(scores):.4f} spread={spread:.4f}") + + +if __name__ == "__main__": + main() diff --git a/tabstar_paper/baselines/_sanity_tabpfn_client.py b/tabstar_paper/baselines/_sanity_tabpfn_client.py new file mode 100644 index 0000000..ebc4cd1 --- /dev/null +++ b/tabstar_paper/baselines/_sanity_tabpfn_client.py @@ -0,0 +1,74 @@ +"""One-shot sanity check: fit the TabPFN-v3 API baseline on a TabSTAR dataset. + +Defaults to `BIN_PROFESSIONAL_FAKE_JOB_POSTING` against prod (api.priorlabs.ai), +subsampled to fit TabPFN's 10k-row pretraining limit. + +Run: + TABPFN_TOKEN= \\ + .TabSTAR-env/bin/python tabstar_paper/baselines/_sanity_tabpfn_client.py +""" +from __future__ import annotations + +import argparse +import os +import time + +import torch +from sklearn.model_selection import train_test_split + +from tabstar.constants import SEED +from tabstar.datasets.all_datasets import OpenMLDatasetID +from tabstar_paper.datasets.downloading import download_dataset + + +os.environ.setdefault("TABPFN_CLIENT_API_URL", "https://api.priorlabs.ai") +os.environ.setdefault("TABPFN_AUTH_API_URL", "https://api.priorlabs.ai") +os.environ.setdefault("TABPFN_CLIENT_FORCE_REUPLOAD", "true") +os.environ.setdefault("TABPFN_CLIENT_FORCE_REFIT", "true") +os.environ.setdefault("TABPFN_CLIENT_TIMEOUT", "3600") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", default="BIN_PROFESSIONAL_FAKE_JOB_POSTING") + parser.add_argument("--max-train", type=int, default=8000) + parser.add_argument("--test-size", type=int, default=2000) + parser.add_argument("--enhanced", action="store_true") + args = parser.parse_args() + + if "TABPFN_TOKEN" not in os.environ: + raise SystemExit("TABPFN_TOKEN env var required (user JWT or IAM identity token)") + + # Import after env defaults so the client picks up the URLs. + from tabstar_paper.baselines.tabpfn_client_baseline import TabPFNClient, TabPFNClientEnhanced + + ds = download_dataset(OpenMLDatasetID[args.dataset]) + x, y = ds.x, ds.y + print(f"📚 dataset={args.dataset} shape={x.shape} task={ds.task_type.name} target={y.name}") + + stratify = y if ds.is_cls else None + x_train, x_test, y_train, y_test = train_test_split( + x, y, train_size=args.max_train, test_size=args.test_size, + random_state=SEED, stratify=stratify, + ) + print(f"split: train={len(x_train)} test={len(x_test)} via {os.environ['TABPFN_CLIENT_API_URL']}") + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + cls = TabPFNClientEnhanced if args.enhanced else TabPFNClient + model = cls(problem_type=ds.task_type, device=device, verbose=True) + + t0 = time.time() + model.fit(x_train, y_train) + fit_s = time.time() - t0 + + t1 = time.time() + score = model.score(x_test, y_test) + score_s = time.time() - t1 + + print(f"score: {score:.4f} (model={cls.MODEL_NAME})") + print(f"fit: {fit_s:.1f}s") + print(f"score: {score_s:.1f}s") + + +if __name__ == "__main__": + main() diff --git a/tabstar_paper/baselines/_sanity_tabpfn_plus.py b/tabstar_paper/baselines/_sanity_tabpfn_plus.py new file mode 100644 index 0000000..f0fa49f --- /dev/null +++ b/tabstar_paper/baselines/_sanity_tabpfn_plus.py @@ -0,0 +1,63 @@ +"""One-shot sanity check: fit the TabPFN-Plus baseline on a TabSTAR dataset. + +Default dataset is `BIN_PROFESSIONAL_FAKE_JOB_POSTING` (12.7k rows, text-heavy +binary classification), subsampled to fit TabPFN's 10k-row pretraining limit. +The bundled imdb CSV is the fallback. + +Run: + .TabSTAR-env/bin/python tabstar_paper/baselines/_sanity_tabpfn_plus.py +""" +from __future__ import annotations + +import argparse +import time + +import torch +from sklearn.model_selection import train_test_split + +from tabstar.constants import SEED +from tabstar.datasets.all_datasets import OpenMLDatasetID +from tabstar_paper.baselines.tabpfn_plus_baseline import TabPFNPlus +from tabstar_paper.datasets.downloading import download_dataset + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", default="BIN_PROFESSIONAL_FAKE_JOB_POSTING") + parser.add_argument("--max-train", type=int, default=8000) + parser.add_argument("--test-size", type=int, default=2000) + args = parser.parse_args() + + ds = download_dataset(OpenMLDatasetID[args.dataset]) + x, y = ds.x, ds.y + print(f"📚 dataset={args.dataset} shape={x.shape} task={ds.task_type.name} " + f"target={y.name} pos_rate={y.eq(y.value_counts().idxmax()).mean():.3f}") + + # Cap train at 8k so we stay under TabPFN's 10k pretraining limit even + # after the wrapper's val split. Test size capped too — predict_proba + # on 12k rows takes a while, and 2k is plenty for a single AUC. + stratify = y if ds.is_cls else None + x_train, x_test, y_train, y_test = train_test_split( + x, y, train_size=args.max_train, test_size=args.test_size, + random_state=SEED, stratify=stratify, + ) + print(f"split: train={len(x_train)} test={len(x_test)}") + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = TabPFNPlus(problem_type=ds.task_type, device=device, verbose=True) + + t0 = time.time() + model.fit(x_train, y_train) + fit_s = time.time() - t0 + + t1 = time.time() + score = model.score(x_test, y_test) + score_s = time.time() - t1 + + print(f"score: {score:.4f} (metric depends on task)") + print(f"fit: {fit_s:.1f}s") + print(f"score: {score_s:.1f}s") + + +if __name__ == "__main__": + main() diff --git a/tabstar_paper/baselines/tabpfn_client_baseline.py b/tabstar_paper/baselines/tabpfn_client_baseline.py new file mode 100644 index 0000000..9cdd0ba --- /dev/null +++ b/tabstar_paper/baselines/tabpfn_client_baseline.py @@ -0,0 +1,70 @@ +"""TabPFN-Client (v3 API) baseline. + +Mirror of `tabpfnv2.py` and `tabpfn_plus_baseline.py`, but routes through the +hosted `tabpfn-client` SDK so the benchmark exercises the production API. +Defaults to v3 via `TabPFNClassifier.create_default_for_version(ModelVersion.V3)`, +matching the tabarena medium-parity setup. + +Required env (set on the caller side): + TABPFN_CLIENT_API_URL e.g. https://api.priorlabs.ai + TABPFN_AUTH_API_URL same as above (or the gapi behind IAM) + TABPFN_TOKEN user JWT (or impersonation identity token) +""" + +from pandas import DataFrame, Series + +from tabstar.constants import SEED +from tabstar_paper.baselines.abstract_model import TabularModel + + +MAX_SAMPLES = 10_000 +MAX_FEATURES = 500 + + +class TabPFNClient(TabularModel): + + MODEL_NAME = "TabPFN-v3 (API) ☁️" + SHORT_NAME = "tabpfn_client" + USE_VAL_SPLIT = False + USE_MEDIAN_FILLING = False + USE_CATEGORICAL_ENCODING = False + # The hosted v3 API runs its own preprocessing/text systems server-side, + # so don't double-encode here. + USE_TEXT_EMBEDDINGS = False + + ENHANCED_FIT_MODE = False + ENHANCED_FIT_TIME_LIMIT_S: float | None = None + N_ESTIMATORS = 4 + + def initialize_model(self) -> None: + return None + + def fit_model(self, x_train: DataFrame, y_train: Series, x_val: DataFrame, y_val: Series): + if len(x_train) > MAX_SAMPLES: + raise RuntimeError(f"TabPFN is not designed to handle datasets larger than {MAX_SAMPLES} samples. ") + if len(x_train.columns) > MAX_FEATURES: + print(f"🚨 Warning: {len(x_train.columns)} features detected, ignoring pretraining limits.") + from tabpfn_client import TabPFNClassifier, TabPFNRegressor # noqa: PLC0415 + from tabpfn_client.constants import ModelVersion # noqa: PLC0415 + + cls = TabPFNClassifier if self.is_cls else TabPFNRegressor + kwargs: dict = dict( + n_estimators=self.N_ESTIMATORS, + random_state=SEED, + ) + if self.ENHANCED_FIT_MODE: + kwargs["enhanced_fit_mode"] = True + if self.ENHANCED_FIT_TIME_LIMIT_S is not None: + kwargs["enhanced_fit_mode_time_limit_s"] = self.ENHANCED_FIT_TIME_LIMIT_S + self.model_ = cls.create_default_for_version(ModelVersion.V3, **kwargs) + self.model_.fit(x_train, y_train) + + +class TabPFNClientEnhanced(TabPFNClient): + """Same as `TabPFNClient` but flips `enhanced_fit_mode=True`.""" + + MODEL_NAME = "TabPFN-v3 (API, enhanced) ☁️" + SHORT_NAME = "tabpfn_client_enhanced" + + ENHANCED_FIT_MODE = True + ENHANCED_FIT_TIME_LIMIT_S = 1200 diff --git a/tabstar_paper/baselines/tabpfn_plus_baseline.py b/tabstar_paper/baselines/tabpfn_plus_baseline.py new file mode 100644 index 0000000..412e86e --- /dev/null +++ b/tabstar_paper/baselines/tabpfn_plus_baseline.py @@ -0,0 +1,109 @@ +"""TabPFN-Plus baseline — `make_tabpfn_plus(...)` with skrub + text systems on, +v3-demo-hf checkpoints. + +Mirror of `tabpfnv2.py`, but routes through the tabpfn_plus wrapper so the +benchmark exercises the post-update preprocessing/text path with **v3 +weights** (the bare `TabPFNClassifier()` defaults to `model_path="auto"`, +which resolves to the v2 fallback URL — wrong for this comparison). + +V3 weights live under ``$TABPFN_V3_MODELS_DIR`` (default ``~/models``) in the +``v3-demo-hf/`` directory and ship as separate binary/multiclass classifier +checkpoints + a regressor — same set tabarena's local-enhanced model uses. + +`use_enhanced` is off in `TabPFNPlus` to match `tabpfnv2` cost; flip via the +`TabPFNPlusEnhanced` subclass. +""" + +import os +from pathlib import Path + +from pandas import DataFrame, Series + +from tabstar.constants import SEED +from tabstar_paper.baselines.abstract_model import TabularModel +from tabpfn import TabPFNClassifier, TabPFNRegressor +from tabpfn_plus import make_tabpfn_plus + + +MAX_SAMPLES = 10_000 +MAX_FEATURES = 500 + + +_V3_MODELS_DIR = Path(os.path.expanduser(os.environ.get("TABPFN_V3_MODELS_DIR", "~/models"))) / "v3-demo-hf" +_V3_BINARY_CKPT = _V3_MODELS_DIR / "20260417-classifier-binary-miizr9pp.ckpt" +_V3_MULTICLASS_CKPT = _V3_MODELS_DIR / "20260417-classifier-multiclass-zcj01hic.ckpt" +_V3_REGRESSOR_CKPT = _V3_MODELS_DIR / "20260417-regressor-2dgwn82n.ckpt" + + +def _resolve_v3_ckpt(*, is_cls: bool, n_classes: int) -> Path: + """Pick the right v3-demo-hf checkpoint for this task. + + Mirrors tabarena's `TabPFNv3OSSDemoHFModel`: dedicated binary head when + `n_classes == 2`, multiclass head otherwise; regressor for non-classification. + """ + if not is_cls: + ckpt = _V3_REGRESSOR_CKPT + elif n_classes == 2: + ckpt = _V3_BINARY_CKPT + else: + ckpt = _V3_MULTICLASS_CKPT + if not ckpt.exists(): + raise FileNotFoundError( + f"v3-demo-hf checkpoint missing: {ckpt}. Set TABPFN_V3_MODELS_DIR or " + f"download the v3-demo-hf weights to ~/models/v3-demo-hf/." + ) + return ckpt + + +class TabPFNPlus(TabularModel): + + MODEL_NAME = "TabPFN-Plus (v3) 🚀" + SHORT_NAME = "tabpfn_plus" + USE_VAL_SPLIT = False + USE_MEDIAN_FILLING = False + USE_CATEGORICAL_ENCODING = False + # tabpfn_plus does its own text handling via `use_text=True`, so we skip + # the abstract model's PCA-on-MiniLM path and let the wrapper own it. + USE_TEXT_EMBEDDINGS = False + + USE_ENHANCED = False + ENHANCED_TIME_BUDGET_S: float | None = None + + def initialize_model(self) -> None: + return None + + def fit_model(self, x_train: DataFrame, y_train: Series, x_val: DataFrame, y_val: Series): + if len(x_train) > MAX_SAMPLES: + raise RuntimeError(f"TabPFN is not designed to handle datasets larger than {MAX_SAMPLES} samples. ") + if len(x_train.columns) > MAX_FEATURES: + print(f"🚨 Warning: {len(x_train.columns)} features detected, ignoring pretraining limits.") + cat_idx = [i for i, c in enumerate(x_train.columns) if c in self.categorical_features] + ckpt = _resolve_v3_ckpt(is_cls=self.is_cls, n_classes=self.d_output) + print(f"📂 v3 checkpoint: {ckpt.name}") + base_cls = TabPFNRegressor if not self.is_cls else TabPFNClassifier + base = base_cls( + random_state=SEED, + categorical_features_indices=cat_idx, + ignore_pretraining_limits=True, + model_path=str(ckpt), + ) + kwargs: dict = dict( + base_estimator=base, + use_preprocessing=True, + use_text=True, + use_enhanced=self.USE_ENHANCED, + ) + if self.USE_ENHANCED and self.ENHANCED_TIME_BUDGET_S is not None: + kwargs["enhanced_time_budget"] = self.ENHANCED_TIME_BUDGET_S + self.model_ = make_tabpfn_plus(**kwargs) + self.model_.fit(x_train, y_train) + + +class TabPFNPlusEnhanced(TabPFNPlus): + """Same as `TabPFNPlus` but flips the enhanced wrapper on.""" + + MODEL_NAME = "TabPFN-Plus (v3, enhanced) 🚀" + SHORT_NAME = "tabpfn_plus_enhanced" + + USE_ENHANCED = True + ENHANCED_TIME_BUDGET_S = 600