From 641beb22033fca9f044a822fb5769bc833fa4ce1 Mon Sep 17 00:00:00 2001
From: XianBW <36835909+XianBW@users.noreply.github.com>
Date: Fri, 14 Mar 2025 17:40:06 +0800
Subject: [PATCH] check 'ensemble' in scores.csv in workflow eval (#682)

---
 .../data_science/ensemble/eval_tests/ensemble_test.txt   | 2 +-
 rdagent/components/coder/data_science/workflow/eval.py   | 9 ++++-----
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/rdagent/components/coder/data_science/ensemble/eval_tests/ensemble_test.txt b/rdagent/components/coder/data_science/ensemble/eval_tests/ensemble_test.txt
index de85f6d5..972f5395 100644
--- a/rdagent/components/coder/data_science/ensemble/eval_tests/ensemble_test.txt
+++ b/rdagent/components/coder/data_science/ensemble/eval_tests/ensemble_test.txt
@@ -22,7 +22,7 @@ def print_preds_info(model_name, data_type, preds):
     else:
         print(f"Model {model_name} {data_type} predictions shape: {preds.shape}")
 
-        if isinstance(preds, pd.DataFrame):
+        if isinstance(preds, (pd.DataFrame, pd.Series)):
             print(preds.head())
         elif isinstance(preds, (np.ndarray, torch.Tensor, tf.Tensor)):
             print(pd.DataFrame(preds).head())
diff --git a/rdagent/components/coder/data_science/workflow/eval.py b/rdagent/components/coder/data_science/workflow/eval.py
index eb830b7f..9df1a07c 100644
--- a/rdagent/components/coder/data_science/workflow/eval.py
+++ b/rdagent/components/coder/data_science/workflow/eval.py
@@ -84,12 +84,11 @@ class WorkflowGeneralCaseSpecEvaluator(CoSTEEREvaluator):
                 model_set_in_folder = set(
                     f[:-3] for f in implementation.file_dict.keys() if re.match(r"^model_(?!test)\w+\.py$", f)
                 )
-                for model in model_set_in_folder:
-                    if model not in model_set_in_scores:
-                        score_check_text += f"\nModel {model} is not evaluated in the scores.csv. The scores.csv has {model_set_in_scores}."
-                        score_ret_code = 1
+                if model_set_in_scores != model_set_in_folder.union({"ensemble"}):
+                    score_check_text += f"\n[Error] The scores dataframe does not contain the correct model names as index.\ncorrect model names are: {model_set_in_folder.union({'ensemble'})}\nscore_df is:\n{score_df}"
+                    score_ret_code = 1
             except Exception as e:
-                score_check_text += f"\nError in checking the scores.csv file: {e}\nscores.csv's content:\n-----\n{score_fp.read_text()}\n-----"
+                score_check_text += f"\n[Error] in checking the scores.csv file: {e}\nscores.csv's content:\n-----\n{score_fp.read_text()}\n-----"
                 score_ret_code = 1
 
         # Check submission file