update

app.py

@@ -1,28 +1,39 @@
 """
 """
 import random
-import gradio
 import config
 from app_util import *
 
-user_simulator_doc = """\
-The agent acts as user simulator.
+user_simulator_pre_doc = """\
+You are a helpful assistant, and the agent acts as user. 
+"""
 
+user_simulator_post_doc = """\
+## How does it work?
 There are maily two types of user simulator:
 - prompt-based user-simulator (role-play)
 - model-based user-simulator
 
 This demo is a 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.
 
-assistant_simulator_doc = """\
-The agent acts as assistant simulator.
+assistant_simulator_pre_doc = """\
+You are a user, and the agent acts as assistant.
+"""
+
+assistant_simulator_post_doc = """\
+"""
+
+self_chat_pre_doc = """\
+Self-chat is a demo which make the model talk to itself. Dual-agent. 
 """
 
-self_chat_doc = """\
-Self-chat is a demo which make the model talk to itself. 
+self_chat_post_doc = """\
+## How does it work?
 It is a combination of user simulator and response generator.
 """
 
@@ -37,12 +48,25 @@ Essentially, it is a form of model compression.
 
 ## 有不用概率的知识蒸馏吗？
 """
+gr.set_static_paths(paths=["assets/"])
 
-with gr.Blocks(head=None) as demo:
+"""
+<div class="avatar-container"><img src="file=assets/man.png" class="avatar-image" alt="user avatar"></div>
+"""
+
+
+css="""
+.image_center {
+      display: block;
+      margin: auto;
+}
+"""
+
+with gr.Blocks(head=None, css=css) as demo:
     # Knowledge Distillation through Self Chatting
     # Distilling the Knowledge from LLM through Self Chatting
     # Generating Synthetic Data through Self Chat
-    gr.HTML("""<h1 align="center">Generating Synthetic Data via Self-Chat</h1>""")
+    gr.HTML("""<h1 align="center">Generating Synthetic Data via Self-Chatting</h1>""")
     with gr.Row():
         with gr.Column(scale=5):
             system = gr.Dropdown(
@@ -56,47 +80,69 @@ with gr.Blocks(head=None) as demo:
 
             chatbot = gr.Chatbot(show_copy_button=True,
                                  show_share_button=True,
-                                 avatar_images=("assets/man.png", "assets/bot.png"),
+                                 # avatar_images=("assets/man.png", "assets/bot.png"),
+                                 avatar_images=("assets/man.png", "assets/女客服.png"),
                                  likeable=True)
 
             # gr.Textbox("For faster inference, you can build locally with  ")
             # ss
-            with gradio.Tab("Self Chat"):
+            with gr.Tab("Self Chat") as tab_dual_agent:
+                gr.Markdown(self_chat_pre_doc)
                 input_text_1 = gr.Textbox(show_label=False, placeholder="...", lines=10, visible=False)
-                generate_btn = gr.Button("🤔️ Self-Chat", variant="primary")
+                generate_btn = gr.Button("🤖Self-Chat🤖", variant="primary")
                 with gr.Row():
-                    retry_btn = gr.Button("🔄  Regenerate", variant="secondary", size="sm", )
-                    undo_btn = gr.Button("↩️ Undo", variant="secondary", size="sm", )
-                    clear_btn = gr.Button("🗑️  Clear", variant="secondary", size="sm", )  # 🧹 Clear History (清除历史)
-                    # stop_btn = gr.Button("停止生成", variant="stop", visible=False)
-                gr.Markdown(self_chat_doc)
-
-            # 也叫 chat-assistant,
-            with gradio.Tab("Response Generator"):
+                    retry_btn = gr.Button("🔄  Regenerate", variant="secondary", size="sm")
+                    undo_btn = gr.Button("↩️ Undo", variant="secondary", size="sm")
+                    # clear_btn = gr.Button("🗑️  Clear", variant="secondary", size="sm")
+                    clear_btn = gr.Button("🧹 Clear History", variant="secondary", size="sm")
+                gr.Markdown(self_chat_post_doc)
+
+            # 也叫 chat-assistant, 🎧，🤖 ，💁，
+            with gr.Tab("Response Generator") as tab_assistant_agent:
+                gr.Markdown(assistant_simulator_pre_doc)
                 with gr.Row():
-                    input_text_2 = gr.Textbox(show_label=False, placeholder="Please type user input", scale=7)
-                    generate_btn_2 = gr.Button("Send", variant="primary")
+                    # gr.HTML(
+                    #     value='<div><img src="/file=./assets/man.png" alt="Big Boat" width="40px" height="40px"></div>',
+                    #     elem_classes=["image_center"]
+                    # )
+                    gr.Image("assets/man.png", interactive=False, show_download_button=False, width=40, height=40,
+                             min_width=40,
+                             show_share_button=False, show_fullscreen_button=False, container=False,
+                             elem_classes=["image_center"])
+                    input_text_2 = gr.Textbox(show_label=False, lines=2, placeholder="Please type user input",
+                                              container=False, scale=12)
+                    generate_btn_2 = gr.Button("Send", variant="primary", min_width=80)
                 with gr.Row():
                     retry_btn_2 = gr.Button("🔄  Regenerate", variant="secondary", size="sm", )
                     undo_btn_2 = gr.Button("↩️ Undo", variant="secondary", size="sm", )
-                    clear_btn_2 = gr.Button("🗑️  Clear", variant="secondary", size="sm", )  # 🧹 Clear History (清除历史)
-                gr.Markdown(assistant_simulator_doc)
+                    clear_btn_2 = gr.Button("🧹 Clear History", variant="secondary", size="sm")
+                gr.Markdown(assistant_simulator_post_doc)
 
             #
-            with gradio.Tab("User Simulator"):
+            with gr.Tab("User Simulator") as tab_user_agent:  # 👨，🔊，
+                gr.Markdown(user_simulator_pre_doc)
                 with gr.Row():
-                    input_text_3 = gr.Textbox(show_label=False, placeholder="Please type assistant response", scale=7)
-                    generate_btn_3 = gr.Button("Send", variant="primary")
+                    # gr.HTML(value='<div class="avatar-container"><img src="file=assets/man.png" class="avatar-image" alt="user avatar"></div>')
+                    # gr.Image("assets/女客服.jpg",
+                    gr.Image("assets/女客服.png",
+                    # gr.Image("assets/男客服.png",
+                             interactive=False, show_download_button=False, width=40, height=40,
+                             min_width=40,
+                             show_share_button=False, show_fullscreen_button=False, container=False, elem_classes=["image_center"])
+                    input_text_3 = gr.Textbox(show_label=False, lines=2, placeholder="Please type assistant response",
+                                              container=False, scale=12)
+                    generate_btn_3 = gr.Button("Send", variant="primary", min_width=80)
                 with gr.Row():
-                    retry_btn_3 = gr.Button("🔄  Regenerate", variant="secondary", size="sm", )
-                    undo_btn_3 = gr.Button("↩️ Undo", variant="secondary", size="sm", )
-                    clear_btn_3 = gr.Button("🗑️  Clear", variant="secondary", size="sm", )  # 🧹 Clear History (清除历史)
-                gr.Markdown(user_simulator_doc)
+                    retry_btn_3 = gr.Button("🔄  Regenerate", variant="secondary", size="sm")
+                    undo_btn_3 = gr.Button("↩️ Undo", variant="secondary", size="sm")
+                    # clear_btn_3 = gr.Button("🗑️  Clear", variant="secondary", size="sm")
+                    clear_btn_3 = gr.Button("🧹 Clear History", variant="secondary", size="sm")  # 🧹 Clear History (清除历史)
+                gr.Markdown(user_simulator_post_doc)
 
         with gr.Column(variant="compact", scale=1, min_width=300):
             # with gr.Column():
             model = gr.Dropdown(
-                ["Qwen2-0.5B-Instruct", "llama3.1", "gemini"],
+                ["Qwen2-0.5B-Instruct", "llama3.1", "gemini", "MiniCPM3-4B"],
                 value="Qwen2-0.5B-Instruct",
                 label="Model",
                 interactive=True,
@@ -124,6 +170,7 @@ with gr.Blocks(head=None) as demo:
                 )
 
     # TODO: gr.State 不能通过API传参。
+    gr_false = gr.State(False)
     history = gr.State([{"role": "system", "content": system_list[0]}])  # 有用信息只有个system，其他和chatbot内容重叠
     system.change(reset_state, inputs=[system], outputs=[chatbot, history])
 
@@ -147,6 +194,8 @@ with gr.Blocks(head=None) as demo:
     undo_btn_2.click(undo_generate, [chatbot, history], outputs=[chatbot, history], show_api=False)
     clear_btn_2.click(reset_state, inputs=[system], outputs=[chatbot, history], show_api=False) \
         .then(reset_user_input, outputs=[input_text_2], show_api=False)
+    tab_assistant_agent.select(generate_assistant_message, [chatbot, history, gr_false], outputs=[chatbot, history],
+              show_progress="full", show_api=False)   # 点击tab，生成response (不warning)
 
     ######## tab3: user-simulator
     generate_btn_3.click(append_assistant_to_history, [input_text_3, chatbot, history], outputs=[chatbot, history],
@@ -160,6 +209,9 @@ with gr.Blocks(head=None) as demo:
     clear_btn_3.click(reset_state, inputs=[system], outputs=[chatbot, history], show_api=False) \
         .then(reset_user_input, outputs=[input_text_3], show_api=False)
 
+    tab_user_agent.select(generate_user_message, [chatbot, history, gr_false], outputs=[chatbot, history],
+              show_progress="full", show_api=False)  # 点击tab，生成user-input
+
     slider_max_new_tokens.change(set_max_new_tokens, inputs=[slider_max_new_tokens])
     slider_temperature.change(set_temperature, inputs=[slider_temperature])
     slider_top_p.change(set_top_p, inputs=[slider_top_p])

app_util.py

@@ -4,7 +4,7 @@ from typing import List, Dict
 from utils.logging_util import logger
 from models.cpp_qwen2 import Qwen2Simulator as Bot
 # from models.hf_qwen2 import Qwen2Simulator as Bot
-
+# from models.mock import MockSimulator as Bot
 
 #
 # def postprocess(self, y):
@@ -31,12 +31,13 @@ system_list = [
     "你是一名作家，擅长写小说。"
 ]
 
-
 bot = Bot(system_list)
 
-def generate_user_message(chatbot, history):
+
+def generate_user_message(chatbot, history, show_warning=True):
     if history and history[-1]["role"] == "user":
-        gr.Warning('You should generate assistant-response.')
+        if show_warning:
+            gr.Warning('You should generate assistant-response.')
         yield chatbot, history
     else:
         chatbot.append(None)
@@ -49,14 +50,15 @@ def generate_user_message(chatbot, history):
         yield chatbot, history
 
 
-def generate_assistant_message(chatbot, history):
+def generate_assistant_message(chatbot, history, show_warning=True):
     """
     auto-mode：query is None
     manual-mode：query 是用户输入
     """
     user_content = history[-1]["content"]
     if history[-1]["role"] != "user":
-        gr.Warning('You should generate or type user-input first.')
+        if show_warning:
+            gr.Warning('You should generate or type user-input first.')
         yield chatbot, history
     else:
         streamer = bot.generate(history, stream=True)
@@ -116,7 +118,6 @@ def append_assistant_to_history(input_content, chatbot, history):
     return chatbot, history
 
 
-
 def undo_generate(chatbot, history):
     if history[-1]["role"] == "user":
         history = history[:-1]
@@ -141,12 +142,14 @@ def reset_state(system):
 def set_max_new_tokens(max_new_tokens):
     bot.generation_kwargs["max_tokens"] = max_new_tokens
 
+
 def set_temperature(temperature):
     bot.generation_kwargs["temperature"] = temperature
+
+
 def set_top_p(top_p):
     bot.generation_kwargs["top_p"] = top_p
 
+
 def set_top_k(top_k):
     bot.generation_kwargs["top_k"] = top_k
-
-

models/mock.py

@@ -0,0 +1,17 @@
+"""
+"""
+
+from models.base_model import Simulator
+
+class MockSimulator(Simulator):
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def strip_stoptokens(self, tokens):
+        return tokens
+
+    def generate(self, history, stream=True):
+        for text in ['hello', 'world']:
+            yield text, [11]
+