update

app.py

@@ -14,9 +14,15 @@ system_list = [
 ]
 
 user_simulator_doc = """\
+
+
+
 There are maily two types of user simulator:
 - prompt-based user-simulator (role-play)
 - model-based user-simulator
+
+In most cases, large language models (LLMs) are used to serve as assistant generator.
+Besides, it can also used as user simulator.
 """
 
 with gr.Blocks() as demo:
@@ -101,16 +107,17 @@ with gr.Blocks() as demo:
                 )
 
     ########
-    history = gr.State([{"role": "system", "content": system_list[0]}])
-    system.change(reset_state, inputs=[system], outputs=[chatbot, history])
-    clear_btn.click(reset_state, inputs=[system], outputs=[chatbot, history])
+    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])
 
-    generate_btn.click(generate, [chatbot, history], outputs=[generated_text_1, chatbot, history],
+    generate_btn.click(generate, [chatbot, history, history_tokens], outputs=[generated_text_1, chatbot, history, history_tokens],
                        show_progress="full")
-    retry_btn.click(undo_generate, [chatbot, history], outputs=[generated_text_1, chatbot, history]) \
-        .then(generate, [chatbot, history], outputs=[generated_text_1, chatbot, history],
+    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],
               show_progress="full")
-    undo_btn.click(undo_generate, [chatbot, history], outputs=[generated_text_1, chatbot, history])
+    undo_btn.click(undo_generate, [chatbot, history, history_tokens], outputs=[generated_text_1, chatbot, history, history_tokens])
 
     slider_max_tokens.change(set_max_tokens, inputs=[slider_max_tokens])
     slider_temperature.change(set_temperature, inputs=[slider_temperature])

app_util.py

@@ -19,54 +19,49 @@ from models.cpp_qwen2 import bot
 # gr.Chatbot.postprocess = postprocess
 
 
-def generate_query(chatbot, history):
+def generate_query(chatbot, history, history_tokens):
     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, stream=True)
-        for query in streamer:
+        streamer = bot.generate_query(history[-1], history_tokens, stream=True)
+        for query, all_tokens in streamer:
             chatbot[-1] = (query, None)
             yield query, chatbot, history
         history.append({"role": "user", "content": query})
-        yield query, chatbot, history
+        yield query, chatbot, history, all_tokens
 
 
-def generate_response(chatbot, history, user_input=None):
+def generate_response(chatbot, history, history_tokens):
     """
     auto-mode：query is None
     manual-mode：query 是用户输入
-    :param chatbot:
-    :param history:
-    :return:
     """
-    if user_input and history[-1]["role"] != "user":
-        history.append({"role": "user", "content": user_input})
+    logger.info(f"generating {json.dumps(history, ensure_ascii=False)}")
     query = history[-1]["content"]
-
     if history[-1]["role"] != "user":
         gr.Warning('You should generate or type user-input first.')
         yield None, chatbot, history
     else:
-        streamer = bot.generate_response(history, stream=True)
-        for response in streamer:
+        streamer = bot.generate_response(history[-1], history_tokens, stream=True)
+        for response, all_tokens in streamer:
             chatbot[-1] = (query, response)
-            yield response, chatbot, history
+            yield response, chatbot, history, all_tokens
 
         history.append({"role": "assistant", "content": response})
         print(f"chatbot is {chatbot}")
         print(f"history is {history}")
-        yield response, chatbot, history
+        yield response, chatbot, history, all_tokens
 
 
-def generate(chatbot, history):
+def generate(chatbot, history, history_tokens):
     logger.info(f"chatbot: {chatbot}; history: {history}")
     streamer = None
     if history[-1]["role"] in ["assistant", "system"]:
-        streamer = generate_query(chatbot, history)
+        streamer = generate_query(chatbot, history, history_tokens)
     elif history[-1]["role"] == "user":
-        streamer = generate_response(chatbot, history)
+        streamer = generate_response(chatbot, history, history_tokens)
     else:
         gr.Warning("bug")
 
@@ -74,7 +69,7 @@ def generate(chatbot, history):
         yield out
 
 
-def undo_generate(chatbot, history):
+def undo_generate(chatbot, history, history_tokens):
     if history[-1]["role"] == "user":
         history = history[:-1]
         chatbot = chatbot[:-1]
@@ -84,7 +79,7 @@ def undo_generate(chatbot, history):
     else:
         pass
     logger.info(f"after undo, {json.dumps(chatbot, ensure_ascii=False)}, {json.dumps(history, ensure_ascii=False)}")
-    return "", chatbot, history
+    return "", chatbot, history, history_tokens
 
 
 def reset_user_input():
@@ -92,7 +87,7 @@ def reset_user_input():
 
 
 def reset_state(system):
-    return [], [{"role": "system", "content": system}]
+    return [], [{"role": "system", "content": system}], []
 
 
 def set_max_tokens(max_tokens):

models/cpp_qwen2.py

@@ -36,14 +36,14 @@ import config
 
 class Qwen2Simulator(Simulator):
 
-    def __init__(self, from_local=False):
-        if from_local:
-            self.hf_tokenizer = AutoTokenizer.from_pretrained(
-                "/workspace/xusong/huggingface/models/Qwen2-0.5B-Instruct/")
-            self.llm = llama_cpp.Llama(  # n_ctx, n_threads
+    def __init__(self):
+        self.hf_tokenizer = AutoTokenizer.from_pretrained(
+            "/workspace/xusong/huggingface/models/Qwen2-0.5B-Instruct/")
 
-                model_path="/workspace/xusong/huggingface/models/Qwen2-0.5B-Instruct-GGUF/qwen2-0_5b-instruct-fp16.gguf",
-                tokenizer=llama_cpp.llama_tokenizer.LlamaHFTokenizer(self.hf_tokenizer),
+        local_path = "/workspace/xusong/huggingface/models/Qwen2-0.5B-Instruct-GGUF/qwen2-0_5b-instruct-fp16.gguf"
+        if os.path.exists(local_path):
+            self.llm = llama_cpp.Llama(  # n_ctx, n_threads
+                model_path=local_path,
                 n_ctx=config.MAX_SEQUENCE_LENGTH,  #
                 # n_threads=None, # 默认会根据cpu数来设置 n_threads
                 use_mlock=True,
@@ -54,7 +54,6 @@ class Qwen2Simulator(Simulator):
             self.llm = llama_cpp.Llama.from_pretrained(
                 repo_id="Qwen/Qwen2-0.5B-Instruct-GGUF",
                 filename="*fp16.gguf",
-                tokenizer=llama_cpp.llama_tokenizer.LlamaHFTokenizer(self.hf_tokenizer),
                 n_ctx=config.MAX_SEQUENCE_LENGTH,
                 use_mlock=True,
                 verbose=False,
@@ -77,68 +76,58 @@ class Qwen2Simulator(Simulator):
             ],
         )
 
-    def generate_query(self, messages, stream=True):
+    def tokenize(self, text):
+        return self.llm.tokenize(text.encode("utf-8"))
+
+    def generate_query(self, message, history_tokens, stream=True):
         """
-        :param messages:
-        :return:
         """
-        assert messages[-1]["role"] != "user"
-        logger.info(f"generating {json.dumps(messages, ensure_ascii=False)}")
-        inputs = self.hf_tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=False,
+        # {% 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 %}",
+
+        input_ids = history_tokens + self.tokenize(
+            f"<|im_start|>{message['role']}\n{message['content']}<|im_end|>\n<|im_start|>user\n"
         )
-        inputs = inputs + "<|im_start|>user\n"
         if stream:
-            return self._stream_generate(inputs)
+            return self._stream_generate(input_ids)
         else:
-            return self._generate(inputs)
-
-    def generate_response(self, messages, stream=True):
-        assert messages[-1]["role"] == "user"
-        logger.info(f"generating {json.dumps(messages, ensure_ascii=False)}")
-        inputs = self.hf_tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True
+            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"
         )
         if stream:
-            return self._stream_generate(inputs)
+            return self._stream_generate(input_ids)
         else:
-            return self._generate(inputs)
-
-    def _generate(self, inputs):
-        logger.info(f"generation_kwargs {self.generation_kwargs}")
-        output = self.llm(
-            inputs,
-            **self.generation_kwargs
-        )
-        output_text = output["choices"][0]["text"]
-        return output_text
+            return self._generate(input_ids)
 
-    def _stream_generate(self, inputs):
+    def _stream_generate(self, input_ids):
         logger.info(f"generation_kwargs {self.generation_kwargs}")
 
         # self.llm.generate  .set_cache   .last_n_tokens_size  .reset  .ctx ._ctx
-        output = self.llm(
-            inputs,
+        output = self.llm.create_completion(
+            input_ids,
             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 not None:
+            if stream["choices"][0]["finish_reason"] is None:
                 generated_text += stream["choices"][0]["text"]
-                if "all_text" in stream["choices"][0]:
-                    yield stream["choices"][0]["all_text"]
+                if "completion_text" in stream["choices"][0]:
+                    yield stream["choices"][0]["completion_text"], stream["choices"][0]["all_tokens"]
                 else:
-                    logger.info("all_text not found")
-                    yield generated_text
+                    logger.info("completion_text not found")
+                    yield generated_text, None
 
 
 bot = Qwen2Simulator()