Function test_balanced_probabilities

tests/test_classifier_interface.py:693–733

Test that balance_probabilities=True works correctly.

()

@pytest.mark.skip(reason="This test is flaky and needs to be fixed.")
def test_balanced_probabilities() -> None:
    """Test that balance_probabilities=True works correctly."""
    n_classes = 3
    n_features = 3

    # Create an IMBALANCED dataset
    X, y = sklearn.datasets.make_classification(
        n_samples=60,
        n_classes=n_classes,
        n_features=n_features,
        n_informative=n_features,
        n_redundant=0,
        weights=[0.8, 0.1, 0.1],
        random_state=42,
    )

    model_unbalanced = TabPFNClassifier(
        balance_probabilities=False,
        random_state=42,
        n_estimators=2,
    )
    model_unbalanced.fit(X, y)
    proba_unbalanced = model_unbalanced.predict_proba(X)

    model_balanced = TabPFNClassifier(
        balance_probabilities=True,
        random_state=42,
        n_estimators=2,
    )
    model_balanced.fit(X, y)
    proba_balanced = model_balanced.predict_proba(X)

    mean_proba_unbalanced = proba_unbalanced.mean(axis=0)
    mean_proba_balanced = proba_balanced.mean(axis=0)

    # Balanced should be MORE uniform than unbalanced
    balanced_deviation = np.std(mean_proba_balanced)
    unbalanced_deviation = np.std(mean_proba_unbalanced)
    assert balanced_deviation < unbalanced_deviation, (
        "Balancing did not make probabilities more uniform"
    )


def test_classifier_in_pipeline(X_y: tuple[np.ndarray, np.ndarray]) -> None: