From 0965baf78588e29d813b61d73f0ebd868a0ce3d0 Mon Sep 17 00:00:00 2001
From: DseidLi <2568818204@qq.com>
Date: Mon, 26 Feb 2024 13:31:26 +0800
Subject: [PATCH] fix bug in needlebench summarizer

---
 opencompass/summarizers/needlebench.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/opencompass/summarizers/needlebench.py b/opencompass/summarizers/needlebench.py
index cf4a897d..27d89be7 100644
--- a/opencompass/summarizers/needlebench.py
+++ b/opencompass/summarizers/needlebench.py
@@ -278,6 +278,12 @@ def merge_dataframes(model_name, dataset_abbrs, parsed_data):
 
     # 沿着列方向合并DataFrame
     merged_df = pd.concat(dfs, axis=1)
+
+    # Check for NaN values and filter out rows with NaN
+    if merged_df.isnull().any().any():
+        print('Warning: Some rows were filtered out due to NaN values. This is often due to mismatched row counts among DataFrames.')
+        merged_df = merged_df.dropna()
+
     return merged_df
 
 def calculate_elementwise_average(merged_df):