Merge pull request #762 from DashAISoftware/fix/predict-cols

Predict save the cols of the full dataset

Author: cristian-tamblay

DashAI/back/job/predict_job.py

@@ -422,7 +422,7 @@ def run(
                         manual_input_data, dataset_trained_path
                     )
 
-                prepared_dataset, y_pred = _run_prediction_pipeline(
+                _, y_pred = _run_prediction_pipeline(
                     task=task,
                     trained_model=trained_model,
                     train_dataset=train_dataset,
@@ -460,9 +460,13 @@ def run(
                 full_path = Path(path) / folder_name
                 full_path.mkdir(parents=True, exist_ok=True)
 
-                # Add predictions to loaded dataset
+                output_col = model_session.output_columns[0]
+                base_columns = [
+                    col for col in loaded_dataset.column_names if col != output_col
+                ]
+                output_dataset = loaded_dataset.select_columns(base_columns)
                 dataset_with_prediction = to_dashai_dataset(
-                    prepared_dataset.add_column(model_session.output_columns[0], y_pred)
+                    output_dataset.add_column(output_col, y_pred)
                 )
 
                 # Filter schema from trained dataset