update

app.py

@@ -20,15 +20,10 @@ python moss_web_demo_gradio.py --model_name fnlp/moss-moon-003-sft --gpu 0,1,2,3
 -
 """
 
-from transformers.generation.utils import logger
 
 import gradio as gr
-import argparse
-import warnings
-import torch
-import os
-# from models.hf_qwen2 import bot
-from models.cpp_qwen2 import bot
+from models.hf_qwen2 import bot
+# from models.cpp_qwen2 import bot
 
 
 #
@@ -83,11 +78,18 @@ def generate_query(chatbot, history):
     if history and history[-1]["role"] == "user":  # 该生成response了
         gr.Warning('You should generate assistant-response.')
         return None, chatbot, history
-    query = bot.generate_query(history)
+    streamer = bot.generate_query(history)
     # chatbot.append((query, ""))
+
+    query = ""
+    for new_text in streamer:
+        print(new_text)
+        query += new_text
+        yield query, chatbot, history
+
     chatbot.append((query, None))
     history.append({"role": "user", "content": query})
-    return query, chatbot, history
+    yield query, chatbot, history
 
 
 def generate_response(query, chatbot, history):
@@ -200,6 +202,8 @@ with gr.Blocks() as demo:
     generate_query_btn.click(generate_query, [chatbot, history], outputs=[user_input, chatbot, history],
                              show_progress="full")
 
+    # generate_query_btn.
+
     gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
     gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature",
               info="Larger temperature increase the randomness"),
@@ -211,5 +215,5 @@ with gr.Blocks() as demo:
         label="Top-p (nucleus sampling)",
     ),
 
-demo.queue().launch(share=False)
+demo.queue().launch(share=False, server_name="0.0.0.0")
 # demo.queue().launch(share=True)

models/cpp_qwen2.py

@@ -16,6 +16,7 @@ python convert_hf_to_gguf.py /workspace/xusong/huggingface/models/Qwen1.5-0.5B-C
 """
 
 import json
+import copy
 from simulator import Simulator
 import llama_cpp
 # import llama_cpp.llama_tokenizer
@@ -45,9 +46,22 @@ class Qwen2Simulator(Simulator):
         logger.info(f"llm has been initialized: {self.llm}")
         # warmup
 
+
+        self.generation_kwargs = dict(
+            temperature=5,
+            # top_p=0.1,
+            top_k=40,
+            max_tokens=20,
+            repeat_penalty=1.1,
+            stop=[
+                "<|im_end|>",
+                "<|im_start|>",
+                "<|endoftext|>",
+            ],
+        )
         ### local
 
-    def generate_query(self, messages):
+    def generate_query(self, messages, stream=True):
         """
         :param messages:
         :return:
@@ -60,11 +74,13 @@ class Qwen2Simulator(Simulator):
             add_generation_prompt=False,
         )
         inputs = inputs + "<|im_start|>user\n"
-        return self._generate(inputs)
-        # for new_text in self._stream_generate(input_ids):
-        #     yield new_text
+        if stream:
+            return self._stream_generate(inputs)
+        else:
+            return self._generate(inputs)
+
 
-    def generate_response(self, messages):
+    def generate_response(self, messages, stream=True):
         assert messages[-1]["role"] == "user"
         logger.info(f"generating {json.dumps(messages)}")
         inputs = self.hf_tokenizer.apply_chat_template(
@@ -72,13 +88,14 @@ class Qwen2Simulator(Simulator):
             tokenize=False,
             add_generation_prompt=True
         )
-
-        return self._generate(inputs)
-        # for new_text in self._stream_generate(input_ids):
-        #     yield new_text
+        if stream:
+            return self._stream_generate(inputs)
+        else:
+            return self._generate(inputs)
 
     def _generate(self, inputs):
         """
+        TODO: chat with cache.
         qwen2-0.5b-chat 有bug：有时user生成结束没有<|im_end|>，示例：
             <|im_start|>system
             you are a helpful assistant<|im_end|>
@@ -91,16 +108,24 @@ class Qwen2Simulator(Simulator):
             <|im_start|>assistant
             I am a 41-year-old man.<|im_end|>
         """
-        # stream=False
         output = self.llm(
             inputs,
-            max_tokens=20,
-            temperature=5,
-            stop=["<|im_end|>", "<|im_start|>"]
+            **self.generation_kwargs
         )
         output_text = output["choices"][0]["text"]
         return output_text
 
+    def _stream_generate(self, inputs):
+        output = self.llm(
+            inputs,
+            stream=True,
+            **self.generation_kwargs
+        )
+        generated_text = ""
+        for out in output:
+            stream = copy.deepcopy(out)
+            generated_text += stream["choices"][0]["text"]
+            yield generated_text
 
 bot = Qwen2Simulator()
 
@@ -117,5 +142,5 @@ if __name__ == "__main__":
         {"role": "user", "content": "hi, what your name"},
         {"role": "assistant", "content": "My name is Jordan"}
     ]
-    output = bot.generate_query(messages)
-    print(output)
+    print(list(bot.generate_query(messages, stream=True)))
+    print(bot.generate_query(messages, stream=False))

models/hf_qwen2.py

@@ -26,10 +26,10 @@ class Qwen2Simulator(Simulator):
             temperature=0.7,
             # repetition_penalty=
             max_length=500,
-            max_new_tokens=200
+            max_new_tokens=20
         )
 
-    def generate_query(self, messages):
+    def generate_query(self, messages, stream=True):
         """
         :param messages:
         :return:
@@ -42,11 +42,23 @@ class Qwen2Simulator(Simulator):
         )
         inputs = inputs + "<|im_start|>user\n"
         input_ids = self.tokenizer.encode(inputs, return_tensors="pt").to(self.model.device)
-        return self._generate(input_ids)
-        # for new_text in self._stream_generate(input_ids):
-        #     yield new_text
 
-    def generate_response(self, messages):
+        streamer = TextIteratorStreamer(tokenizer=self.tokenizer, skip_prompt=True, timeout=120.0,
+                                        skip_special_tokens=True)
+
+        stream_generation_kwargs = dict(
+            input_ids=input_ids,
+            streamer=streamer
+        ).update(self.generation_kwargs)
+        thread = Thread(target=self.model.generate, kwargs=stream_generation_kwargs)
+        thread.start()
+
+        for new_text in streamer:
+            print(new_text)
+            yield new_text
+        # return self._generate(input_ids)
+
+    def generate_response(self, messages, stream=True):
         assert messages[-1]["role"] == "user"
         input_ids = self.tokenizer.apply_chat_template(
             messages,
@@ -54,46 +66,49 @@ class Qwen2Simulator(Simulator):
             return_tensors="pt",
             add_generation_prompt=True
         ).to(self.model.device)
-        return self._generate(input_ids)
-        # for new_text in self._stream_generate(input_ids):
-        #     yield new_text
 
-    def _generate(self, input_ids):
-        input_ids_length = input_ids.shape[-1]
-        response = self.model.generate(input_ids=input_ids, **self.generation_kwargs)
-        return self.tokenizer.decode(response[0][input_ids_length:], skip_special_tokens=True)
-
-    def _stream_generate(self, input_ids):
-        streamer = TextIteratorStreamer(tokenizer=self.tokenizer, skip_prompt=True, timeout=60.0,
-                                        skip_special_tokens=True)
+        streamer = TextIteratorStreamer(
+            tokenizer=self.tokenizer,
+            # skip_prompt=True,
+            # timeout=120.0,
+            # skip_special_tokens=True
+        )
 
-        stream_generation_kwargs = dict(
+        generation_kwargs = dict(
             input_ids=input_ids,
             streamer=streamer
         ).update(self.generation_kwargs)
-        thread = Thread(target=self.model.generate, kwargs=stream_generation_kwargs)
+        print(generation_kwargs)
+
+        thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
         thread.start()
 
         for new_text in streamer:
+            print(new_text)
             yield new_text
 
+    def _generate(self, input_ids):
+        input_ids_length = input_ids.shape[-1]
+        response = self.model.generate(input_ids=input_ids, **self.generation_kwargs)
+        return self.tokenizer.decode(response[0][input_ids_length:], skip_special_tokens=True)
+
 
 bot = Qwen2Simulator(r"E:\data_model\Qwen2-0.5B-Instruct")
 # bot = Qwen2Simulator("Qwen/Qwen2-0.5B-Instruct")
 
 
 if __name__ == "__main__":
-    # messages = [
-    #     {"role": "system", "content": "you are a helpful assistant"},
-    #     {"role": "user", "content": "hi, what your name"}
-    # ]
-    # output = bot.generate_response(messages)
-    # print(output)
-
     messages = [
         {"role": "system", "content": "you are a helpful assistant"},
-        {"role": "user", "content": "hi, what your name"},
-        {"role": "assistant", "content": "My name is Jordan"}
+        {"role": "user", "content": "hi, what your name"}
     ]
-    output = bot.generate_query(messages)
-    print(output)
+    streamer = bot.generate_response(messages)
+    # print(output)
+
+    # messages = [
+    #     {"role": "system", "content": "you are a helpful assistant"},
+    #     {"role": "user", "content": "hi, what your name"},
+    #     {"role": "assistant", "content": "My name is Jordan"}
+    # ]
+    # streamer = bot.generate_query(messages)
+    print(list(streamer))