Merge pull request #855 from DashAISoftware/fix/regression-prediction-crashes

Fix regression prediction crashes

Author: cristian-tamblay

DashAI/back/job/predict_job.py

@@ -12,6 +12,7 @@
 from DashAI.back.job.base_job import BaseJob, JobError
 from DashAI.back.models.base_model import BaseModel
 from DashAI.back.tasks.base_task import BaseTask
+from DashAI.back.tasks.regression_task import RegressionTask
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import sessionmaker
@@ -482,6 +483,14 @@ def run(
                     if key in model_session.input_columns + model_session.output_columns
                 }
 
+                # Regression models predict continuous values regardless of
+                # the training target's original dtype (e.g. a target column
+                # that happened to hold only integer-looking values), so the
+                # output column's saved schema must reflect that instead of
+                # inheriting the training dataset's type.
+                if isinstance(task, RegressionTask):
+                    filtered_schema[output_col] = {"type": "Float", "dtype": "float64"}
+
                 # Store num of rows, columns, and column names
                 dataset_with_prediction.compute_base_metadata()
 

DashAI/back/models/scikit_learn/sklearn_like_model.py

@@ -79,6 +79,12 @@ def train(self, x_train, y_train, x_validation=None, y_validation=None):
         """
         x_processed = self.prepare_dataset(x_train, is_fit=True).to_pandas()
         y_processed = self.prepare_output(y_train, is_fit=True).to_pandas()
+        # Every task using this base class has outputs_cardinality 1, so this
+        # is always a single column. Passed as a DataFrame, some estimators
+        # (e.g. LinearRegression) keep predictions 2D to match; squeezing to a
+        # Series here keeps fit/predict shapes 1D for every estimator alike.
+        if y_processed.shape[1] == 1:
+            y_processed = y_processed.iloc[:, 0]
         return super().fit(x_processed, y_processed)
 
     def predict(self, x: "DashAIDataset"):