Support DataFrame model outputs for predict and predict_proba (#334)

Author: oegedijk

RELEASE_NOTES.md

@@ -5,6 +5,8 @@
 
 ### Bug Fixes
 - Allow FeatureInputComponent (what-if inputs) to customize numeric ranges and rounding, and apply min/max/step to inputs.
+- Improve compatibility with AutoGluon/custom wrappers by coercing pandas `DataFrame` outputs from `predict_proba`/`predict` to numpy arrays before indexing in classifier/regression helper paths.
+- Harden one-vs-all scorer handling so `make_one_vs_all_scorer` also accepts classifiers whose `predict_proba` returns a pandas `DataFrame`.
 
 ## Version 0.5.6:
 

TODO.md

@@ -9,7 +9,6 @@
 - [M][Explainers][#273] categorical columns with NaNs: sorting and column preservation.
 
 **Next**
-- [S][Explainers][#270] Autogluon integration (coerce predict_proba to ndarray).
 - [M][Hub][#269] add_dashboard endpoint fails after first request (Flask blueprint lifecycle).
 - [M/L][Components][#262] add filters for random transaction selection in whatif tab.
 - [S][Methods][#220] get_contrib_df accepts list/array input.

explainerdashboard/explainer_methods.py

@@ -829,6 +829,8 @@ def _scorer(clf, X, y):
         warnings.filterwarnings("ignore", category=UserWarning)
         y_pred = clf.predict_proba(X)
         warnings.filterwarnings("default", category=UserWarning)
+        y_pred = _ensure_numeric_predictions(y_pred)
+        y_pred = np.asarray(y_pred)
         score = sign * partial_metric(y, y_pred)
         return score
 

explainerdashboard/explainers.py

@@ -914,15 +914,19 @@ def get_col_value_plus_prediction(
             if self.is_classifier:
                 if pos_label is None:
                     pos_label = self.pos_label
-                pred_probas_raw = self.model.predict_proba(model_input)[0]
+                pred_probas_raw = self.model.predict_proba(model_input)
                 pred_probas_raw = _ensure_numeric_predictions(pred_probas_raw)
-                prediction = np.asarray(pred_probas_raw)[pos_label].squeeze()
+                pred_probas = np.asarray(pred_probas_raw).squeeze()
+                if pred_probas.ndim > 1:
+                    pred_probas = pred_probas[0]
+                prediction = pred_probas[pos_label].squeeze()
                 if self.model_output == "probability":
                     prediction = 100 * prediction
             elif self.is_regression:
-                pred_raw = self.model.predict(model_input)[0]
+                pred_raw = self.model.predict(model_input)
                 pred_raw = _ensure_numeric_predictions(pred_raw)
-                prediction = np.asarray(pred_raw).squeeze()
+                pred_array = np.asarray(pred_raw).squeeze()
+                prediction = pred_array.flat[0]
             return col_value, prediction
         else:
             raise ValueError("You need to pass either index or X_row!")
@@ -3793,9 +3797,11 @@ def prediction_result_df(
                 X_row = X_cats_to_X(X_row, self.onehot_dict, self.X.columns)
             if self.shap == "skorch":
                 X_row = X_row.values.astype("float32")
-            pred_probas_raw = self.model.predict_proba(X_row)[0, :]
+            pred_probas_raw = self.model.predict_proba(X_row)
             pred_probas_raw = _ensure_numeric_predictions(pred_probas_raw)
             pred_probas = np.asarray(pred_probas_raw).squeeze()
+            if pred_probas.ndim > 1:
+                pred_probas = pred_probas[0]
 
         preds_df = pd.DataFrame(dict(label=self.labels, probability=pred_probas))
         if logodds and all(preds_df.probability < 1 - np.finfo(np.float64).eps):

tests/test_classifier_base.py

@@ -14,6 +14,21 @@
 from explainerdashboard.explainer_methods import IndexNotFoundError
 
 
+class DataFramePredictProbaWrapper:
+    def __init__(self, model):
+        self.model = model
+        self.classes_ = model.classes_
+
+    def predict(self, X):
+        return self.model.predict(X)
+
+    def predict_proba(self, X):
+        probas = self.model.predict_proba(X)
+        return pd.DataFrame(
+            probas, columns=self.classes_, index=getattr(X, "index", None)
+        )
+
+
 def test_explainer_with_dataframe_y(fitted_rf_classifier_model, classifier_data):
     _, _, X_test, y_test = classifier_data
     explainer = ClassifierExplainer(
@@ -332,6 +347,20 @@ def test_prediction_result_df(precalculated_rf_classifier_explainer):
     assert isinstance(df, pd.DataFrame)
 
 
+def test_prediction_result_df_with_dataframe_predict_proba(
+    fitted_rf_classifier_model, classifier_data
+):
+    _, _, X_test, y_test = classifier_data
+    wrapped_model = DataFramePredictProbaWrapper(fitted_rf_classifier_model)
+    explainer = ClassifierExplainer(wrapped_model, X_test.head(50), y_test.head(50))
+
+    df = explainer.prediction_result_df(0)
+    _, prediction = explainer.get_col_value_plus_prediction("Age", index=0)
+
+    assert isinstance(df, pd.DataFrame)
+    assert np.isscalar(prediction)
+
+
 def test_pdp_df(precalculated_rf_classifier_explainer):
     assert isinstance(precalculated_rf_classifier_explainer.pdp_df("Age"), pd.DataFrame)
     assert isinstance(

tests/test_dtype_alignment.py

@@ -2,11 +2,13 @@
 import pandas as pd
 
 from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin
+from sklearn.metrics import roc_auc_score
 from sklearn.metrics import r2_score
 
 from explainerdashboard.explainer_methods import (
     align_categorical_dtypes,
     get_pdp_df,
+    make_one_vs_all_scorer,
     permutation_importances,
 )
 
@@ -37,6 +39,12 @@ def predict_proba(self, X):
         return np.tile(np.array([0.2, 0.8]), (len(X), 1))
 
 
+class DataFrameProbaClassifier(ClassifierMixin, BaseEstimator):
+    def predict_proba(self, X):
+        probs = np.tile(np.array([0.2, 0.8]), (len(X), 1))
+        return pd.DataFrame(probs, index=getattr(X, "index", None), columns=[0, 1])
+
+
 def test_permutation_importances_preserves_categorical_dtypes():
     X = pd.DataFrame(
         {
@@ -80,3 +88,14 @@ def test_align_categorical_dtypes_matches_reference():
     aligned = align_categorical_dtypes(target, reference)
 
     assert aligned["cat"].dtype == reference["cat"].dtype
+
+
+def test_make_one_vs_all_scorer_accepts_dataframe_predict_proba():
+    scorer = make_one_vs_all_scorer(roc_auc_score, pos_label=1)
+    clf = DataFrameProbaClassifier()
+    X = pd.DataFrame({"feature": [0, 1, 2, 3]})
+    y = np.array([0, 1, 0, 1])
+
+    score = scorer(clf, X, y)
+
+    assert isinstance(score, float)

tests/test_regression_base.py

@@ -4,6 +4,17 @@
 
 import plotly.graph_objects as go
 
+from explainerdashboard import RegressionExplainer
+
+
+class DataFramePredictWrapper:
+    def __init__(self, model):
+        self.model = model
+
+    def predict(self, X):
+        preds = self.model.predict(X)
+        return pd.DataFrame({"prediction": preds}, index=getattr(X, "index", None))
+
 
 def test_explainer_len(precalculated_rf_regression_explainer, testlen):
     assert len(precalculated_rf_regression_explainer) == testlen
@@ -55,6 +66,18 @@ def test_prediction_result_df(precalculated_rf_regression_explainer):
     assert isinstance(df, pd.DataFrame)
 
 
+def test_get_col_value_plus_prediction_with_dataframe_predict(
+    fitted_rf_regression_model, regression_data
+):
+    _, _, X_test, y_test = regression_data
+    wrapped_model = DataFramePredictWrapper(fitted_rf_regression_model)
+    explainer = RegressionExplainer(wrapped_model, X_test.head(50), y_test.head(50))
+
+    _, prediction = explainer.get_col_value_plus_prediction("Age", index=0)
+
+    assert np.isscalar(prediction)
+
+
 def test_preds(precalculated_rf_regression_explainer):
     assert isinstance(precalculated_rf_regression_explainer.preds, np.ndarray)