From 77ddee64ce73bf0af0608450ba4bae6375764b3c Mon Sep 17 00:00:00 2001
From: XianBW <36835909+XianBW@users.noreply.github.com>
Date: Wed, 19 Feb 2025 21:26:09 +0800
Subject: [PATCH] fix: change ensemble test (#622)

* change ensemble test

* remove redundant info

* fix bug

* add model name duplication check in ensemble test

* change text
---
 .../ensemble/eval_tests/ensemble_test.txt          | 14 ++++----------
 1 file changed, 4 insertions(+), 10 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 1d70116e..296c7ed1 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
@@ -90,16 +90,10 @@ if isinstance(final_pred, (list, np.ndarray, pd.DataFrame, torch.Tensor, tf.Tens
 assert Path("scores.csv").exists(), "scores.csv is not generated"
 score_df = pd.read_csv("scores.csv", index_col=0)
 model_set_in_scores = set(score_df.index)
-for model in {{model_names}}:
-    if model not in model_set_in_scores:
-        print(f"\nModel {model} is not evaluated in the scores.csv.")
 
-print("Please check the format of the scores dataframe:")
-print(score_df)
-
-if "ensemble" not in model_set_in_scores:
-    print("\nError: Ensemble results are missing or incorrectly saved in `scores.csv`. "
-          "Ensure that the ensemble's performance is recorded as 'ensemble' in the file.")
+assert model_set_in_scores == set({{model_names}}).union({"ensemble"}), (
+    f"The scores dataframe does not contain the correct model names as index.\ncorrect model names are: {{model_names}} + ['ensemble']\nscore_df is:\n{score_df}"
+)
+assert score_df.index.is_unique, "The scores dataframe has duplicate model names."
 
 print("Ensemble test end.")
-print(f"Final prediction shape: {final_pred.shape}")