Add predict proba to classifier

alteryx · Jan 31, 2023 · 2ae5e19 · 2ae5e19
1 parent d33ba2a
commit 2ae5e19
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/evalml/pipelines/components/estimators/classifiers/catboost_classifier.py b/evalml/pipelines/components/estimators/classifiers/catboost_classifier.py
@@ -132,7 +132,7 @@ def predict(self, X):
             X (pd.DataFrame): Data of shape [n_samples, n_features].
 
         Returns:
-            pd.DataFrame: Predicted values.
+            pd.Series: Predicted values.
         """
         X = infer_feature_types(X)
         X = handle_float_categories_for_catboost(X)
@@ -147,6 +147,20 @@ def predict(self, X):
         predictions.index = X.index
         return predictions
 
+    def predict_proba(self, X):
+        """Make prediction probabilities using the fitted CatBoost classifier.
+
+        Args:
+            X (pd.DataFrame): Data of shape [n_samples, n_features].
+
+        Returns:
+            pd.DataFrame: Predicted probability values.
+        """
+        X = infer_feature_types(X)
+        X = handle_float_categories_for_catboost(X)
+        predictions = super().predict_proba(X)
+        return predictions
+
     @property
     def feature_importance(self):
         """Feature importance of fitted CatBoost classifier."""

diff --git a/evalml/tests/component_tests/test_catboost_classifier.py b/evalml/tests/component_tests/test_catboost_classifier.py
@@ -61,3 +61,5 @@ def test_catboost_classifier_double_categories_in_X(categorical_floats_df):
     assert isinstance(fitted, CatBoostClassifier)
     predictions = clf.predict(X)
     assert isinstance(predictions, pd.Series)
+    predictions = clf.predict_proba(X)
+    assert isinstance(predictions, pd.DataFrame)