Upload app.py

app.py

@@ -2,12 +2,17 @@ import hashlib
 import os
 import re
 import json
+import threading
 
 import gradio as gr
 
 from chat_template import ChatTemplate
 from llama_cpp_python_streamingllm import StreamingLLM
 
+#  ========== 全局锁，确保只能进行一个会话 ==========
+lock = threading.Lock()
+session_active = False
+
 #  ========== 让聊天界面的文本框等高 ==========
 custom_css = r'''
 #area > div {
@@ -208,6 +213,9 @@ def btn_submit_com(_n_keep, _n_discard,
                    _top_p, _min_p, _typical_p,
                    _tfs_z, _mirostat_mode, _mirostat_eta,
                    _mirostat_tau, _role, _max_tokens):
+    with lock:
+        if not session_active:
+            raise RuntimeError
     # ========== 初始化输出模版 ==========
     t_bot = chat_template(_role)
     completion_tokens = []  # 有可能多个 tokens 才能构成一个 utf-8 编码的文字
@@ -267,10 +275,15 @@ def btn_submit_com(_n_keep, _n_discard,
 
 # ========== 显示用户消息 ==========
 def btn_submit_usr(message: str, history):
-    # print('btn_submit_usr', message, history)
-    if history is None:
-        history = []
-    return "", history + [[message.strip(), '']], gr.update(interactive=False)
+    global session_active
+    with lock:
+        if session_active:
+            raise RuntimeError
+        session_active = True
+        # print('btn_submit_usr', message, history)
+        if history is None:
+            history = []
+        return "", history + [[message.strip(), '']], gr.update(interactive=False)
 
 
 # ========== 模型流式响应 ==========
@@ -281,6 +294,9 @@ def btn_submit_bot(history, _n_keep, _n_discard,
                    _tfs_z, _mirostat_mode, _mirostat_eta,
                    _mirostat_tau, _usr, _char,
                    _rag, _max_tokens):
+    with lock:
+        if not session_active:
+            raise RuntimeError
     # ========== 需要临时注入的内容 ==========
     rag_idx = None
     if len(_rag) > 0:
@@ -336,6 +352,9 @@ def btn_submit_vo(_n_keep, _n_discard,
                   _top_p, _min_p, _typical_p,
                   _tfs_z, _mirostat_mode, _mirostat_eta,
                   _mirostat_tau, _max_tokens):
+    with lock:
+        if not session_active:
+            raise RuntimeError
     global vo_idx
     vo_idx = model.venv_create()  # 创建隔离环境
     # ========== 模型输出旁白 ==========
@@ -356,6 +375,9 @@ def btn_submit_suggest(_n_keep, _n_discard,
                        _top_p, _min_p, _typical_p,
                        _tfs_z, _mirostat_mode, _mirostat_eta,
                        _mirostat_tau, _usr, _max_tokens):
+    with lock:
+        if not session_active:
+            raise RuntimeError
     model.venv_create()  # 创建隔离环境
     # ========== 模型输出 ==========
     _tmp = btn_submit_com(_n_keep, _n_discard,
@@ -371,6 +393,15 @@ def btn_submit_suggest(_n_keep, _n_discard,
     yield _h, str((model.n_tokens, model.venv))
 
 
+def btn_submit_finish():
+    global session_active
+    with lock:
+        if not session_active:
+            raise RuntimeError
+        session_active = False
+        return gr.update(interactive=True)
+
+
 # ========== 聊天页面 ==========
 with gr.Blocks() as chatting:
     with gr.Row(equal_height=True):
@@ -396,7 +427,7 @@ with gr.Blocks() as chatting:
         fn=btn_submit_usr, api_name="submit",
         inputs=[msg, chatbot],
         outputs=[msg, chatbot, btn_submit]
-    ).then(
+    ).success(
         fn=btn_submit_bot,
         inputs=[chatbot, setting_n_keep, setting_n_discard,
                 setting_temperature, setting_repeat_penalty, setting_frequency_penalty,
@@ -406,7 +437,7 @@ with gr.Blocks() as chatting:
                 setting_mirostat_tau, role_usr, role_char,
                 rag, setting_max_tokens],
         outputs=[chatbot, s_info]
-    ).then(
+    ).success(
         fn=btn_submit_vo,
         inputs=[setting_n_keep, setting_n_discard,
                 setting_temperature, setting_repeat_penalty, setting_frequency_penalty,
@@ -415,7 +446,7 @@ with gr.Blocks() as chatting:
                 setting_tfs_z, setting_mirostat_mode, setting_mirostat_eta,
                 setting_mirostat_tau, setting_max_tokens],
         outputs=[vo, s_info]
-    ).then(
+    ).success(
         fn=btn_submit_suggest,
         inputs=[setting_n_keep, setting_n_discard,
                 setting_temperature, setting_repeat_penalty, setting_frequency_penalty,
@@ -424,28 +455,31 @@ with gr.Blocks() as chatting:
                 setting_tfs_z, setting_mirostat_mode, setting_mirostat_eta,
                 setting_mirostat_tau, role_usr, setting_max_tokens],
         outputs=[msg, s_info]
-    ).then(
-        fn=lambda: gr.update(interactive=True),
+    ).success(
+        fn=btn_submit_finish,
         outputs=btn_submit
     )
 
     # ========== 用于调试 ==========
-    btn_com1.click(fn=lambda: model.str_detokenize(model._input_ids), outputs=rag)
+    # btn_com1.click(fn=lambda: model.str_detokenize(model._input_ids), outputs=rag)
 
 
     @btn_com2.click(inputs=setting_cache_path,
-                    outputs=s_info)
+                    outputs=[s_info, btn_submit])
     def btn_com2(_cache_path):
-        _tmp = model.load_session(setting_cache_path.value)
-        print(f'load cache from {setting_cache_path.value} {_tmp}')
-        global vo_idx
-        vo_idx = 0
-        model.venv = [0]
-        return str((model.n_tokens, model.venv))
+        with lock:
+            _tmp = model.load_session(setting_cache_path.value)
+            print(f'load cache from {setting_cache_path.value} {_tmp}')
+            global vo_idx
+            vo_idx = 0
+            model.venv = [0]
+            global session_active
+            session_active = False
+            return str((model.n_tokens, model.venv)), gr.update(interactive=True)
 
     # ========== 开始运行 ==========
 demo = gr.TabbedInterface([chatting, setting, role],
                           ["聊天", "设置", '角色'],
                           css=custom_css)
 gr.close_all()
-demo.queue().launch(share=False)
+demo.queue(max_size=1).launch(max_threads=1, share=False)