diff --git a/changelog/1349.fixed.md b/changelog/1349.fixed.md new file mode 100644 index 000000000..71fea8faa --- /dev/null +++ b/changelog/1349.fixed.md @@ -0,0 +1 @@ +Build the categorical feature schema from the number of features rather than the number of training rows, so a categorical column at a high index is reported to the model as categorical. \ No newline at end of file diff --git a/src/tabpfn/classifier.py b/src/tabpfn/classifier.py index 959d97f7c..5880b52be 100644 --- a/src/tabpfn/classifier.py +++ b/src/tabpfn/classifier.py @@ -1082,9 +1082,11 @@ def fit_from_preprocessed( self.n_classes_ = max(int(t.max().item()) for t in y_preprocessed) + 1 self.classes_ = torch.arange(self.n_classes_) + # X_preprocessed is the collated batch: (n_datasets, n_train_rows, + # n_columns), so the column count is the last dimension. feature_schema = convert_batch_of_cat_ix_to_schema( batch_of_cat_indices=cat_ix, - num_features=X_preprocessed[0].shape[1], + num_features=X_preprocessed[0].shape[-1], ) # Preprocessed tensors hold no dates or strings either, so these fit diff --git a/src/tabpfn/regressor.py b/src/tabpfn/regressor.py index abc77ed41..e89e805bb 100644 --- a/src/tabpfn/regressor.py +++ b/src/tabpfn/regressor.py @@ -1147,9 +1147,11 @@ def fit_from_preprocessed( self.inference_precision, self.devices_ ) + # X_preprocessed is the collated batch: (n_datasets, n_train_rows, + # n_columns), so the column count is the last dimension. feature_schema = convert_batch_of_cat_ix_to_schema( batch_of_cat_indices=cat_ix, - num_features=X_preprocessed[0].shape[1], + num_features=X_preprocessed[0].shape[-1], ) # Preprocessed tensors hold no dates or strings either, so these fit diff --git a/tests/test_classifier_interface.py b/tests/test_classifier_interface.py index 32d3ae5ab..79feb2af0 100644 --- a/tests/test_classifier_interface.py +++ b/tests/test_classifier_interface.py @@ -24,6 +24,7 @@ from tabpfn import TabPFNClassifier from tabpfn.architectures import tabpfn_v2_5 +from tabpfn.architectures.interface import PerformanceOptions from tabpfn.base import ModelSpecs, initialize_tabpfn_model from tabpfn.constants import ModelVersion from tabpfn.inference_config import DEFAULT_SOFTMAX_TEMPERATURE, InferenceConfig @@ -34,6 +35,10 @@ ) from tabpfn.model_loading import ModelSource, prepend_cache_path from tabpfn.preprocessing import PreprocessorConfig +from tabpfn.preprocessing.datamodel import FeatureModality +from tabpfn.preprocessing.ensemble import ( + generate_classification_ensemble_configs, +) from tabpfn.utils import infer_devices from .utils import ( @@ -1742,3 +1747,142 @@ def mk(seed: int, n: int = 60, f: int = 5) -> tuple[np.ndarray, np.ndarray]: fitted.predict_proba_batched(X_list, y_list, X_tests) after = fitted.predict_proba(a_x[:5]) np.testing.assert_array_equal(before, after) + + +# ============================================================================= +# fit_from_preprocessed: the feature schema is built from the column count +# ============================================================================= + + +N_BATCH_ROWS = 5 +N_BATCH_COLS = 8 +BATCH_CAT_COL = N_BATCH_COLS - 1 # index 7 >= the 5 rows of a collated batch + + +def _collated_batch( + n_datasets: int = 2, n_estimators: int = 2 +) -> tuple[ + ModelSpecs, + list[torch.Tensor], + list[torch.Tensor], + list[list[list[int]]], + list[list[object]], +]: + """The (n_datasets, n_train_rows, n_columns) tensors the collator produces.""" + specs = _create_dummy_classifier_model_specs() + flat_configs = generate_classification_ensemble_configs( + num_estimators=n_estimators, + add_fingerprint_feature=False, + polynomial_features="no", + feature_shift_decoder="shuffle", + preprocessor_configs=[ + PreprocessorConfig( + "none", categorical_name="numeric", max_features_per_estimator=500 + ) + ], + class_shift_method="shuffle", + n_classes=2, + random_state=0, + num_models=1, + outlier_removal_std=None, + ) + configs: list[list[object]] = [ + [flat_configs[estimator] for estimator in range(n_estimators)] + for _ in range(n_datasets) + ] + torch.manual_seed(0) + X_preprocessed = [ + torch.randn(n_datasets, N_BATCH_ROWS, N_BATCH_COLS), + ] + y_preprocessed = [ + torch.randint(0, 2, (n_datasets, N_BATCH_ROWS)).float(), + ] + cat_ix = [[[BATCH_CAT_COL] for _ in range(n_estimators)] for _ in range(n_datasets)] + return specs, X_preprocessed, y_preprocessed, cat_ix, configs + + +def test__fit_from_preprocessed__feature_schema_spans_the_columns() -> None: + """The schema is as wide as the batch's columns, not its rows. + + X_preprocessed[0] is the collated (n_datasets, n_train_rows, n_columns) + tensor, so `shape[1]` is the row count. A schema built from it is narrower + than the data whenever there are more columns than rows, which silently + drops every categorical column whose index is >= n_train_rows and tells the + model those columns are numerical. + """ + specs, X_preprocessed, y_preprocessed, cat_ix, configs = _collated_batch() + + clf = TabPFNClassifier( + model_path=specs, device="cpu", n_estimators=2, random_state=0 + ) + clf.fit_from_preprocessed( + X_preprocessed, + y_preprocessed, + cat_ix, + configs, + performance_options=PerformanceOptions(), + ) + + assert X_preprocessed[0].shape[1] == N_BATCH_ROWS + assert X_preprocessed[0].shape[-1] == N_BATCH_COLS + for dataset_schemas in clf.executor_.feature_schema_list: + for schema in dataset_schemas: + assert len(schema.features) == N_BATCH_COLS + assert schema.indices_for(FeatureModality.CATEGORICAL) == [BATCH_CAT_COL] + + +def test__predict_proba_batched__tells_the_model_about_high_categorical_columns() -> ( + None +): + """Batched inference reports the same categorical columns as the single path. + + A wide dataset pushes the categorical past the row count after + preprocessing (the SVD/fingerprint steps add columns), so a row-wide schema + omits it and the model is told the column is numerical -- while + `predict_proba`, which uses each member's own schema, still reports it. + """ + n_rows = 100 # MIN_NUMBER_OF_SAMPLES_FOR_CATEGORICAL_INFERENCE + n_cols = 250 # wide enough that the categorical lands beyond n_train_rows + rng = np.random.RandomState(0) + X = rng.randn(n_rows, n_cols).astype(object) + X[:, -1] = rng.choice(["a", "b", "c", "d"], size=n_rows) + y = (X[:, -1] == "a").astype(int) + X_test = np.repeat(X[:1], 2, axis=0).astype(object) + X_test[:, -1] = ["a", "b"] + + clf = TabPFNClassifier( + model_path=_create_dummy_classifier_model_specs(), + device="cpu", + n_estimators=1, + random_state=0, + ) + clf.fit(X, y) + + member_cats = clf.executor_.ensemble_members[0].feature_schema.indices_for( + FeatureModality.CATEGORICAL + ) + # Precondition: this is the case the bug needs, so the test cannot pass + # vacuously on a schema that happened to keep the column. + assert [c for c in member_cats if c >= n_rows], ( + f"expected a categorical column at index >= {n_rows}, got {member_cats}" + ) + + seen: list[object] = [] + original_forward = type(clf.models_[0]).forward + + def spy(self: object, x: torch.Tensor, y: torch.Tensor, **kwargs: object) -> object: + seen.append(kwargs.get("categorical_inds")) + return original_forward(self, x, y, **kwargs) # type: ignore[arg-type] + + type(clf.models_[0]).forward = spy # type: ignore[method-assign] + try: + clf.predict_proba(X_test) + single = list(seen) + seen.clear() + clf.predict_proba_batched([X], [y], [X_test]) + batched = list(seen) + finally: + type(clf.models_[0]).forward = original_forward # type: ignore[method-assign] + + assert single, "the spy recorded no forward calls for predict_proba" + assert batched == single diff --git a/tests/test_regressor_interface.py b/tests/test_regressor_interface.py index fc4578f48..02d314587 100644 --- a/tests/test_regressor_interface.py +++ b/tests/test_regressor_interface.py @@ -26,6 +26,7 @@ import tabpfn.regressor as regressor_module from tabpfn import TabPFNRegressor from tabpfn.architectures import tabpfn_v2_5 +from tabpfn.architectures.interface import PerformanceOptions from tabpfn.architectures.shared.bar_distribution import FullSupportBarDistribution from tabpfn.base import ModelSpecs, initialize_tabpfn_model from tabpfn.constants import ModelVersion @@ -39,6 +40,8 @@ ) from tabpfn.model_loading import ModelSource, prepend_cache_path from tabpfn.preprocessing import PreprocessorConfig +from tabpfn.preprocessing.datamodel import FeatureModality +from tabpfn.preprocessing.ensemble import generate_regression_ensemble_configs from tabpfn.settings import settings from tabpfn.utils import infer_devices from tabpfn.validation import ensure_compatible_predict_input_sklearn @@ -2100,3 +2103,139 @@ def test__eval_metric_and_tuning_config__round_trip_through_sklearn_clone( round_tripped = TabPFNRegressor().set_params(**params) assert round_tripped.eval_metric == eval_metric assert round_tripped.tuning_config == tuning_config + + +# ============================================================================= +# fit_from_preprocessed: the feature schema is built from the column count +# ============================================================================= + + +N_BATCH_ROWS = 5 +N_BATCH_COLS = 8 +BATCH_CAT_COL = N_BATCH_COLS - 1 # index 7 >= the 5 rows of a collated batch + + +def _collated_batch( + n_datasets: int = 2, n_estimators: int = 2 +) -> tuple[ + ModelSpecs, + list[torch.Tensor], + list[torch.Tensor], + list[list[list[int]]], + list[list[object]], +]: + """The (n_datasets, n_train_rows, n_columns) tensors the collator produces.""" + specs = _create_dummy_regressor_model_specs() + flat_configs = generate_regression_ensemble_configs( + num_estimators=n_estimators, + add_fingerprint_feature=False, + polynomial_features="no", + feature_shift_decoder="shuffle", + preprocessor_configs=[ + PreprocessorConfig( + "none", categorical_name="numeric", max_features_per_estimator=500 + ) + ], + target_transforms=[None], + random_state=0, + num_models=1, + outlier_removal_std=None, + ) + configs: list[list[object]] = [ + [flat_configs[estimator] for estimator in range(n_estimators)] + for _ in range(n_datasets) + ] + torch.manual_seed(0) + X_preprocessed = [ + torch.randn(n_datasets, N_BATCH_ROWS, N_BATCH_COLS), + ] + y_preprocessed = [ + torch.randn(n_datasets, N_BATCH_ROWS), + ] + cat_ix = [[[BATCH_CAT_COL] for _ in range(n_estimators)] for _ in range(n_datasets)] + return specs, X_preprocessed, y_preprocessed, cat_ix, configs + + +def test__fit_from_preprocessed__feature_schema_spans_the_columns() -> None: + """The schema is as wide as the batch's columns, not its rows. + + X_preprocessed[0] is the collated (n_datasets, n_train_rows, n_columns) + tensor, so `shape[1]` is the row count. A schema built from it is narrower + than the data whenever there are more columns than rows, which silently + drops every categorical column whose index is >= n_train_rows and tells the + model those columns are numerical. + """ + specs, X_preprocessed, y_preprocessed, cat_ix, configs = _collated_batch() + + reg = TabPFNRegressor( + model_path=specs, device="cpu", n_estimators=2, random_state=0 + ) + reg.fit_from_preprocessed( + X_preprocessed, + y_preprocessed, + cat_ix, + configs, + performance_options=PerformanceOptions(), + ) + + assert X_preprocessed[0].shape[1] == N_BATCH_ROWS + assert X_preprocessed[0].shape[-1] == N_BATCH_COLS + for dataset_schemas in reg.executor_.feature_schema_list: + for schema in dataset_schemas: + assert len(schema.features) == N_BATCH_COLS + assert schema.indices_for(FeatureModality.CATEGORICAL) == [BATCH_CAT_COL] + + +def test__predict_batched__tells_the_model_about_high_categorical_columns() -> None: + """Batched inference reports the same categorical columns as the single path. + + A wide dataset pushes the categorical past the row count after + preprocessing (the SVD/fingerprint steps add columns), so a row-wide schema + omits it and the model is told the column is numerical -- while `predict`, + which uses each member's own schema, still reports it. + """ + n_rows = 100 # wide enough to be a realistic regression context + n_cols = 250 # wide enough that the categorical lands beyond n_train_rows + rng = np.random.RandomState(0) + X = rng.randn(n_rows, n_cols).astype(object) + X[:, -1] = rng.choice(["a", "b", "c", "d"], size=n_rows) + y = rng.randn(n_rows) + X_test = np.repeat(X[:1], 2, axis=0).astype(object) + X_test[:, -1] = ["a", "b"] + + reg = TabPFNRegressor( + model_path=_create_dummy_regressor_model_specs(), + device="cpu", + n_estimators=1, + random_state=0, + ) + reg.fit(X, y) + + member_cats = reg.executor_.ensemble_members[0].feature_schema.indices_for( + FeatureModality.CATEGORICAL + ) + # Precondition: this is the case the bug needs, so the test cannot pass + # vacuously on a schema that happened to keep the column. + assert [c for c in member_cats if c >= n_rows], ( + f"expected a categorical column at index >= {n_rows}, got {member_cats}" + ) + + seen: list[object] = [] + original_forward = type(reg.models_[0]).forward + + def spy(self: object, x: torch.Tensor, y: torch.Tensor, **kwargs: object) -> object: + seen.append(kwargs.get("categorical_inds")) + return original_forward(self, x, y, **kwargs) # type: ignore[arg-type] + + type(reg.models_[0]).forward = spy # type: ignore[method-assign] + try: + reg.predict(X_test) + single = list(seen) + seen.clear() + reg.predict_batched([X], [y], [X_test]) + batched = list(seen) + finally: + type(reg.models_[0]).forward = original_forward # type: ignore[method-assign] + + assert single, "the spy recorded no forward calls for predict" + assert batched == single