diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 0672242..ba6a346 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -528,7 +528,20 @@ class MLFlow(SientiaMonitoring): ) # Keep heavy model fitting outside MLflow run timing. - wrapper.retrain(data) + prediction_data = wrapper.retrain(data) + + # Merge prediction data with retrain data + evaluation_data = pd.merge( + data, prediction_data, left_index=True, right_index=True, how='left' + ) + + # Rename target column to "target" + evaluation_data.rename(columns={target: 'target'}, inplace=True) + + # Reset index and put as column "timestamp" + evaluation_data['timestamp'] = evaluation_data.index + evaluation_data.reset_index(drop=True, inplace=True) + evaluation_data.sort_values(by='timestamp', inplace=True, ascending=True) run_name = f'{model_name}-retrain-{datetime.now().strftime("%Y%m%d%H%M%S")}' @@ -542,8 +555,12 @@ class MLFlow(SientiaMonitoring): tmp_dir = tempfile.mkdtemp(prefix='laborious_retrain_') try: raw_csv = Path(tmp_dir) / 'retrain_input.csv' + evaluation_csv = Path(tmp_dir) / 'evaluation_data.csv' data.to_csv(raw_csv, index=False) + evaluation_data.to_csv(evaluation_csv, index=False) + mlflow.log_artifact(str(raw_csv)) + mlflow.log_artifact(str(evaluation_csv)) finally: rmtree(tmp_dir, ignore_errors=True)