update

app.py

@@ -108,16 +108,15 @@ with gr.Blocks() as demo:
 
     ########
     history = gr.State([{"role": "system", "content": system_list[0]}])  # 有用信息只有个system，其他和chatbot内容重叠
-    history_tokens = gr.State([])
-    system.change(reset_state, inputs=[system], outputs=[chatbot, history, history_tokens])
-    clear_btn.click(reset_state, inputs=[system], outputs=[chatbot, history, history_tokens])
+    system.change(reset_state, inputs=[system], outputs=[chatbot, history])
+    clear_btn.click(reset_state, inputs=[system], outputs=[chatbot, history])
 
-    generate_btn.click(generate, [chatbot, history, history_tokens], outputs=[generated_text_1, chatbot, history, history_tokens],
+    generate_btn.click(generate, [chatbot, history], outputs=[generated_text_1, chatbot, history],
                        show_progress="full")
-    retry_btn.click(undo_generate, [chatbot, history, history_tokens], outputs=[generated_text_1, chatbot, history, history_tokens]) \
-        .then(generate, [chatbot, history, history_tokens], outputs=[generated_text_1, chatbot, history, history_tokens],
+    retry_btn.click(undo_generate, [chatbot, history], outputs=[generated_text_1, chatbot, history]) \
+        .then(generate, [chatbot, history], outputs=[generated_text_1, chatbot, history],
               show_progress="full")
-    undo_btn.click(undo_generate, [chatbot, history, history_tokens], outputs=[generated_text_1, chatbot, history, history_tokens])
+    undo_btn.click(undo_generate, [chatbot, history], outputs=[generated_text_1, chatbot, history])
 
     slider_max_tokens.change(set_max_tokens, inputs=[slider_max_tokens])
     slider_temperature.change(set_temperature, inputs=[slider_temperature])

app_util.py

@@ -19,21 +19,21 @@ from models.cpp_qwen2 import bot
 # gr.Chatbot.postprocess = postprocess
 
 
-def generate_query(chatbot, history, history_tokens):
+def generate_query(chatbot, history):
     if history and history[-1]["role"] == "user":
         gr.Warning('You should generate assistant-response.')
         yield None, chatbot, history
     else:
         chatbot.append(None)
-        streamer = bot.generate_query(history[-1], history_tokens, stream=True)
-        for query, all_tokens in streamer:
+        streamer = bot.generate_query(history[-1], stream=True)
+        for query, query_tokens in streamer:
             chatbot[-1] = (query, None)
             yield query, chatbot, history
-        history.append({"role": "user", "content": query})
-        yield query, chatbot, history, all_tokens
+        history.append({"role": "user", "content": query, "tokens": query_tokens})
+        yield query, chatbot, history
 
 
-def generate_response(chatbot, history, history_tokens):
+def generate_response(chatbot, history):
     """
     auto-mode：query is None
     manual-mode：query 是用户输入
@@ -44,24 +44,24 @@ def generate_response(chatbot, history, history_tokens):
         gr.Warning('You should generate or type user-input first.')
         yield None, chatbot, history
     else:
-        streamer = bot.generate_response(history[-1], history_tokens, stream=True)
-        for response, all_tokens in streamer:
+        streamer = bot.generate_response(history[-1], stream=True)
+        for response, response_tokens in streamer:
             chatbot[-1] = (query, response)
-            yield response, chatbot, history, all_tokens
+            yield response, chatbot, history
 
-        history.append({"role": "assistant", "content": response})
+        history.append({"role": "assistant", "content": response, "tokens": response_tokens})
         print(f"chatbot is {chatbot}")
         print(f"history is {history}")
-        yield response, chatbot, history, all_tokens
+        yield response, chatbot, history
 
 
-def generate(chatbot, history, history_tokens):
+def generate(chatbot, history):
     logger.info(f"chatbot: {chatbot}; history: {history}")
     streamer = None
     if history[-1]["role"] in ["assistant", "system"]:
-        streamer = generate_query(chatbot, history, history_tokens)
+        streamer = generate_query(chatbot, history)
     elif history[-1]["role"] == "user":
-        streamer = generate_response(chatbot, history, history_tokens)
+        streamer = generate_response(chatbot, history)
     else:
         gr.Warning("bug")
 
@@ -69,7 +69,7 @@ def generate(chatbot, history, history_tokens):
         yield out
 
 
-def undo_generate(chatbot, history, history_tokens):
+def undo_generate(chatbot, history):
     if history[-1]["role"] == "user":
         history = history[:-1]
         chatbot = chatbot[:-1]
@@ -79,7 +79,7 @@ def undo_generate(chatbot, history, history_tokens):
     else:
         pass
     logger.info(f"after undo, {json.dumps(chatbot, ensure_ascii=False)}, {json.dumps(history, ensure_ascii=False)}")
-    return "", chatbot, history, history_tokens
+    return "", chatbot, history
 
 
 def reset_user_input():

models/cpp_qwen2.py

@@ -81,7 +81,7 @@ class Qwen2Simulator(Simulator):
     def tokenize(self, text):
         return self.llm.tokenize(text.encode("utf-8"))
 
-    def generate_query(self, message, history_tokens, stream=True):
+    def generate_query(self, history, stream=True):
         """
         """
         # {% for message in messages %}
@@ -92,18 +92,28 @@ class Qwen2Simulator(Simulator):
         # {% endfor %}
         # {% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
 
-        input_ids = history_tokens + self.tokenize(
-            f"<|im_start|>{message['role']}\n{message['content']}<|im_end|>\n<|im_start|>user\n"
-        )
+        # TODO: 判断是否有结束符，
+
+        input_ids = []
+        for message in history:
+            if "tokens" not in message:
+                message["tokens"] = self.tokenize(message["content"])
+            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, message, history_tokens, stream=True):
-        input_ids = history_tokens + self.tokenize(
-            f"<|im_start|>{message['role']}\n{message['content']}<|im_end|>\n<|im_start|>assistant\n"
-        )
+    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")
         if stream:
             return self._stream_generate(input_ids)
         else:
@@ -118,18 +128,12 @@ class Qwen2Simulator(Simulator):
             stream=True,
             **self.generation_kwargs
         )
-        generated_text = ""
         # TODO: 检测finish reason，如果是length，则shift，并继续生成。
         # TODO: 返回 token_id,
         for out in output:
             stream = copy.deepcopy(out)
             if stream["choices"][0]["finish_reason"] is None:
-                generated_text += stream["choices"][0]["text"]
-                if "completion_text" in stream["choices"][0]:
-                    yield stream["choices"][0]["completion_text"], stream["choices"][0]["all_tokens"]
-                else:
-                    logger.info("completion_text not found")
-                    yield generated_text, None
+                yield stream["choices"][0]["completion_text"], stream["choices"][0]["completion_tokens"]
 
 
 bot = Qwen2Simulator()
@@ -144,11 +148,10 @@ if __name__ == "__main__":
 
     message = {"role": "system", "content": "你是一个导游。"}
     print(message)
-    for generated_text, all_tokens in bot.generate_query(message, [], stream=True):
+    for generated_text, all_tokens in bot.generate_query(message, stream=True):
         print(generated_text, all_tokens)
 
     message = {"role": "user", "content": generated_text}
     print(message)
     for generated_text, all_tokens in bot.generate_response(message, all_tokens, stream=True):
         print(generated_text, all_tokens)
-