From 673d75509dae235c5154f305dfee019bb88ac6ad Mon Sep 17 00:00:00 2001
From: svlandeg <svlandeg@github.com>
Date: Fri, 16 Jan 2026 12:27:50 +0100
Subject: [PATCH] avoid keeping less data than the necessary examples for
 few-shot

---
 scripts/base_eval.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/scripts/base_eval.py b/scripts/base_eval.py
index bd83ff3..f3d604c 100644
--- a/scripts/base_eval.py
+++ b/scripts/base_eval.py
@@ -96,7 +96,8 @@ def evaluate_model(model, tokenizer, device, max_per_task=-1):
         shuffle_rng = random.Random(1337)
         shuffle_rng.shuffle(data)
         if max_per_task > 0:
-            data = data[:max_per_task]
+            data_cutoff = max(task_meta['num_fewshot']+1, max_per_task)
+            data = data[:data_cutoff]
 
         # run the evaluation for this task
         accuracy = evaluate_task(model, tokenizer, data, device, task_meta)