update

models/cpp_qwen2.py

@@ -81,18 +81,12 @@ class Qwen2Simulator(Simulator):
     def tokenize(self, text):
         return self.llm.tokenize(text.encode("utf-8"))
 
-    def generate_query(self, history, stream=True):
-        """
-        """
-        # {% for message in messages %}
-        #   {% if loop.first and messages[0]['role'] != 'system' %}
-        #     {{ '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}
-        #   {% endif %}
-        #   {{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}
-        # {% endfor %}
-        # {% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
-
-        # TODO: 判断是否有结束符，
+
+    def generate(self, history, stream=True):
+        if history[-1]['role'] in ["user"]:
+            start_tokens = self.tokenize("<|im_start|>assistant\n")
+        elif history[-1]['role'] in ["assistant", "system"]:
+            start_tokens = self.tokenize("<|im_start|>user\n")
 
         input_ids = []
         for message in history:
@@ -101,19 +95,7 @@ class Qwen2Simulator(Simulator):
             input_ids += self.tokenize(f"<|im_start|>{message['role']}\n") \
                          + message["tokens"] \
                          + self.tokenize("<|im_end|>\n")
-        input_ids += self.tokenize("<|im_start|>user\n")
-        if stream:
-            return self._stream_generate(input_ids)
-        else:
-            return self._generate(input_ids)
-
-    def generate_response(self, history, stream=True):
-        input_ids = []
-        for message in history:
-            input_ids += self.tokenize(f"<|im_start|>{message['role']}\n") \
-                         + message["tokens"] \
-                         + self.tokenize("<|im_end|>\n")
-        input_ids += self.tokenize("<|im_start|>assistant\n")
+        input_ids += start_tokens
         if stream:
             return self._stream_generate(input_ids)
         else:
@@ -151,7 +133,9 @@ if __name__ == "__main__":
     for generated_text, generated_tokens in bot.generate_query(messages, stream=True):
         print(generated_text, generated_tokens)
 
-    messages.append({"role": "user", "content": generated_text, "tokens": generated_tokens})
-    print(messages)
-    for generated_text, all_tokens in bot.generate_response(messages, stream=True):
-        print(generated_text, all_tokens)
+    for i in range(3):
+        messages.append({"role": "user" if i % 2 == 0 else "assistant", "content": generated_text, "tokens": generated_tokens})
+        print("requesting", messages)
+        for generated_text, all_tokens in bot.generate(messages, stream=True):
+            print(generated_text, all_tokens)
+