From 910f2bb00dfd6d04046f73e3625916044a737bb8 Mon Sep 17 00:00:00 2001
From: "clementine@huggingface.co" <clementine@huggingface.co>
Date: Thu, 18 Jul 2024 09:40:20 +0000
Subject: [PATCH] fix

---
 src/lighteval/models/base_model.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/lighteval/models/base_model.py b/src/lighteval/models/base_model.py
index f1ba61517..8e69cdbbd 100644
--- a/src/lighteval/models/base_model.py
+++ b/src/lighteval/models/base_model.py
@@ -554,7 +554,10 @@ def greedy_until(
                     # There will be truncation of at least one sample, maximum generation size will be one
                     max_new_tokens = 1
                 else:  # We can't allow generation of more than max_length
-                    max_new_tokens = min(self.max_length - context_size, max_new_tokens)
+                    if max_new_tokens is None:  # If generation size is not set, we go all the way
+                        max_new_tokens = self.max_length - context_size
+                    else:
+                        max_new_tokens = min(self.max_length - context_size, max_new_tokens)
 
                 prepared_batch = Batch(
                     input_ids=tokenized["input_ids"],