From 0624ed180f89faca9be9f0a2821b130f5513b1c4 Mon Sep 17 00:00:00 2001
From: Yuante Li <104308117+WinstonLiyt@users.noreply.github.com>
Date: Thu, 20 Mar 2025 17:35:39 +0800
Subject: [PATCH] fix: fix a bug in ensemble test script (#713)

* fix a bug in ensemble test script

* fix

* add a comment
---
 .../coder/data_science/ensemble/eval_tests/ensemble_test.txt   | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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 972f5395..66d715b4 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,10 +22,11 @@ def print_preds_info(model_name, data_type, preds):
     else:
         print(f"Model {model_name} {data_type} predictions shape: {preds.shape}")
 
+        print("Showing a preview of the predictions (first few entries only):")
         if isinstance(preds, (pd.DataFrame, pd.Series)):
             print(preds.head())
         elif isinstance(preds, (np.ndarray, torch.Tensor, tf.Tensor)):
-            print(pd.DataFrame(preds).head())
+            print(preds[:2])
         elif isinstance(preds, list):
             print(pd.DataFrame(preds[:5]))
         else: