From 11e8acc88bcb0dab3ff72f5be1ac1ef2f20143d9 Mon Sep 17 00:00:00 2001
From: Yuante Li <104308117+WinstonLiyt@users.noreply.github.com>
Date: Fri, 4 Apr 2025 22:44:55 +0800
Subject: [PATCH] feat: add a check for whether values in score_df are NaN
 (#756)

* add a check for whether values in score_df are NaN

* fix ci

* change raise to assert
---
 .../data_science/ensemble/eval_tests/ensemble_test.txt      | 5 +++++
 rdagent/components/coder/data_science/pipeline/eval.py      | 6 ++++++
 rdagent/components/coder/data_science/workflow/eval.py      | 6 ++++++
 3 files changed, 17 insertions(+)

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 5b6237dd..ee48452e 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
@@ -128,5 +128,10 @@ assert model_set_in_scores == set({{model_names}}).union({"ensemble"}), (
 assert score_df.index.is_unique, "The scores dataframe has duplicate model names."
 assert score_df.columns.tolist() == ["{{metric_name}}"], f"The column names of the scores dataframe should be ['{{metric_name}}'], but is '{score_df.columns.tolist()}'"
 
+# Check for NaN values in score_df
+assert not score_df.isnull().values.any(), (
+    f"The scores dataframe contains NaN values at the following locations:\n{score_df[score_df.isnull().any(axis=1)]}"
+)
+
 
 print("Ensemble test end.")
diff --git a/rdagent/components/coder/data_science/pipeline/eval.py b/rdagent/components/coder/data_science/pipeline/eval.py
index 184c4991..9521fd10 100644
--- a/rdagent/components/coder/data_science/pipeline/eval.py
+++ b/rdagent/components/coder/data_science/pipeline/eval.py
@@ -82,6 +82,12 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator):
                     score_check_text += f"\n[Error] The scores dataframe does not contain the correct column names.\nCorrect columns is: ['{self.scen.metric_name}']\nBut got: {score_df.columns.tolist()}"
                     score_ret_code = 1
 
+                # Check if scores contain NaN (values)
+                if score_df.isnull().values.any():
+                    nan_locations = score_df[score_df.isnull().any(axis=1)]
+                    score_check_text += f"\n[Error] The scores dataframe contains NaN values at the following locations:\n{nan_locations}"
+                    score_ret_code = 1
+
             except Exception as e:
                 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
diff --git a/rdagent/components/coder/data_science/workflow/eval.py b/rdagent/components/coder/data_science/workflow/eval.py
index 5313700a..b7675176 100644
--- a/rdagent/components/coder/data_science/workflow/eval.py
+++ b/rdagent/components/coder/data_science/workflow/eval.py
@@ -106,6 +106,12 @@ class WorkflowGeneralCaseSpecEvaluator(CoSTEEREvaluator):
                     score_check_text += f"\n[Error] The scores dataframe does not contain the correct column names.\nCorrect columns is: ['{self.scen.metric_name}']\nBut got: {score_df.columns.tolist()}"
                     score_ret_code = 1
 
+                # Check if scores contain NaN (values)
+                if score_df.isnull().values.any():
+                    nan_locations = score_df[score_df.isnull().any(axis=1)]
+                    score_check_text += f"\n[Error] The scores dataframe contains NaN values at the following locations:\n{nan_locations}"
+                    score_ret_code = 1
+
             except Exception as e:
                 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