diff --git a/changelog/369.fixed.md b/changelog/369.fixed.md new file mode 100644 index 0000000..fc05d9b --- /dev/null +++ b/changelog/369.fixed.md @@ -0,0 +1 @@ +`TabPFNRegressor.predict(output_type="full")` now matches a local prediction: masked-out bars come back as `-inf` rather than `NaN`, and test sets above the server's full-output row cap are split across calls instead of raising. diff --git a/src/tabpfn_client/estimator.py b/src/tabpfn_client/estimator.py index ca22182..4c97051 100644 --- a/src/tabpfn_client/estimator.py +++ b/src/tabpfn_client/estimator.py @@ -29,6 +29,7 @@ from tabpfn_client.utils import model_limit_from_version, model_version_from_path from tabpfn_client.service_wrapper import InferenceClient from tabpfn_client.api_models import ( + ModelLimit, RegressorTabPFNConfig, ClassifierTabPFNConfig, RegressorPredictParams, @@ -794,6 +795,16 @@ def predict( predict_params=predict_params, ) + # A full-output response carries one logit per histogram bar for every + # test row, so the server caps the rows a single response may cover. + # Splitting the call here keeps the returned arrays identical to what + # one unrestricted call would have produced. + rows_per_call = ( + _full_output_row_limit(tabpfn_config.model_path) + if output_type == "full" + else None + ) + chunked = rows_per_call is not None and X.shape[0] > rows_per_call validate_test_set( X, output_type, @@ -801,8 +812,8 @@ def predict( train_rows=self._last_train_X.shape[0] if self._last_train_X is not None else None, + split_full_output=chunked, ) - X_clean = _clean_text_features(X) # NOTE(@trace_id) # If this instance reuses a previous fit via a directly-assigned @@ -816,34 +827,58 @@ def predict( ): self.client_options.headers["sentry-trace"] = self._last_trace_id - def predict_task() -> PredictionResult: - return InferenceClient.predict( - X_clean, - fitted_train_set_id=self.model_id_, - task_config=task_config, - client_options=self.client_options, - ) + def predict_rows(X_rows: Any) -> PredictionResult: + X_clean = _clean_text_features(X_rows) - result = run_task(predict_task, "Predicting") - # Unpack and store metadata - self._last_meta = result.metadata + def predict_task() -> PredictionResult: + return InferenceClient.predict( + X_clean, + fitted_train_set_id=self.model_id_, + task_config=task_config, + client_options=self.client_options, + ) + + return run_task(predict_task, "Predicting") + + if chunked: + rows_per_call = cast(int, rows_per_call) + results = [ + predict_rows(_row_slice(X, start, start + rows_per_call)) + for start in range(0, X.shape[0], rows_per_call) + ] + # Metadata describes the request, and every chunk shares the same + # config; the last one stands for the whole prediction. + self._last_meta = results[-1].metadata + output = _merge_full_outputs( + [cast("dict[str, np.ndarray]", r.y_pred) for r in results] + ) + else: + result = predict_rows(X) + # Unpack and store metadata + self._last_meta = result.metadata + output = result.y_pred - output = result.y_pred if output_type == "quantiles" and isinstance(output, np.ndarray): return list(output) if output.ndim == 2 else [output] if output_type == "full": + # `criterion` is a bar distribution rather than an array, so the + # full output is looser than the declared `dict[str, np.ndarray]`. + full = cast("dict[str, Any]", output) + if "logits" in full: + full["logits"] = _restore_masked_logits(full["logits"]) try: from tabpfn.regressor import FullSupportBarDistribution # type: ignore import torch # type: ignore - output["criterion"] = FullSupportBarDistribution( - borders=torch.tensor(output["borders"]) + full["criterion"] = FullSupportBarDistribution( + borders=torch.tensor(full["borders"]) ) except ImportError: logger.warning( "Optional dependencies 'tabpfn' and 'torch' are required to " "construct the criterion when output_type='full'. Skipping criterion." ) + return full return output @@ -906,23 +941,39 @@ def validate_train_set( ) +def _limit_for_model_path(model_path: str | None) -> ModelLimit | None: + """Return the row/cell caps for `model_path`, or None if unknown.""" + api_settings = ServiceClient.get_settings() + if api_settings is None: + return None + if not model_path: + return api_settings.model_limits[api_settings.default_model_version] + model_version = model_version_from_path(model_path) + return model_limit_from_version(model_version, api_settings.model_limits) + + +def _full_output_row_limit(model_path: str | None) -> int | None: + """Rows a single `output_type="full"` response may cover, None if unknown.""" + limit = _limit_for_model_path(model_path) + return limit.test_set_max_rows_w_full_regression_output if limit else None + + def validate_test_set( X: pd.DataFrame | np.ndarray, output_type: str | None, model_path: str | None = None, train_rows: int | None = None, + split_full_output: bool = False, ): - """Check the integrity of the test data.""" + """Check the integrity of the test data. - api_settings = ServiceClient.get_settings() - if api_settings is None: - return + `split_full_output` marks that the caller will honour the full-output row + cap by splitting the request, so that cap is not enforced here. + """ - if not model_path: - limit = api_settings.model_limits[api_settings.default_model_version] - else: - model_version = model_version_from_path(model_path) - limit = model_limit_from_version(model_version, api_settings.model_limits) + limit = _limit_for_model_path(model_path) + if limit is None: + return max_rows = limit.test_set_max_rows if train_rows: @@ -944,7 +995,7 @@ def validate_test_set( f"The number of test cells ({n_cells}) exceeds the maximum of {limit.test_set_max_cells}. " "Split the test set across multiple calls to reduce the number of cells." ) - if output_type == "full": + if output_type == "full" and not split_full_output: if X.shape[0] > limit.test_set_max_rows_w_full_regression_output: raise ValueError( f"The number of test rows ({X.shape[0]}) exceeds the maximum of {limit.test_set_max_rows_w_full_regression_output} " @@ -952,6 +1003,57 @@ def validate_test_set( ) +def _restore_masked_logits(logits: np.ndarray) -> np.ndarray: + """Return `logits` with bars outside a row's support back at -inf. + + The response encoding carries no representation for -inf, so those bars + arrive as null and land in the array as NaN. + """ + logits = np.asarray(logits, dtype=float) + return np.where(np.isnan(logits), -np.inf, logits) + + +def _row_slice(X: Any, start: int, stop: int) -> Any: + """Return rows `[start:stop)` of `X`, keeping its container type.""" + if isinstance(X, pd.DataFrame): + return X.iloc[start:stop] + return X[start:stop] + + +# Axis along which each `output_type="full"` array runs over test rows. None +# marks a row-independent array: `borders` describes the histogram of the +# fitted target, so every chunk returns the same one. +_FULL_OUTPUT_ROW_AXIS: dict[str, int | None] = { + "mean": 0, + "median": 0, + "mode": 0, + "logits": 0, + "quantiles": 1, + "borders": None, +} + + +def _merge_full_outputs(parts: list[dict[str, np.ndarray]]) -> dict[str, np.ndarray]: + """Stitch per-chunk full-output predictions back into one result.""" + unknown = sorted(set(parts[0]) - set(_FULL_OUTPUT_ROW_AXIS)) + if unknown: + raise RuntimeError( + f"Cannot combine full regression output across calls: the server " + f"returned unrecognised field(s) {unknown}. Upgrade tabpfn-client, " + f"or split the test set yourself and merge the results." + ) + merged: dict[str, np.ndarray] = {} + for key, axis in _FULL_OUTPUT_ROW_AXIS.items(): + if key not in parts[0]: + continue + merged[key] = ( + parts[0][key] + if axis is None + else np.concatenate([p[key] for p in parts], axis=axis) + ) + return merged + + @overload def _clean_text_features(X: pd.DataFrame) -> pd.DataFrame: ... @overload diff --git a/tests/unit/test_tabpfn_regressor.py b/tests/unit/test_tabpfn_regressor.py index 95d3510..17163d3 100644 --- a/tests/unit/test_tabpfn_regressor.py +++ b/tests/unit/test_tabpfn_regressor.py @@ -982,3 +982,132 @@ def test_paper_version_behavior(self, mock_predict, mock_fit): tabpfn_false.fit(X, y) y_pred_false = tabpfn_false.predict(test_X) self.assertIsNotNone(y_pred_false) + + +class TestFullOutputChunking(unittest.TestCase): + """`output_type="full"` splits into per-chunk calls above the server cap.""" + + FULL_OUTPUT_MAX_ROWS = 4 + N_BARS = 6 + N_QUANTILES = 9 + + def setUp(self): + # skip init + config.Config.is_initialized = True + payload = _api_settings_payload() + for limit in [payload["max_model_limit"], *payload["model_limits"].values()]: + limit["test_set_max_rows_w_full_regression_output"] = ( + self.FULL_OUTPUT_MAX_ROWS + ) + ServiceClient._api_settings = GetSettingsResponse(**payload) + ServiceClient._api_settings_ts = time.monotonic() + + self.borders = np.linspace(0.0, 1.0, self.N_BARS + 1) + self.regressor = TabPFNRegressor(model_path="v2.5_default") + self.regressor.model_id_ = UUID("00000000-0000-0000-0000-000000000000") + self.regressor._last_train_X = np.random.randn(5, 2) + + def tearDown(self): + ServiceClient._api_settings = None + ServiceClient._api_settings_ts = 0.0 + # undo setUp + config.reset() + + def _full_output(self, n_rows: int, offset: int) -> dict[str, np.ndarray]: + """A full-output payload whose values encode each row's global index.""" + rows = np.arange(offset, offset + n_rows, dtype=float) + return { + "mean": rows, + "median": rows + 0.1, + "mode": rows + 0.2, + "quantiles": rows + np.arange(self.N_QUANTILES)[:, None], + "logits": rows[:, None] + np.arange(self.N_BARS)[None, :], + "borders": self.borders, + } + + def _predict(self, n_test_rows: int): + """Predict `n_test_rows` rows, serving each request from its own chunk.""" + served = [] + + def fake_predict(X, **kwargs): + offset = sum(len(x) for x in served) + served.append(X) + return PredictionResult( + y_pred=cast(Any, self._full_output(len(X), offset)), metadata={} + ) + + with patch.object(InferenceClient, "predict", side_effect=fake_predict): + output = cast( + "dict[str, Any]", + self.regressor.predict( + np.random.randn(n_test_rows, 2), output_type="full" + ), + ) + return output, served + + def test_stays_a_single_call_at_the_limit(self): + _, served = self._predict(self.FULL_OUTPUT_MAX_ROWS) + self.assertEqual([len(x) for x in served], [self.FULL_OUTPUT_MAX_ROWS]) + + def test_splits_into_chunks_of_at_most_the_limit(self): + _, served = self._predict(self.FULL_OUTPUT_MAX_ROWS * 2 + 1) + self.assertEqual( + [len(x) for x in served], + [self.FULL_OUTPUT_MAX_ROWS, self.FULL_OUTPUT_MAX_ROWS, 1], + ) + + def test_merged_output_matches_a_single_unsplit_call(self): + n_rows = self.FULL_OUTPUT_MAX_ROWS * 2 + 1 + output, _ = self._predict(n_rows) + expected = self._full_output(n_rows, 0) + for key, value in expected.items(): + np.testing.assert_allclose(output[key], value, err_msg=key) + + def test_rejects_output_it_cannot_merge(self): + def fake_predict(X, **kwargs): + payload = self._full_output(len(X), 0) + payload["surprise"] = np.zeros(len(X)) + return PredictionResult(y_pred=cast(Any, payload), metadata={}) + + with patch.object(InferenceClient, "predict", side_effect=fake_predict): + with self.assertRaisesRegex(RuntimeError, "surprise"): + self.regressor.predict( + np.random.randn(self.FULL_OUTPUT_MAX_ROWS + 1, 2), + output_type="full", + ) + + def test_masked_bars_come_back_as_negative_infinity(self): + """Null logits mark bars outside a row's support, i.e. -inf.""" + + def fake_predict(X, **kwargs): + payload = self._full_output(len(X), 0) + payload["logits"] = payload["logits"].copy() + payload["logits"][:, -2:] = np.nan + return PredictionResult(y_pred=cast(Any, payload), metadata={}) + + with patch.object(InferenceClient, "predict", side_effect=fake_predict): + output = cast( + "dict[str, Any]", + self.regressor.predict( + np.random.randn(self.FULL_OUTPUT_MAX_ROWS + 1, 2), + output_type="full", + ), + ) + + logits = output["logits"] + self.assertFalse(np.isnan(logits).any()) + self.assertTrue(np.isneginf(logits[:, -2:]).all()) + self.assertTrue(np.isfinite(logits[:, :-2]).all()) + + def test_other_output_types_keep_the_row_limit(self): + """The cap only applies to `full`, so `mean` is never split.""" + served = [] + + def fake_predict(X, **kwargs): + served.append(X) + return PredictionResult(y_pred=np.zeros(len(X)), metadata={}) + + n_rows = self.FULL_OUTPUT_MAX_ROWS * 3 + with patch.object(InferenceClient, "predict", side_effect=fake_predict): + self.regressor.predict(np.random.randn(n_rows, 2), output_type="mean") + self.assertEqual([len(x) for x in served], [n_rows])