diff --git a/lale/lib/lale/auto_pipeline.py b/lale/lib/lale/auto_pipeline.py index ce8a29c63..bdfcd2532 100644 --- a/lale/lib/lale/auto_pipeline.py +++ b/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, }, } diff --git a/test/test_core_pipeline.py b/test/test_core_pipeline.py index 8ea7eb560..e620ccc30 100644 --- a/test/test_core_pipeline.py +++ b/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)