feat: multiple context, close Akegarasu#20

YunCheng66 · Mar 22, 2023 · 284402e · 284402e
1 parent d34eeab
commit 284402e
Showing 1 changed file with 16 additions and 22 deletions.
diff --git a/modules/ui.py b/modules/ui.py
@@ -3,15 +3,15 @@
 import gradio as gr
 
 from modules import options
-from modules.context import ctx
+from modules.context import Context
 from modules.model import infer
 
 css = "style.css"
 script_path = "scripts"
 _gradio_template_response_orig = gr.routes.templates.TemplateResponse
 
 
-def predict(query, max_length, top_p, temperature, use_stream_chat):
+def predict(ctx, query, max_length, top_p, temperature, use_stream_chat):
     ctx.limit_round()
     flag = True
     for _, output in infer(
@@ -32,12 +32,12 @@ def predict(query, max_length, top_p, temperature, use_stream_chat):
     yield ctx.rh, ""
 
 
-def clear_history():
+def clear_history(ctx):
     ctx.clear()
     return gr.update(value=[])
 
 
-def apply_max_round_click(max_round):
+def apply_max_round_click(ctx, max_round):
     ctx.max_rounds = max_round
     return f"Applied: max round {ctx.max_rounds}"
 
@@ -46,7 +46,8 @@ def create_ui():
     reload_javascript()
 
     with gr.Blocks(css=css, analytics_enabled=False) as chat_interface:
-        prompt = "输入你的内容..."
+        _ctx = Context()
+        state = gr.State(_ctx)
         with gr.Row():
             with gr.Column(scale=3):
                 gr.Markdown("""<h2><center>ChatGLM WebUI</center></h2>""")
@@ -68,7 +69,7 @@ def create_ui():
                 with gr.Row():
                     with gr.Column(variant="panel"):
                         with gr.Row():
-                            clear = gr.Button("清空对话")
+                            clear_history_btn = gr.Button("清空对话")
 
                         with gr.Row():
                             sync_his_btn = gr.Button("同步对话")
@@ -83,7 +84,7 @@ def create_ui():
             with gr.Column(scale=7):
                 chatbot = gr.Chatbot(elem_id="chat-box", show_label=False).style(height=800)
                 with gr.Row():
-                    input_message = gr.Textbox(placeholder=prompt, show_label=False, lines=2, elem_id="chat-input")
+                    input_message = gr.Textbox(placeholder="输入你的内容...", show_label=False, lines=2, elem_id="chat-input")
                     clear_input = gr.Button("🗑️", elem_id="del-btn")
 
                 with gr.Row():
@@ -93,6 +94,7 @@ def create_ui():
                     revoke_btn = gr.Button("撤回")
 
         submit.click(predict, inputs=[
+            state,
             input_message,
             max_length,
             top_p,
@@ -102,22 +104,14 @@ def create_ui():
             chatbot,
             input_message
         ])
-
-        revoke_btn.click(lambda x: ctx.revoke(), inputs=[input_message], outputs=[chatbot])
-
-        clear.click(clear_history, outputs=[chatbot])
+        revoke_btn.click(lambda ctx: ctx.revoke(), inputs=[state], outputs=[chatbot])
+        clear_history_btn.click(clear_history, inputs=[state], outputs=[chatbot])
         clear_input.click(lambda x: "", inputs=[input_message], outputs=[input_message])
-
-        save_his_btn.click(ctx.save_history, outputs=[cmd_output])
-        save_md_btn.click(ctx.save_as_md, outputs=[cmd_output])
-        load_his_btn.upload(ctx.load_history, inputs=[
-            load_his_btn,
-        ], outputs=[
-            chatbot
-        ])
-        sync_his_btn.click(lambda x: ctx.rh, inputs=[input_message], outputs=[chatbot])
-
-        apply_max_rounds.click(apply_max_round_click, inputs=[max_rounds], outputs=[cmd_output])
+        save_his_btn.click(lambda ctx: ctx.save_history(), inputs=[state], outputs=[cmd_output])
+        save_md_btn.click(lambda ctx: ctx.save_as_md(), inputs=[state], outputs=[cmd_output])
+        load_his_btn.upload(lambda ctx, f: ctx.load_history(f), inputs=[state, load_his_btn], outputs=[chatbot])
+        sync_his_btn.click(lambda ctx: ctx.rh, inputs=[state], outputs=[chatbot])
+        apply_max_rounds.click(apply_max_round_click, inputs=[state, max_rounds], outputs=[cmd_output])
 
     with gr.Blocks(css=css, analytics_enabled=False) as settings_interface:
         with gr.Row():