Skip to content

Enable fit_with_cache for predictions on API with K/V #347

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
safaricd wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Enable fit_with_cache for predictions on API with K/V #347

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
3d67a48
Enable fit_with_cache for predictions on API with K/V
safaricd Jul 23, 2026
bc2e672
Fit with cache support
safaricd Jul 23, 2026
4a9ca2f
Fix failing unit tests
safaricd Jul 23, 2026
0b0b8e2
Fix cursor findings on ENG-880: state desync, auth init, and CI issues
safaricd Jul 23, 2026
6e86ac9
Trim FitMode to values gapi actually wires
safaricd Jul 24, 2026
a1787b2
Apply PR review nits from @simo-prior
safaricd Jul 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion src/tabpfn_client/api_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,14 @@ class ClassifierPredictParams(BaseModel):
) = None


class FitMode(str, Enum):
# Only the two values that gapi wires end-to-end. The server-side enum
# also has internal `low_memory` / `batched` modes, but those are not
# accepted by the public API — see `FitRequest._validate_fit_mode`.
FIT_PREPROCESSORS = "fit_preprocessors"
FIT_WITH_CACHE = "fit_with_cache"


class ClassifierTabPFNConfig(BaseModel):
n_estimators: int | None = None
categorical_features_indices: list[int] | None = None
Expand All @@ -65,6 +73,7 @@ class ClassifierTabPFNConfig(BaseModel):
Annotated[Literal["autocast", "auto"] | str, Field(union_mode="left_to_right")] | None
) = None
ignore_pretraining_limits: bool | None = None
fit_mode: Annotated[FitMode | UnknownEnum, Field(union_mode="left_to_right")] | None = None
model_path: str | None = None
balance_probabilities: bool | None = None

Expand Down Expand Up @@ -161,6 +170,7 @@ class RegressorTabPFNConfig(BaseModel):
Annotated[Literal["autocast", "auto"] | str, Field(union_mode="left_to_right")] | None
) = None
ignore_pretraining_limits: bool | None = None
fit_mode: Annotated[FitMode | UnknownEnum, Field(union_mode="left_to_right")] | None = None
model_path: str | None = None


Expand Down Expand Up @@ -226,7 +236,7 @@ class FitRequest(BaseModel):
)
async_mode: bool | None = Field(
default=None,
description="Submit the fit and return immediately with status=pending; poll GET /tabpfn/fit/{fitted_train_set_id} for the outcome. Defaults to the blocking behaviour.",
description="Submit the fit and return immediately with status=pending; poll GET /fit/{fitted_train_set_id} for the outcome. Defaults to the blocking behaviour.",
)


Expand All @@ -239,6 +249,7 @@ class FitStatusResponse(BaseModel):
fitted_train_set_id: UUID
status: Annotated[FitStatus | UnknownEnum, Field(union_mode="left_to_right")]
error: str | None = None
error_code: str | None = None


class GetModelLimitsResponse(BaseModel):
Expand Down
Loading
Loading