Merge pull request #3519 from pycaret/fix_predict_model_transformed_m…

…etrics
pycaret · May 1, 2023 · eba5ee9 · eba5ee9
2 parents 8608d5e + b1f9008
commit eba5ee9
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/pycaret/internal/pycaret_experiment/supervised_experiment.py b/pycaret/internal/pycaret_experiment/supervised_experiment.py
@@ -4981,16 +4981,21 @@ def replace_labels_in_column(label_encoder, labels: pd.Series) -> pd.Series:
             pred = pred.values
 
         try:
+            # This is a classifier
             score = estimator.predict_proba(X_test_)
 
             if len(np.unique(pred)) <= 2:
                 pred_prob = score[:, 1]
             else:
                 pred_prob = score
 
+            y_test_metrics = y_test_
+
         except Exception:
+            # This is not a classifier
             score = None
             pred_prob = None
+            y_test_metrics = y_test_untransformed
 
         if probability_threshold is not None and pred_prob is not None:
             try:
@@ -5005,7 +5010,7 @@ def replace_labels_in_column(label_encoder, labels: pd.Series) -> pd.Series:
         if y_test_ is not None and self._setup_ran:
             # model name
             full_name = self._get_model_name(estimator)
-            metrics = self._calculate_metrics(y_test_, pred, pred_prob)  # type: ignore
+            metrics = self._calculate_metrics(y_test_metrics, pred, pred_prob)  # type: ignore
             df_score = pd.DataFrame(metrics, index=[0])
             df_score.insert(0, "Model", full_name)
             df_score = df_score.round(round)

diff --git a/tests/test_supervised_predict_model.py b/tests/test_supervised_predict_model.py
@@ -1,3 +1,4 @@
+import numpy as np
 import pandas as pd
 
 import pycaret.classification
@@ -51,5 +52,22 @@ def test_regression_predict_model():
     )
     lr_model = pycaret.regression.create_model("lr")
     predictions = pycaret.regression.predict_model(lr_model, data=unseen_data)
+    metrics = pycaret.regression.pull()
     # Check that columns of raw data are contained in columns of returned dataframe
     assert all(item in predictions.columns for item in unseen_data.columns)
+
+    pycaret.regression.setup(
+        data,
+        target="medv",
+        ignore_features=["crim", "zn"],
+        remove_multicollinearity=True,
+        multicollinearity_threshold=0.95,
+        transform_target=True,
+        html=False,
+        session_id=123,
+        n_jobs=1,
+    )
+    lr_model = pycaret.regression.create_model("lr")
+    pycaret.regression.predict_model(lr_model, data=unseen_data)
+    metrics_transformed = pycaret.regression.pull()
+    assert np.isclose(metrics["R2"], metrics_transformed["R2"], atol=0.15)