From 967b9c0a498c1c527ae0bbfb6350535eee6b13e2 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: 11bf282dcc0ee257f2c28f46cc1a8edcf62421dc --- 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