chore: tabpfn_client + tabpfn_plus baselines

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=<jwt> \\
+    .TabSTAR-env/bin/python tabstar_paper/baselines/_sanity_folds.py client
+
+    # both
+    TABPFN_TOKEN=<jwt> \\
+    .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()

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=<jwt> \\
+    .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()

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()

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