max_tokensの指定の追加

llm-jp · Jan 14, 2025 · 9d24b41 · 9d24b41
1 parent c5d576b
commit 9d24b41
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/docker_sample/src/main.py b/docker_sample/src/main.py
@@ -1,7 +1,7 @@
 import os
 import json
 import argparse
-from vllm import LLM
+from vllm import LLM, SamplingParams
 
 
 def main():
@@ -27,7 +27,11 @@ def main():
         messages = [{"role": "user", "content": d["text"]}]
         messages_list.append(messages)
 
-    outputs = llm.chat(messages_list)
+    sampling_params = SamplingParams(
+        max_tokens=1024,
+    )
+
+    outputs = llm.chat(messages_list, sampling_params=sampling_params)
     for i, output in enumerate(outputs):
         data[i]["response"] = output.outputs[0].text