[Python] TabPFN feature selection raises KeyError(f"None of [{key}] are in the [{axis_name}]")

I trained a tabPFN model, which I then tried applying a sequential feature selector for important feature selection. I've been getting this error

KeyError(f"None of [{key}] are in the [{axis_name}]")


on the following code:

Offending code:

from tabpfn_extensions import interpretability
from sklearn.feature_selection import SequentialFeatureSelector

try:
sfs = interpretability.feature_selection.feature_selection(
estimator=PFN,
X=X_train,
y=y_train,
n_features_to_select=5,
feature_names=category_column_indexes,
error_score='raise'
)
except KeyError as e:
print("KeyError:", e)
print("X_train.columns:", X_train.columns)
print("category_column_indexes:", category_column_indexes)

selected_features = [X_train.columns for i, selected in enumerate(sfs.get_support()) if selected]

print("\nSelected features:")
for feature in selected_features:
print(f"- {feature}")


Full error:

ValueError:
All the 5 fits failed.
It is very likely that your model is misconfigured.
You can try to debug the error by setting error_score='raise'.

Below are more details about the failures:
--------------------------------------------------------------------------------
5 fits failed with the following error:
Traceback (most recent call last):
File "/python3.12/site-packages/sklearn/model_selection/_validation.py", line 888, in _fit_and_score
estimator.fit(X_train, y_train, **fit_params)
File "/TabPFN/src/tabpfn/classifier.py", line 453, in fit
X = _fix_dtypes(X, cat_indices=self.categorical_features_indices)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/TabPFN/src/tabpfn/utils.py", line 480, in _fix_dtypes
X[cat_indices] = X[cat_indices].astype("category")
~^^^^^^^^^^^^^
File "/python3.12/site-packages/pandas/core/frame.py", line 4108, in __getitem__
indexer = self.columns._get_indexer_strict(key, "columns")[1]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/python3.12/site-packages/pandas/core/indexes/base.py", line 6200, in _get_indexer_strict
self._raise_if_missing(keyarr, indexer, axis_name)
File "/python3.12/site-packages/pandas/core/indexes/base.py", line 6249, in _raise_if_missing
raise KeyError(f"None of [{key}] are in the [{axis_name}]")
KeyError: "None of [Index([ 1, 3, 5, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21,\n 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39,\n 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57,\n 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75,\n 76, 77, 78, 79, 80, 81, 82, 83, 84, 85],\n dtype='int64')] are in the [columns]"


For reference my data frame consists of four numerical columns (index 0, 2, 4, 6) and the remainder are categorical data. Here is how I initialised the classifier.

PFN = TabPFNClassifier(
random_state=42,
categorical_features_indices=category_column_indexes,
softmax_temperature=0.5,
n_estimators=94,
balance_probabilities=False,
average_before_softmax=True,
inference_config=ModelInterfaceConfig(
OUTLIER_REMOVAL_STD=1000,
REGRESSION_Y_PREPROCESS_TRANSFORMS=(None,),
FINGERPRINT_FEATURE=False,
PREPROCESS_TRANSFORMS=(PreprocessorConfig("none",),)
)
)


I tried commenting out feature_names=category_column_indexes , I also tried inputting the category_column_names , I also tried training the TabPFN with categorical columns only, but get the same error each time.

I tried inserting print statements in the classifier source code itself and this is what I got

self.categorical_features_indices: [1, 3, 5, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85]
use_iloc: False
is_numeric_indices: True
columns_are_numeric: True
X.columns: RangeIndex(start=0, stop=86, step=1)
cat_indices: [1, 3, 5, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85]

Continue reading...