Fix max_tokens handling in vllm_vlms.py (EleutherAI#2637)

jkaniecki · baberabb · web-flow · commit 370e2f9e5bbe · 2025-01-21T21:55:23.000+05:00
* Update vllm_vlms.py

* pre-commit

---------

Co-authored-by: Baber &lt;baber@hey.com&gt;
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

Original file line number	Diff line number	Diff line change
`@@ -271,7 +271,9 @@ def _collate(x):`
`271`	`271`	`left_truncate_len=max_ctx_len,`
`272`	`272`	`)`
`273`	`273`
`274`		`- cont = self._model_generate(inputs, stop=until, generate=True, **kwargs)`
	`274`	`+ cont = self._model_generate(`
	`275`	`+ inputs, stop=until, generate=True, max_tokens=max_gen_toks, **kwargs`
	`276`	`+ )`
`275`	`277`
`276`	`278`	`for output, context in zip(cont, contexts):`
`277`	`279`	`generated_text = output.outputs[0].text`