Class TestClassifierInputFormats

tests/quick_test_v2.py:155–199

Test classification with different input data formats.

class TestClassifierInputFormats:
    """Test classification with different input data formats."""

    def test_pandas_input(self, clf_data):
        X_train, X_test, y_train, _ = clf_data
        X_train_df = pd.DataFrame(X_train)
        X_test_df = pd.DataFrame(X_test)
        clf = TabPFNClassifier.create_default_for_version(
            ModelVersion.V2_5, n_estimators=3
        )
        preds, _ = _fit_and_predict_classifier(clf, X_train_df, y_train, X_test_df)
        assert preds.shape == (len(X_test),)

    def test_small_train_set(self, clf_data):
        """Fit with very few training samples."""
        X_train, X_test, y_train, _ = clf_data
        clf = TabPFNClassifier.create_default_for_version(
            ModelVersion.V2_5, n_estimators=3
        )
        preds, _ = _fit_and_predict_classifier(clf, X_train[:10], y_train[:10], X_test)
        assert preds.shape == (len(X_test),)

    def test_single_test_sample(self, clf_data):
        """Predict on a single test sample."""
        X_train, X_test, y_train, _ = clf_data
        clf = TabPFNClassifier.create_default_for_version(
            ModelVersion.V2_5, n_estimators=3
        )
        clf.fit(X_train, y_train)
        preds = clf.predict(X_test[:1])
        probas = clf.predict_proba(X_test[:1])
        assert preds.shape == (1,)
        assert probas.shape == (1, 2)

    def test_few_features(self):
        """Classification with a small number of features."""
        rng = np.random.RandomState(42)
        X_train = rng.randn(50, 2)
        y_train = (X_train[:, 0] > 0).astype(int)
        X_test = rng.randn(10, 2)
        clf = TabPFNClassifier.create_default_for_version(
            ModelVersion.V2_5, n_estimators=3
        )
        preds, _ = _fit_and_predict_classifier(clf, X_train, y_train, X_test)
        assert preds.shape == (10,)


# ---------------------------------------------------------------------------