diff --git a/src/tabpfn_client/api_models.py b/src/tabpfn_client/api_models.py index f7ecc4f..6e41680 100644 --- a/src/tabpfn_client/api_models.py +++ b/src/tabpfn_client/api_models.py @@ -118,6 +118,9 @@ class ModelLimit(BaseModel): max_cols: int test_set_max_rows_w_full_regression_output: int test_set_max_cells: int + # Optional until all deployed servers send it; estimator.py falls back to + # its local default when absent. + predict_row_pairs_budget: int | None = None class ModelVersion(str, Enum): diff --git a/src/tabpfn_client/estimator.py b/src/tabpfn_client/estimator.py index 811057e..98db184 100644 --- a/src/tabpfn_client/estimator.py +++ b/src/tabpfn_client/estimator.py @@ -66,6 +66,13 @@ THINKING_TIMEOUT_MAX_S = 40 * 60 +# Prediction compute scales with n_train_rows * n_test_rows, so the API caps +# their product. The effective per-call test row limit therefore shrinks as +# the fitted training set grows (e.g. 1M training rows -> 250k test rows). +# The server sends its budget via get_model_limits (`predict_row_pairs_budget`); +# this is only the fallback for servers that predate that field. +FALLBACK_PREDICT_ROW_PAIRS_BUDGET = 250_000 * 1_000_000 + _VALID_THINKING_EFFORT_LEVELS = frozenset({"medium", "high"}) @@ -404,7 +411,14 @@ def _predict( predict_params=predict_params, ) - validate_test_set(X, output_type, tabpfn_config.model_path) + validate_test_set( + X, + output_type, + tabpfn_config.model_path, + train_rows=self._last_train_X.shape[0] + if self._last_train_X is not None + else None, + ) X_clean = _clean_text_features(X) if ( @@ -742,7 +756,14 @@ def predict( predict_params=predict_params, ) - validate_test_set(X, output_type, tabpfn_config.model_path) + validate_test_set( + X, + output_type, + tabpfn_config.model_path, + train_rows=self._last_train_X.shape[0] + if self._last_train_X is not None + else None, + ) X_clean = _clean_text_features(X) if ( @@ -921,6 +942,7 @@ def validate_test_set( X: pd.DataFrame | np.ndarray, output_type: str | None, model_path: str | None = None, + train_rows: int | None = None, ): """Check the integrity of the test data.""" @@ -934,9 +956,14 @@ def validate_test_set( model_version = model_version_from_path(model_path) limit = model_limit_from_version(model_version, limits.model_limits) - if X.shape[0] > limit.test_set_max_rows: + max_rows = limit.test_set_max_rows + if train_rows: + budget = limit.predict_row_pairs_budget or FALLBACK_PREDICT_ROW_PAIRS_BUDGET + max_rows = min(max_rows, budget // train_rows) + + if X.shape[0] > max_rows: raise ValueError( - f"The number of test rows ({X.shape[0]}) exceeds the maximum of {limit.test_set_max_rows}. " + f"The number of test rows ({X.shape[0]}) exceeds the maximum of {max_rows}. " "Split the test set across multiple calls to reduce the number of rows." ) if X.shape[1] > limit.max_cols: diff --git a/tests/unit/test_validate_test_set.py b/tests/unit/test_validate_test_set.py new file mode 100644 index 0000000..a5fa6e8 --- /dev/null +++ b/tests/unit/test_validate_test_set.py @@ -0,0 +1,91 @@ +from typing import Any +from unittest.mock import patch + +import numpy as np +import pytest + +from tabpfn_client.api_models import GetModelLimitsResponse +from tabpfn_client.client import ServiceClient +from tabpfn_client.estimator import FALLBACK_PREDICT_ROW_PAIRS_BUDGET, validate_test_set + + +def _limits( + test_set_max_rows: int = 1_000_000, + predict_row_pairs_budget: int | None = FALLBACK_PREDICT_ROW_PAIRS_BUDGET, +) -> GetModelLimitsResponse: + model_limit: dict[str, Any] = { + "train_set_max_rows": 1_000_000, + "train_set_max_cells": 100_000_000, + "test_set_max_rows": test_set_max_rows, + "test_set_max_cells": 100_000_000, + "test_set_max_rows_w_full_regression_output": 400, + "max_cols": 2_000, + "max_classes": 10, + } + if predict_row_pairs_budget is not None: + model_limit["predict_row_pairs_budget"] = predict_row_pairs_budget + return GetModelLimitsResponse.model_validate( + { + "default_model_version": "v3", + "max_model_limit": model_limit, + "model_limits": {"v3": model_limit}, + "dataset_max_size_bytes": 100_000_000, + } + ) + + +def _X(n_rows: int) -> np.ndarray: + return np.zeros((n_rows, 1)) + + +def _validate( + n_test_rows: int, + train_rows: int | None = None, + limits: GetModelLimitsResponse | None = None, +): + with patch.object( + ServiceClient, "get_model_limits", return_value=limits or _limits() + ): + validate_test_set(_X(n_test_rows), None, train_rows=train_rows) + + +def test_static_limit_applies_without_train_rows(): + _validate(1_000_000) + with pytest.raises(ValueError, match="exceeds the maximum of 1000000"): + _validate(1_000_001) + + +def test_adaptive_limit_shrinks_with_train_rows(): + # 1M train rows -> 250k test rows per call + _validate(250_000, train_rows=1_000_000) + with pytest.raises(ValueError, match="exceeds the maximum of 250000"): + _validate(250_001, train_rows=1_000_000) + + +def test_adaptive_limit_never_exceeds_static_limit(): + # Small train sets leave the static cap binding + small_train = FALLBACK_PREDICT_ROW_PAIRS_BUDGET // 1_000_000 // 10 + _validate(1_000_000, train_rows=small_train) + with pytest.raises(ValueError, match="exceeds the maximum of 1000000"): + _validate(1_000_001, train_rows=small_train) + + +def test_server_sent_budget_takes_precedence(): + # Half the fallback budget -> 125k test rows for 1M train rows + limits = _limits(predict_row_pairs_budget=125_000 * 1_000_000) + _validate(125_000, train_rows=1_000_000, limits=limits) + with pytest.raises(ValueError, match="exceeds the maximum of 125000"): + _validate(125_001, train_rows=1_000_000, limits=limits) + + +def test_fallback_budget_applies_when_server_omits_it(): + # Older servers don't send predict_row_pairs_budget + limits = _limits(predict_row_pairs_budget=None) + _validate(250_000, train_rows=1_000_000, limits=limits) + with pytest.raises(ValueError, match="exceeds the maximum of 250000"): + _validate(250_001, train_rows=1_000_000, limits=limits) + + +def test_no_limits_available_skips_validation(): + with patch.object(ServiceClient, "get_model_limits", return_value=None): + validate_test_set(_X(2_000_000), None, train_rows=1_000_000)