Fix max token setting in EncoderDecoder generator

RUC-NLPIR · Jul 26, 2024 · 96b20c8 · 96b20c8
1 parent ed7ef8e
commit 96b20c8
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/flashrag/generator/generator.py b/flashrag/generator/generator.py
@@ -94,11 +94,12 @@ def generate(self, input_list: List, batch_size=None, **params):
             stopping_criteria = [StopWordCriteria(tokenizer=self.tokenizer, prompts=input_list, stop_words=stop_sym)]
             generation_params['stopping_criteria'] = stopping_criteria
 
-        if 'max_tokens' in generation_params:
-            if 'max_tokens' in params:
-                generation_params['max_new_tokens'] = params.pop('max_tokens')
-            else:
-                generation_params['max_new_tokens'] = generation_params.pop('max_tokens')
+        max_tokens = params.pop('max_tokens', None) or params.pop('max_new_tokens', None)
+        if max_tokens is not None:
+            generation_params['max_new_tokens'] = max_tokens
+        else:
+            generation_params['max_new_tokens'] = generation_params.get('max_new_tokens', generation_params.pop('max_tokens', None))
+        generation_params.pop('max_tokens', None)
 
         responses = []
         for idx in trange(0, len(input_list), batch_size, desc='Generation process: '):