From b75b89eaa5f0b2b0bbb8794512e29aeab516a911 Mon Sep 17 00:00:00 2001 From: Dominik Safaric Date: Sat, 18 Jul 2026 21:24:51 +0200 Subject: [PATCH 1/5] Include options for KV cache and model_path --- src/tabpfn_client/endpoints/__init__.py | 11 + src/tabpfn_client/endpoints/estimator.py | 251 +++++++++++++++++++++++ 2 files changed, 262 insertions(+) create mode 100644 src/tabpfn_client/endpoints/__init__.py create mode 100644 src/tabpfn_client/endpoints/estimator.py diff --git a/src/tabpfn_client/endpoints/__init__.py b/src/tabpfn_client/endpoints/__init__.py new file mode 100644 index 0000000..f2fad8c --- /dev/null +++ b/src/tabpfn_client/endpoints/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Prior Labs GmbH 2026. +# Licensed under the Apache License, Version 2.0 +"""HTTP client for a self-hosted TabPFN inference container. + +from tabpfn_client.endpoints import TabPFNClassifier, TabPFNRegressor +""" + +from tabpfn_client.endpoints.estimator import TabPFNClassifier, TabPFNRegressor + + +__all__ = ["TabPFNClassifier", "TabPFNRegressor"] diff --git a/src/tabpfn_client/endpoints/estimator.py b/src/tabpfn_client/endpoints/estimator.py new file mode 100644 index 0000000..4b41b7f --- /dev/null +++ b/src/tabpfn_client/endpoints/estimator.py @@ -0,0 +1,251 @@ +# Copyright (c) Prior Labs GmbH 2026. +# Licensed under the Apache License, Version 2.0 +"""scikit-learn estimators for a TabPFN inference container reached over HTTP. + +Mirrors the `tabpfn_client.TabPFNClassifier` / `TabPFNRegressor` surface; +each `predict*` call POSTs to the user-supplied `endpoint_url` (the full +scoring URL, including the `/predict` path). `fit()` does not call the +endpoint — it just stores `X` / `y` on the estimator. The training data +is shipped to the endpoint on the next `predict*` call, where the actual +fit runs. + +Auth is optional: when `api_key` is set it is sent as `Bearer `, +otherwise no `Authorization` header is added. + +Cached predict is a per-call option: pass `model_id` to any `predict*` +call to send `context.model_id` and omit the training data, letting the +endpoint reuse an already-fit model. `fit()` is not required in that +case. When the endpoint returns a `model_id` in its response, it is +exposed as `self.model_id_` so callers can pass it forward. `model_path` +is similarly optional and only sent when set; some deployments reject +overrides. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Literal, Optional + +import httpx +import numpy as np +import pandas as pd +from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin +from sklearn.utils.validation import check_is_fitted + + +def _to_jsonable(X: Any) -> list: + """Coerce numpy / pandas inputs to plain Python lists for JSON.""" + if isinstance(X, pd.DataFrame): + return X.values.tolist() + if isinstance(X, pd.Series): + return X.tolist() + return np.asarray(X).tolist() + + +class _EndpointBase(BaseEstimator): + """Shared HTTP plumbing for the endpoint-backed TabPFN estimators.""" + + _TASK: str = "" # overridden by subclasses + + def __init__( + self, + endpoint_url: str, + api_key: Optional[str] = None, + extra_headers: Optional[Dict[str, str]] = None, + model_path: Optional[str] = None, + n_estimators: int = 8, + softmax_temperature: float = 0.9, + balance_probabilities: bool = False, + average_before_softmax: bool = False, + inference_precision: Literal["autocast", "auto"] = "auto", + random_state: Optional[int] = 0, + inference_config: Optional[Dict[str, Any]] = None, + n_preprocessing_jobs: int = 4, + memory_saving_mode: Optional[bool | Literal["auto"]] = None, + categorical_features_indices: Optional[List[int]] = None, + timeout_s: float = 300.0, + ): + self.endpoint_url = endpoint_url + self.api_key = api_key + self.extra_headers = extra_headers + self.model_path = model_path + self.n_estimators = n_estimators + self.softmax_temperature = softmax_temperature + self.balance_probabilities = balance_probabilities + self.average_before_softmax = average_before_softmax + self.inference_precision = inference_precision + self.random_state = random_state + self.inference_config = inference_config + self.n_preprocessing_jobs = n_preprocessing_jobs + self.memory_saving_mode = memory_saving_mode + self.categorical_features_indices = categorical_features_indices + self.timeout_s = timeout_s + + def _build_tabpfn_config(self) -> Dict[str, Any]: + cfg: Dict[str, Any] = { + "n_estimators": self.n_estimators, + "softmax_temperature": self.softmax_temperature, + "average_before_softmax": self.average_before_softmax, + "inference_precision": self.inference_precision, + "random_state": self.random_state, + "inference_config": self.inference_config, + "n_preprocessing_jobs": self.n_preprocessing_jobs, + } + if self.memory_saving_mode is not None: + cfg["memory_saving_mode"] = self.memory_saving_mode + if self.categorical_features_indices is not None: + cfg["categorical_features_indices"] = self.categorical_features_indices + if self.model_path is not None: + cfg["model_path"] = self.model_path + if self._TASK == "classification": + cfg["balance_probabilities"] = self.balance_probabilities + return cfg + + def _headers(self) -> Dict[str, str]: + headers: Dict[str, str] = { + "Content-Type": "application/json", + "Accept": "application/json", + } + # The API key is optional because some deployments do not require it + if self.api_key is not None: + headers["Authorization"] = f"Bearer {self.api_key}" + if self.extra_headers: + headers.update(self.extra_headers) + return headers + + def _http_client(self) -> httpx.Client: + # Cache the httpx.Client so repeated predict* calls reuse the TCP / + # TLS connection (keep-alive) instead of redoing the handshake on + # every request. + client = getattr(self, "_cached_client", None) + if client is not None: + return client + client = httpx.Client(timeout=self.timeout_s) + self._cached_client = client + return client + + def __getstate__(self) -> Dict[str, Any]: + # httpx.Client isn't pickleable; strip the cache for sklearn pickling. + state = self.__dict__.copy() + state.pop("_cached_client", None) + return state + + def fit(self, X: Any, y: Any) -> "_EndpointBase": + X_arr = X if isinstance(X, pd.DataFrame) else np.asarray(X) + y_arr = y if isinstance(y, (pd.DataFrame, pd.Series)) else np.asarray(y) + if X_arr.shape[0] != y_arr.shape[0]: + raise ValueError( + f"X and y must have the same number of samples; " + f"got X={X_arr.shape}, y={y_arr.shape}" + ) + self.X_train_ = X_arr + self.y_train_ = y_arr + if self._TASK == "classification": + self.classes_ = np.unique(y_arr) + return self + + def _invoke( + self, + X_test: Any, + output_type: str, + predict_params: Optional[Dict[str, Any]] = None, + model_id: Optional[str] = None, + ) -> Dict[str, Any]: + params: Dict[str, Any] = {"output_type": output_type} + if predict_params: + params.update(predict_params) + + body: Dict[str, Any] = { + "task_config": { + "task": self._TASK, + "tabpfn_config": self._build_tabpfn_config(), + "predict_params": params, + }, + "X_test": _to_jsonable(X_test), + } + + # Allow passing model_id to reference a previously fitted + # model and run predict by skipping the ICL step + if model_id is not None: + body["context"] = {"model_id": model_id} + else: + check_is_fitted(self, ["X_train_", "y_train_"]) + # y_train on the wire is 2D (n_samples, 1). + y_arr = np.asarray(self.y_train_) + if y_arr.ndim == 1: + y_arr = y_arr.reshape(-1, 1) + body["X_train"] = _to_jsonable(self.X_train_) + body["y_train"] = y_arr.tolist() + + resp = self._http_client().post( + self.endpoint_url, + json=body, + headers=self._headers(), + ) + resp.raise_for_status() + payload = resp.json() + + returned_id = payload.get("model_id") + if returned_id is not None: + self.model_id_ = returned_id + + return payload + + +class TabPFNClassifier(_EndpointBase, ClassifierMixin): + """TabPFN classifier backed by a self-hosted inference endpoint. + + Example: + from tabpfn_client.endpoints import TabPFNClassifier + clf = TabPFNClassifier( + endpoint_url="https:///predict", + api_key="", + ) + clf.fit(X_train, y_train) + clf.predict(X_test) + clf.predict_proba(X_test) + """ + + _TASK = "classification" + + def predict(self, X: Any, model_id: Optional[str] = None) -> np.ndarray: + result = self._invoke(X, output_type="preds", model_id=model_id) + return np.asarray(result["prediction"]) + + def predict_proba(self, X: Any, model_id: Optional[str] = None) -> np.ndarray: + result = self._invoke(X, output_type="probas", model_id=model_id) + return np.asarray(result["prediction"]) + + +class TabPFNRegressor(_EndpointBase, RegressorMixin): + """TabPFN regressor backed by a self-hosted inference endpoint. + + Example: + from tabpfn_client.endpoints import TabPFNRegressor + reg = TabPFNRegressor( + endpoint_url="https:///predict", + api_key="", + ) + reg.fit(X_train, y_train) + reg.predict(X_test) + reg.predict(X_test, output_type="quantiles", quantiles=[0.1, 0.5, 0.9]) + """ + + _TASK = "regression" + + def predict( + self, + X: Any, + output_type: str = "mean", + quantiles: Optional[list] = None, + model_id: Optional[str] = None, + ) -> np.ndarray: + predict_params: Dict[str, Any] = {} + if quantiles is not None: + predict_params["quantiles"] = quantiles + result = self._invoke( + X, + output_type=output_type, + predict_params=predict_params, + model_id=model_id, + ) + return np.asarray(result["prediction"]) From e3f3a855a00d58f55c6922d942b280cdb74a8af0 Mon Sep 17 00:00:00 2001 From: Dominik Safaric Date: Sat, 18 Jul 2026 21:47:58 +0200 Subject: [PATCH 2/5] Minor fixes --- src/tabpfn_client/endpoints/estimator.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/src/tabpfn_client/endpoints/estimator.py b/src/tabpfn_client/endpoints/estimator.py index 4b41b7f..3b80706 100644 --- a/src/tabpfn_client/endpoints/estimator.py +++ b/src/tabpfn_client/endpoints/estimator.py @@ -15,10 +15,11 @@ Cached predict is a per-call option: pass `model_id` to any `predict*` call to send `context.model_id` and omit the training data, letting the endpoint reuse an already-fit model. `fit()` is not required in that -case. When the endpoint returns a `model_id` in its response, it is -exposed as `self.model_id_` so callers can pass it forward. `model_path` -is similarly optional and only sent when set; some deployments reject -overrides. +case. To make the endpoint build a reusable cache in the first place, +set `fit_mode="fit_with_cache"` on the constructor; the endpoint then +returns a `model_id` which is exposed as `self.model_id_` so callers +can pass it forward. `model_path` is similarly optional and only sent +when set; some deployments reject overrides. """ from __future__ import annotations @@ -38,6 +39,8 @@ def _to_jsonable(X: Any) -> list: return X.values.tolist() if isinstance(X, pd.Series): return X.tolist() + if isinstance(X, list): + return X return np.asarray(X).tolist() @@ -52,6 +55,9 @@ def __init__( api_key: Optional[str] = None, extra_headers: Optional[Dict[str, str]] = None, model_path: Optional[str] = None, + fit_mode: Optional[ + Literal["fit_preprocessors", "low_memory", "fit_with_cache", "batched"] + ] = None, n_estimators: int = 8, softmax_temperature: float = 0.9, balance_probabilities: bool = False, @@ -68,6 +74,7 @@ def __init__( self.api_key = api_key self.extra_headers = extra_headers self.model_path = model_path + self.fit_mode = fit_mode self.n_estimators = n_estimators self.softmax_temperature = softmax_temperature self.balance_probabilities = balance_probabilities @@ -96,6 +103,8 @@ def _build_tabpfn_config(self) -> Dict[str, Any]: cfg["categorical_features_indices"] = self.categorical_features_indices if self.model_path is not None: cfg["model_path"] = self.model_path + if self.fit_mode is not None: + cfg["fit_mode"] = self.fit_mode if self._TASK == "classification": cfg["balance_probabilities"] = self.balance_probabilities return cfg @@ -139,6 +148,10 @@ def fit(self, X: Any, y: Any) -> "_EndpointBase": ) self.X_train_ = X_arr self.y_train_ = y_arr + # New training data invalidates any cached model_id captured from a + # prior predict; otherwise `predict(X, model_id=self.model_id_)` would + # keep hitting the old server-side model. + self.model_id_ = None if self._TASK == "classification": self.classes_ = np.unique(y_arr) return self From e25c6fe11aa9e8a6aa9d1f3715520a033deb3fad Mon Sep 17 00:00:00 2001 From: Dominik Safaric Date: Tue, 21 Jul 2026 10:23:40 +0200 Subject: [PATCH 3/5] Leave sklearn interface intact --- .../{endpoints => hosted}/__init__.py | 4 +- .../{endpoints => hosted}/estimator.py | 55 ++++++++++--------- 2 files changed, 32 insertions(+), 27 deletions(-) rename src/tabpfn_client/{endpoints => hosted}/__init__.py (58%) rename src/tabpfn_client/{endpoints => hosted}/estimator.py (83%) diff --git a/src/tabpfn_client/endpoints/__init__.py b/src/tabpfn_client/hosted/__init__.py similarity index 58% rename from src/tabpfn_client/endpoints/__init__.py rename to src/tabpfn_client/hosted/__init__.py index f2fad8c..d8db77a 100644 --- a/src/tabpfn_client/endpoints/__init__.py +++ b/src/tabpfn_client/hosted/__init__.py @@ -2,10 +2,10 @@ # Licensed under the Apache License, Version 2.0 """HTTP client for a self-hosted TabPFN inference container. -from tabpfn_client.endpoints import TabPFNClassifier, TabPFNRegressor +from tabpfn_client.hosted import TabPFNClassifier, TabPFNRegressor """ -from tabpfn_client.endpoints.estimator import TabPFNClassifier, TabPFNRegressor +from tabpfn_client.hosted.estimator import TabPFNClassifier, TabPFNRegressor __all__ = ["TabPFNClassifier", "TabPFNRegressor"] diff --git a/src/tabpfn_client/endpoints/estimator.py b/src/tabpfn_client/hosted/estimator.py similarity index 83% rename from src/tabpfn_client/endpoints/estimator.py rename to src/tabpfn_client/hosted/estimator.py index 3b80706..63e7fb5 100644 --- a/src/tabpfn_client/endpoints/estimator.py +++ b/src/tabpfn_client/hosted/estimator.py @@ -12,14 +12,19 @@ Auth is optional: when `api_key` is set it is sent as `Bearer `, otherwise no `Authorization` header is added. -Cached predict is a per-call option: pass `model_id` to any `predict*` -call to send `context.model_id` and omit the training data, letting the -endpoint reuse an already-fit model. `fit()` is not required in that -case. To make the endpoint build a reusable cache in the first place, -set `fit_mode="fit_with_cache"` on the constructor; the endpoint then -returns a `model_id` which is exposed as `self.model_id_` so callers -can pass it forward. `model_path` is similarly optional and only sent -when set; some deployments reject overrides. +Cached predict is opt-in via the `model_id` constructor argument. +When set, `predict*` sends `context.model_id` and omits the training +data, letting the endpoint reuse an already-fit model — `fit()` is not +required. To make the endpoint build a reusable cache in the first +place, set `fit_mode="fit_with_cache"`; the endpoint then returns a +`model_id` in its response, captured on the fitted attribute +`self.model_id_` so callers can construct a follow-up estimator with +`model_id=prior.model_id_`. `model_path` is similarly optional and +only sent when set; some deployments reject overrides. + +`model_id` lives on the constructor (not on `predict`) so the sklearn +estimator contract stays intact: `predict(X)` / `predict_proba(X)` work +with `Pipeline`, `GridSearchCV`, `cross_validate`, etc. """ from __future__ import annotations @@ -44,7 +49,7 @@ def _to_jsonable(X: Any) -> list: return np.asarray(X).tolist() -class _EndpointBase(BaseEstimator): +class _HostedBase(BaseEstimator): """Shared HTTP plumbing for the endpoint-backed TabPFN estimators.""" _TASK: str = "" # overridden by subclasses @@ -54,6 +59,7 @@ def __init__( endpoint_url: str, api_key: Optional[str] = None, extra_headers: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, model_path: Optional[str] = None, fit_mode: Optional[ Literal["fit_preprocessors", "low_memory", "fit_with_cache", "batched"] @@ -73,6 +79,7 @@ def __init__( self.endpoint_url = endpoint_url self.api_key = api_key self.extra_headers = extra_headers + self.model_id = model_id self.model_path = model_path self.fit_mode = fit_mode self.n_estimators = n_estimators @@ -138,7 +145,7 @@ def __getstate__(self) -> Dict[str, Any]: state.pop("_cached_client", None) return state - def fit(self, X: Any, y: Any) -> "_EndpointBase": + def fit(self, X: Any, y: Any) -> "_HostedBase": X_arr = X if isinstance(X, pd.DataFrame) else np.asarray(X) y_arr = y if isinstance(y, (pd.DataFrame, pd.Series)) else np.asarray(y) if X_arr.shape[0] != y_arr.shape[0]: @@ -161,7 +168,6 @@ def _invoke( X_test: Any, output_type: str, predict_params: Optional[Dict[str, Any]] = None, - model_id: Optional[str] = None, ) -> Dict[str, Any]: params: Dict[str, Any] = {"output_type": output_type} if predict_params: @@ -176,10 +182,11 @@ def _invoke( "X_test": _to_jsonable(X_test), } - # Allow passing model_id to reference a previously fitted - # model and run predict by skipping the ICL step - if model_id is not None: - body["context"] = {"model_id": model_id} + # When the constructor was given a model_id, reference the previously + # fitted server-side model and skip the ICL step. Otherwise the + # estimator must have been fit() first — ship the training data. + if self.model_id is not None: + body["context"] = {"model_id": self.model_id} else: check_is_fitted(self, ["X_train_", "y_train_"]) # y_train on the wire is 2D (n_samples, 1). @@ -204,11 +211,11 @@ def _invoke( return payload -class TabPFNClassifier(_EndpointBase, ClassifierMixin): +class TabPFNClassifier(_HostedBase, ClassifierMixin): """TabPFN classifier backed by a self-hosted inference endpoint. Example: - from tabpfn_client.endpoints import TabPFNClassifier + from tabpfn_client.hosted import TabPFNClassifier clf = TabPFNClassifier( endpoint_url="https:///predict", api_key="", @@ -220,20 +227,20 @@ class TabPFNClassifier(_EndpointBase, ClassifierMixin): _TASK = "classification" - def predict(self, X: Any, model_id: Optional[str] = None) -> np.ndarray: - result = self._invoke(X, output_type="preds", model_id=model_id) + def predict(self, X: Any) -> np.ndarray: + result = self._invoke(X, output_type="preds") return np.asarray(result["prediction"]) - def predict_proba(self, X: Any, model_id: Optional[str] = None) -> np.ndarray: - result = self._invoke(X, output_type="probas", model_id=model_id) + def predict_proba(self, X: Any) -> np.ndarray: + result = self._invoke(X, output_type="probas") return np.asarray(result["prediction"]) -class TabPFNRegressor(_EndpointBase, RegressorMixin): +class TabPFNRegressor(_HostedBase, RegressorMixin): """TabPFN regressor backed by a self-hosted inference endpoint. Example: - from tabpfn_client.endpoints import TabPFNRegressor + from tabpfn_client.hosted import TabPFNRegressor reg = TabPFNRegressor( endpoint_url="https:///predict", api_key="", @@ -250,7 +257,6 @@ def predict( X: Any, output_type: str = "mean", quantiles: Optional[list] = None, - model_id: Optional[str] = None, ) -> np.ndarray: predict_params: Dict[str, Any] = {} if quantiles is not None: @@ -259,6 +265,5 @@ def predict( X, output_type=output_type, predict_params=predict_params, - model_id=model_id, ) return np.asarray(result["prediction"]) From 03c33b127a7b194e54b75ac65679dd7b083e911e Mon Sep 17 00:00:00 2001 From: Dominik Safaric Date: Tue, 21 Jul 2026 13:42:40 +0200 Subject: [PATCH 4/5] Update client version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index dd9252b..bcee834 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "tabpfn-client" -version = "0.3.3" +version = "0.3.4" requires-python = ">=3.10" description = "API access for TabPFN: Foundation model for tabular data" From 80968b9552de1a861c7eb905ac98cf67166fdbf7 Mon Sep 17 00:00:00 2001 From: Dominik Safaric Date: Tue, 21 Jul 2026 14:11:15 +0200 Subject: [PATCH 5/5] Fix inconsistency in fit vs model_id reusability --- src/tabpfn_client/hosted/estimator.py | 28 +++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/src/tabpfn_client/hosted/estimator.py b/src/tabpfn_client/hosted/estimator.py index 63e7fb5..c920d5f 100644 --- a/src/tabpfn_client/hosted/estimator.py +++ b/src/tabpfn_client/hosted/estimator.py @@ -13,9 +13,12 @@ otherwise no `Authorization` header is added. Cached predict is opt-in via the `model_id` constructor argument. -When set, `predict*` sends `context.model_id` and omits the training -data, letting the endpoint reuse an already-fit model — `fit()` is not -required. To make the endpoint build a reusable cache in the first +When set (and `fit()` has not been called), `predict*` sends +`context.model_id` and omits the training data, letting the endpoint +reuse an already-fit model. Calling `fit()` on the estimator supersedes +the constructor `model_id` — the freshly-stored training data is +shipped on the next `predict*`, matching sklearn's re-fit semantics. +To make the endpoint build a reusable cache in the first place, set `fit_mode="fit_with_cache"`; the endpoint then returns a `model_id` in its response, captured on the fitted attribute `self.model_id_` so callers can construct a follow-up estimator with @@ -182,19 +185,24 @@ def _invoke( "X_test": _to_jsonable(X_test), } - # When the constructor was given a model_id, reference the previously - # fitted server-side model and skip the ICL step. Otherwise the - # estimator must have been fit() first — ship the training data. - if self.model_id is not None: - body["context"] = {"model_id": self.model_id} - else: - check_is_fitted(self, ["X_train_", "y_train_"]) + # Precedence: a completed fit() wins over the constructor's model_id. + # Rationale — if the caller ran fit(new_X, new_y) on an estimator that + # was constructed with model_id set, they have re-trained; the cached + # id is now stale relative to the new data. Falls back to the cached + # path only when fit() was never called. + has_training_data = hasattr(self, "X_train_") and hasattr(self, "y_train_") + if has_training_data: # y_train on the wire is 2D (n_samples, 1). y_arr = np.asarray(self.y_train_) if y_arr.ndim == 1: y_arr = y_arr.reshape(-1, 1) body["X_train"] = _to_jsonable(self.X_train_) body["y_train"] = y_arr.tolist() + elif self.model_id is not None: + body["context"] = {"model_id": self.model_id} + else: + # Neither path available — raise the standard sklearn error. + check_is_fitted(self, ["X_train_", "y_train_"]) resp = self._http_client().post( self.endpoint_url,