FIX: add AutoPipeline predict_proba

lale/lib/lale/auto_pipeline.py

@@ -280,6 +280,11 @@ def predict(self, X, **predict_params):
         result = best_pipeline.predict(X, **predict_params)
         return result
 
+    def predict_proba(self, X, **predict_proba_params):
+        best_pipeline = self._pipelines[self._name_of_best]
+        result = best_pipeline.predict_proba(X, **predict_proba_params)
+        return result
+
     def summary(self):
         """Table summarizing the trial results (name, tid, loss, time, log_loss, status).
         Returns
@@ -394,6 +399,8 @@ def get_pipeline(
     },
 }
 
+_input_predict_proba_schema = _input_predict_schema
+
 _output_predict_schema = {
     "anyOf": [
         {"type": "array", "items": {"type": "number"}},
@@ -402,6 +409,11 @@ def get_pipeline(
     ]
 }
 
+_output_predict_proba_schema = {
+    "type": "array",
+    "items": {"type": "array", "items": {"type": "number"}},
+}
+
 _combined_schemas = {
     "description": """Automatically find a pipeline for a dataset.
 
@@ -421,6 +433,8 @@ def get_pipeline(
         "input_fit": _input_fit_schema,
         "input_predict": _input_predict_schema,
         "output_predict": _output_predict_schema,
+        "input_predict_proba": _input_predict_proba_schema,
+        "output_predict_proba": _output_predict_proba_schema,
     },
 }
 

test/test_core_pipeline.py

@@ -809,6 +809,28 @@ def test_sklearn_iris(self):
         all_X, all_y = sklearn.datasets.load_iris(return_X_y=True)
         self._fit_predict("classification", all_X, all_y)
 
+    def test_predict_proba_with_roc_auc(self):
+        from lale.lib.lale import AutoPipeline
+
+        all_X, all_y = sklearn.datasets.load_breast_cancer(return_X_y=True)
+        train_X, test_X, train_y, _test_y = train_test_split(
+            all_X, all_y, test_size=0.2, random_state=42
+        )
+
+        trainable = AutoPipeline(
+            prediction_type="classification",
+            scoring="roc_auc",
+            max_evals=5,
+            max_opt_time=60,
+            max_eval_time=30,
+            cv=3,
+        )
+        trained = trainable.fit(train_X, train_y)
+        predicted_proba = trained.predict_proba(test_X)
+
+        self.assertEqual(predicted_proba.shape[0], test_X.shape[0])
+        self.assertEqual(predicted_proba.shape[1], 2)
+
     def test_sklearn_digits(self):
         # classification, numbers but some appear categorical, no missing values
         all_X, all_y = sklearn.datasets.load_digits(return_X_y=True)