From 3e7c74ec7abcc9c24e0c41155615ef7dc2aca3a6 Mon Sep 17 00:00:00 2001
From: zR <2448370773@qq.com>
Date: Sun, 14 Apr 2024 19:00:46 +0800
Subject: [PATCH] =?UTF-8?q?128k=E6=A8=A1=E5=9E=8B=E6=97=A0=E6=B3=95?=
 =?UTF-8?q?=E6=8E=A8=E7=90=86=E7=9A=84=E9=97=AE=E9=A2=98=E4=BF=AE=E5=A4=8D?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 demo/vllm_based_demo.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/demo/vllm_based_demo.py b/demo/vllm_based_demo.py
index efefb37..cf471c4 100644
--- a/demo/vllm_based_demo.py
+++ b/demo/vllm_based_demo.py
@@ -56,6 +56,7 @@ def vllm_gen(dialog: List, top_p: float, temperature: float, max_dec_len: int):
     """
     assert len(dialog) % 2 == 1
     prompt = tokenizer.apply_chat_template(dialog, tokenize=False, add_generation_prompt=False)
+    token_ids = tokenizer.convert_tokens_to_ids(["<|im_end|>"])
     params_dict = {
         "n": 1,
         "best_of": 1,
@@ -67,8 +68,8 @@ def vllm_gen(dialog: List, top_p: float, temperature: float, max_dec_len: int):
         "use_beam_search": False,
         "length_penalty": 1,
         "early_stopping": False,
-        "stop": None,
-        "stop_token_ids": None,
+        "stop": "<|im_end|>",
+        "stop_token_ids": token_ids,
         "ignore_eos": False,
         "max_tokens": max_dec_len,
         "logprobs": None,