From 910f2bb00dfd6d04046f73e3625916044a737bb8 Mon Sep 17 00:00:00 2001 From: "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"],