""" Uniform loading of every deployable model, for inference ======================================================== The three families are trained and stored differently and predict differently. A conventional classifier is a pickled fitted estimator. A deep network is a state dictionary that must be poured back into a rebuilt architecture. A foundation model is an in-context learner with no fitted parameters at all: the checkpoint is fixed and pre-trained, and what the study "fitted" is the corpus it conditions on, so the artefact stores that context rather than weights. Serving them through one interface is what lets the application offer all three in a single menu, and lets the optimiser score candidates without knowing which kind it is holding. Batching is the reason this module exposes `predict_proba` on a matrix rather than a row. The cost of a foundation model is dominated by the pass over its 2,646 context rows, which is paid once per call however many candidates are scored in it, so a batch of 32 costs what a batch of 1 costs. Measured on this corpus: one row 32 rows per candidate in a batch Bagged Trees 83 ms 80 ms 2.5 ms TabICL (GPU) 997 ms 1 040 ms 33 ms TabICL (CPU) 17 000 ms 16 900 ms 530 ms Scoring candidates singly, as the previous release did, therefore made the foundation models appear unusable when they are not: it charged the context pass once per candidate instead of once per batch. """ from __future__ import annotations import json import os from dataclasses import dataclass, field from pathlib import Path from typing import Callable import joblib import numpy as np from mlate import config as cfg # The exact TabPFN checkpoint the study used, so a served prediction matches a # reported one. See 05_deep_learning/foundation_models.py for why this is a # filename rather than a version string. TABPFN_CHECKPOINT = "tabpfn-v2.6-classifier-v2.6_default.ckpt" N_ESTIMATORS = 8 FAMILY_DIRS = {"ml": ("classifiers", True), "dl": ("deep", True), "foundation": ("foundation", False)} # What a user is choosing between, in time rather than in architecture. Taken # from the measurements in the module docstring and stated per 150-trial search # over both targets, batched, on a CPU host. The application shows the real # elapsed time once a search has run; this is only what it can say beforehand. COST_HINT = {"ml": "fast, a few seconds", "dl": "fast, a few seconds", "foundation": "slow on CPU, roughly 3 minutes"} def device() -> str: """ CUDA when it is genuinely available, CPU otherwise. Resolved once here rather than left to each library's own default, because the two disagree: TabPFN refuses a CPU run with more than 1,000 context rows unless told to proceed, while TabICL simply runs and is twenty times slower than on a GPU. A single answer keeps the menu's cost estimate honest for whichever host the application is deployed to. """ try: import torch return "cuda" if torch.cuda.is_available() else "cpu" except Exception: return "cpu" # Keys a TabPFN token may arrive under. Read from the environment only - a # Hugging Face Space supplies it as a secret, a local checkout as an # uncommitted .env - and never written to a source file or logged. A missing # token is not an error: the package then uses locally downloaded weights. TABPFN_TOKEN_KEYS = ("TABPFN_TOKEN", "TABPFN_ACCESS_TOKEN", "PRIORLABS_TOKEN", "PRIOR_LABS_TOKEN") def _apply_tabpfn_token() -> bool: for key in TABPFN_TOKEN_KEYS: value = os.environ.get(key) if value: os.environ["TABPFN_TOKEN"] = value.strip() return True return False @dataclass class Servable: """One model, loaded and ready to score a batch of candidate rows.""" name: str task: str family: str path: Path classes: np.ndarray weighted_f1: float = 0.0 _predict: Callable[[np.ndarray], np.ndarray] | None = field( default=None, repr=False) @property def cost_hint(self) -> str: return COST_HINT.get(self.family, "") def predict_proba(self, X: np.ndarray) -> np.ndarray: """Class probabilities for every row of X, in one call.""" return np.asarray(self._predict(np.asarray(X, dtype=float)), dtype=float) # ── per-family loading ─────────────────────────────────────────────────────── def _load_sklearn(path: Path) -> Servable: bundle = joblib.load(path) est = bundle["estimator"] classes = np.asarray(bundle["classes"], dtype=float) def predict(X): if hasattr(est, "predict_proba"): return est.predict_proba(X) # Five conventional classifiers expose no probabilities. A one-hot on # the predicted label keeps the interface uniform; the expected value # it produces is then a step function rather than a smooth one, which # is a property of the model and is stated where it matters. idx = np.asarray(est.predict(X)).astype(int) out = np.zeros((len(idx), len(classes))) out[np.arange(len(idx)), idx] = 1.0 return out return Servable(bundle["model"], bundle["task"], "ml", path, classes, _predict=predict) def _load_deep(path: Path) -> Servable: import torch from mlate import deep as deep_mod bundle = torch.load(path, map_location="cpu", weights_only=False) classes = np.asarray(bundle["classes"], dtype=float) net = deep_mod.build(bundle["architecture"], bundle["input_dim"], len(classes), bundle["params"]) net.load_state_dict(bundle["state_dict"]) net.eval() def predict(X): with torch.no_grad(): logits = net(torch.tensor(X, dtype=torch.float32)) return torch.softmax(logits, dim=1).numpy() return Servable(bundle["architecture"], bundle["task"], "dl", path, classes, _predict=predict) def _load_foundation(path: Path) -> Servable: """ Rebuild the in-context learner and re-supply the corpus it was given. `fit` here is not training. It hands the model the stored context and returns; the checkpoint is unchanged. The call is deferred until the first prediction so that opening the menu does not pay for every model in it. """ bundle = joblib.load(path) classes = np.asarray(bundle["classes"], dtype=float) name = bundle["model"] state: dict = {} def fitted(): if "clf" in state: return state["clf"] dev = device() if name.startswith("TabPFN"): from tabpfn import TabPFNClassifier kwargs = dict(device=dev, n_estimators=bundle["n_estimators"], # 153 features exceeds the pre-training shape, and # 2,646 context rows exceed what TabPFN will run on # CPU by default. Both limits are advisory and both # are lifted by this flag; the second is why a CPU # host is slow rather than why it refuses. ignore_pretraining_limits=True, model_path=bundle.get("checkpoint") or TABPFN_CHECKPOINT, random_state=cfg.RANDOM_STATE) if name.endswith("(thinking)"): from tabpfn.inference_tuning import ClassifierTuningConfig kwargs["eval_metric"] = "f1" kwargs["tuning_config"] = ClassifierTuningConfig( calibrate_temperature=True, tune_decision_thresholds=True) os.environ.setdefault("TABPFN_ALLOW_CPU_LARGE_DATASET", "1") os.environ.setdefault("TABPFN_NO_BROWSER", "1") _apply_tabpfn_token() clf = TabPFNClassifier(**kwargs) elif name == "TabICL": from tabicl import TabICLClassifier clf = TabICLClassifier(device=dev, n_estimators=bundle["n_estimators"], allow_auto_download=True, random_state=cfg.RANDOM_STATE) else: raise KeyError(name) clf.fit(bundle["context_X"], bundle["context_y"]) state["clf"] = clf return clf return Servable(name, bundle["task"], "foundation", path, classes, _predict=lambda X: fitted().predict_proba(X)) LOADERS = {"ml": _load_sklearn, "dl": _load_deep, "foundation": _load_foundation} def load(path: Path, family: str) -> Servable: return LOADERS[family](Path(path)) # ── discovery ──────────────────────────────────────────────────────────────── def manifest_scores(models_root: Path, task: str, protocol: str) -> dict[str, float]: """ Benchmarked weighted F1 per model, for ranking the menu. Foundation models carry no split protocol - they are not refitted per split, so the manifest records them as "n/a" - and are scored against the same held-out test partition as the rest. """ p = models_root / "deployment_manifest.json" if not p.exists(): return {} out = {} for r in json.loads(p.read_text(encoding="utf-8"))["models"]: if r.get("task") != task: continue if r.get("protocol") not in (protocol, "n/a"): continue out[r["model"]] = (r.get("metrics") or {}).get("weighted_f1", 0.0) return out def discover(models_root: Path, task: str, protocol: str = "random") -> list[Servable]: """ Every artefact the application can offer for one target, best first. Nothing is loaded here beyond the manifest and each artefact's identity; the estimators themselves are opened when a model is chosen. Ranking is by benchmarked weighted F1 so that the head of the menu and the head of the performance table are the same model, which they were not while the menu saw only the conventional classifiers. """ scores = manifest_scores(models_root, task, protocol) found: list[Servable] = [] for family, (sub, per_protocol) in FAMILY_DIRS.items(): d = models_root / sub / (protocol if per_protocol else "") if not d.exists(): continue for path in sorted(d.glob(f"{task}__*")): if path.suffix not in (".pkl", ".pt"): continue found.append(Stub(path, family, task, scores)) found.sort(key=lambda s: -s.weighted_f1) return found @dataclass class Stub: """ A menu entry that has not been opened yet. Exists so that listing the models costs a filename parse rather than unpickling ninety estimators and rebuilding six networks. `open()` returns the real thing. """ path: Path family: str task: str _scores: dict = field(default_factory=dict, repr=False) name: str = "" weighted_f1: float = 0.0 def __post_init__(self): if not self.name: self.name = _display_name(self.path, self.task, self._scores) self.weighted_f1 = self._scores.get(self.name, 0.0) @property def cost_hint(self) -> str: return COST_HINT.get(self.family, "") def open(self) -> Servable: s = load(self.path, self.family) s.weighted_f1 = self.weighted_f1 return s def _display_name(path: Path, task: str, scores: dict) -> str: """ The model's benchmark name, recovered from its filename. File names are slugs - "cell_response__k_nn_distance_weighted.pkl" - while the manifest and the performance table use the printed name, "K-NN (distance weighted)". Matching on the slug rather than opening the artefact keeps discovery cheap; an unmatched slug falls back to a readable form of itself rather than failing, so a newly added model still appears. """ stem = path.stem[len(task) + 2:] for known in scores: if _slug(known) == stem: return known return stem.replace("_", " ").title() def _slug(name: str) -> str: out = [] for ch in name.lower(): out.append(ch if ch.isalnum() else "_") slug = "".join(out) while "__" in slug: slug = slug.replace("__", "_") return slug.strip("_")