From 2c011060b172d361283f7866b7000f8692f82cd1 Mon Sep 17 00:00:00 2001 From: hiyouga Date: Wed, 15 May 2024 02:17:54 +0800 Subject: [PATCH] fix bug in vllm engine Former-commit-id: 38f02a2c5b52cba6908c2d3c2a455677f8574faf --- src/llmtuner/chat/vllm_engine.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/llmtuner/chat/vllm_engine.py b/src/llmtuner/chat/vllm_engine.py index faf8c9fe..aaaad2f1 100644 --- a/src/llmtuner/chat/vllm_engine.py +++ b/src/llmtuner/chat/vllm_engine.py @@ -100,8 +100,9 @@ class VllmEngine(BaseEngine): max_new_tokens = input_kwargs.pop("max_new_tokens", None) stop = input_kwargs.pop("stop", None) + max_tokens = self.generating_args["max_new_tokens"] or self.generating_args["max_length"] if max_length: - max_tokens = max_length - prompt_length + max_tokens = max_length - prompt_length if max_length > prompt_length else 1 if max_new_tokens: max_tokens = max_new_tokens