Fix max_tokens handling in vllm_vlms.py (#2637)

* Update vllm_vlms.py * pre-commit --------- Co-authored-by: Baber <[email protected]>
EleutherAI · Jan 21, 2025 · 370e2f9 · 370e2f9
1 parent b2c090c
commit 370e2f9
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/lm_eval/models/vllm_vlms.py b/lm_eval/models/vllm_vlms.py
@@ -271,7 +271,9 @@ def _collate(x):
                 left_truncate_len=max_ctx_len,
             )
 
-            cont = self._model_generate(inputs, stop=until, generate=True, **kwargs)
+            cont = self._model_generate(
+                inputs, stop=until, generate=True, max_tokens=max_gen_toks, **kwargs
+            )
 
             for output, context in zip(cont, contexts):
                 generated_text = output.outputs[0].text