update

models/cpp_qwen2.py

@@ -64,6 +64,12 @@ class Qwen2Simulator(Simulator):
                     f"n_threads={self.llm.n_threads}, n_ctx={self.llm.n_ctx}, "
                     f"env[CACHE]={os.environ.get('CACHE', None)}")
 
+        self.stop_words = [
+            "<|im_end|>",
+            "<|im_start|>",
+            "<|endoftext|>",
+        ]
+        self.stop_tokens = self.tokenize(self.stop_words)
         self.generation_kwargs = dict(
             temperature=config.DEFAULT_TEMPERATURE,
             top_p=config.DEFAULT_TOP_P,
@@ -71,28 +77,32 @@ class Qwen2Simulator(Simulator):
             max_tokens=config.DEFAULT_MAX_TOKENS,
             repeat_penalty=1.1,
             # qwen2-0.5b-chat 有时内容生成结束没有<|im_end|>，直接跟 <|im_start|>
-            stop=[
-                "<|im_end|>",
-                "<|im_start|>",
-                "<|endoftext|>",
-            ],
+            stop=self.stop_words,
         )
 
+        self.user_start_tokens = self.tokenize("<|im_start|>user\n")
+        self.assistant_start_tokens = self.tokenize("<|im_start|>assistant\n")
+
     def tokenize(self, text):
         return self.llm.tokenize(text.encode("utf-8"))
 
+    def _strip_stoptokens(self, tokens):
+        while tokens and tokens[0] in self.stop_tokens:
+            tokens.pop(0)
+        while tokens and tokens[-1] in self.stop_tokens:
+            tokens.pop()
 
     def generate(self, history, stream=True):
         if history[-1]['role'] in ["user"]:
-            start_tokens = self.tokenize("<|im_start|>assistant\n")
+            start_tokens = self.assistant_start_tokens
         elif history[-1]['role'] in ["assistant", "system"]:
-            start_tokens = self.tokenize("<|im_start|>user\n")
+            start_tokens = self.user_start_tokens
 
         input_ids = []
         for message in history:
-            if "tokens" not in message:
+            if "tokens" not in message:  # tokens
                 message["tokens"] = self.tokenize(message["content"])
-            input_ids += self.tokenize(f"<|im_start|>{message['role']}\n") \
+            input_ids += self._strip_stoptokens(self.tokenize(f"<|im_start|>{message['role']}\n")) \
                          + message["tokens"] \
                          + self.tokenize("<|im_end|>\n")
         input_ids += start_tokens
@@ -120,7 +130,6 @@ class Qwen2Simulator(Simulator):
                 print(f'finish_reason with text: {stream["choices"][0]["text"]}')
 
 
-
 bot = Qwen2Simulator()
 
 if __name__ == "__main__":
@@ -131,11 +140,10 @@ if __name__ == "__main__":
     for generated_text, generated_tokens in bot.generate(messages, stream=True):
         print(generated_text, generated_tokens)
 
-
     for i in range(3):
-        messages.append({"role": "user" if i % 2 == 0 else "assistant", "content": generated_text, "tokens": generated_tokens})
+        messages.append(
+            {"role": "user" if i % 2 == 0 else "assistant", "content": generated_text, "tokens": generated_tokens})
         print("######## requesting", messages)
         for generated_text, generated_tokens in bot.generate(messages, stream=True):
             pass
             # print(generated_text, all_tokens)
-