fix bug in vllm engine

2024-05-15 02:17:54 +08:00 · 2024-05-15 02:17:54 +08:00 · 11bf282dcc
parent 144801db09
commit 11bf282dcc
1 changed files with 2 additions and 1 deletions
--- 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