Create app.py

app.py

@@ -0,0 +1,436 @@
+import hashlib
+import os
+import re
+import json
+
+import gradio as gr
+
+from chat_template import ChatTemplate
+from llama_cpp_python_streamingllm import StreamingLLM
+
+#  ========== 让聊天界面的文本框等高 ==========
+custom_css = r'''
+#area > div {
+    height: 100%;
+}
+#RAG-area {
+    flex-grow: 1;
+}
+#RAG-area > label {
+    height: 100%;
+    display: flex;
+    flex-direction: column;
+}
+#RAG-area > label > textarea {
+    flex-grow: 1;
+    max-height: 20vh;
+}
+#VO-area {
+    flex-grow: 1;
+}
+#VO-area > label {
+    height: 100%;
+    display: flex;
+    flex-direction: column;
+}
+#VO-area > label > textarea {
+    flex-grow: 1;
+    max-height: 20vh;
+}
+#prompt > label > textarea {
+    max-height: 63px;
+}
+'''
+
+
+#  ========== 适配 SillyTavern 的模版 ==========
+def text_format(text: str, _env=None, **env):
+    if _env is not None:
+        for k, v in _env.items():
+            text = text.replace(r'{{' + k + r'}}', v)
+    for k, v in env.items():
+        text = text.replace(r'{{' + k + r'}}', v)
+    return text
+
+
+# ========== 哈希函数 ==========
+def x_hash(x: str):
+    return hashlib.sha1(x.encode('utf-8')).hexdigest()
+
+
+# ========== 读取配置文件 ==========
+with open('rp_config.json', encoding='utf-8') as f:
+    tmp = f.read()
+with open('rp_sample_config.json', encoding='utf-8') as f:
+    cfg = json.load(f)
+cfg['setting_cache_path']['value'] += x_hash(tmp)
+cfg.update(json.loads(tmp))
+
+#  ========== 给引号加粗 ==========
+reg_q = re.compile(r'“(.+?)”')
+
+
+def chat_display_format(text: str):
+    return reg_q.sub(r' **\g<0>** ', text)
+
+
+#  ========== 温度、采样之类的设置 ==========
+with gr.Blocks() as setting:
+    with gr.Row():
+        setting_path = gr.Textbox(label="模型路径", max_lines=1, scale=2, **cfg['setting_path'])
+        setting_cache_path = gr.Textbox(label="缓存路径", max_lines=1, scale=2, **cfg['setting_cache_path'])
+        setting_seed = gr.Number(label="随机种子", scale=1, **cfg['setting_seed'])
+        setting_n_gpu_layers = gr.Number(label="n_gpu_layers", scale=1, **cfg['setting_n_gpu_layers'])
+    with gr.Row():
+        setting_ctx = gr.Number(label="上下文大小（Tokens）", **cfg['setting_ctx'])
+        setting_max_tokens = gr.Number(label="最大响应长度（Tokens）", interactive=True, **cfg['setting_max_tokens'])
+        setting_n_keep = gr.Number(value=10, label="n_keep", interactive=False)
+        setting_n_discard = gr.Number(label="n_discard", interactive=True, **cfg['setting_n_discard'])
+    with gr.Row():
+        setting_temperature = gr.Number(label="温度", interactive=True, **cfg['setting_temperature'])
+        setting_repeat_penalty = gr.Number(label="重复惩罚", interactive=True, **cfg['setting_repeat_penalty'])
+        setting_frequency_penalty = gr.Number(label="频率惩罚", interactive=True, **cfg['setting_frequency_penalty'])
+        setting_presence_penalty = gr.Number(label="存在惩罚", interactive=True, **cfg['setting_presence_penalty'])
+        setting_repeat_last_n = gr.Number(label="惩罚范围", interactive=True, **cfg['setting_repeat_last_n'])
+    with gr.Row():
+        setting_top_k = gr.Number(label="Top-K", interactive=True, **cfg['setting_top_k'])
+        setting_top_p = gr.Number(label="Top P", interactive=True, **cfg['setting_top_p'])
+        setting_min_p = gr.Number(label="Min P", interactive=True, **cfg['setting_min_p'])
+        setting_typical_p = gr.Number(label="Typical", interactive=True, **cfg['setting_typical_p'])
+        setting_tfs_z = gr.Number(label="TFS", interactive=True, **cfg['setting_tfs_z'])
+    with gr.Row():
+        setting_mirostat_mode = gr.Number(label="Mirostat 模式", **cfg['setting_mirostat_mode'])
+        setting_mirostat_eta = gr.Number(label="Mirostat 学习率", interactive=True, **cfg['setting_mirostat_eta'])
+        setting_mirostat_tau = gr.Number(label="Mirostat 目标熵", interactive=True, **cfg['setting_mirostat_tau'])
+
+    #  ========== 加载模型 ==========
+    model = StreamingLLM(model_path=setting_path.value,
+                         seed=setting_seed.value,
+                         n_gpu_layers=setting_n_gpu_layers.value,
+                         n_ctx=setting_ctx.value)
+    setting_ctx.value = model.n_ctx()
+
+# ========== 聊天的模版 默认 chatml ==========
+chat_template = ChatTemplate(model)
+
+# ========== 展示角色卡 ==========
+with gr.Blocks() as role:
+    with gr.Row():
+        role_usr = gr.Textbox(label="用户名称", max_lines=1, interactive=False, **cfg['role_usr'])
+        role_char = gr.Textbox(label="角色名称", max_lines=1, interactive=False, **cfg['role_char'])
+
+    role_char_d = gr.Textbox(lines=10, label="故事描述", **cfg['role_char_d'])
+    role_chat_style = gr.Textbox(lines=10, label="回复示例", **cfg['role_chat_style'])
+
+    # model.eval_t([1])  # 这个暖机的 bos [1] 删了就不正常了
+    if os.path.exists(setting_cache_path.value):
+        # ========== 加载角色卡-缓存 ==========
+        tmp = model.load_session(setting_cache_path.value)
+        print(f'load cache from {setting_cache_path.value} {tmp}')
+        tmp = chat_template('system',
+                            text_format(role_char_d.value,
+                                        char=role_char.value,
+                                        user=role_usr.value))
+        setting_n_keep.value = len(tmp)
+        tmp = chat_template(role_char.value,
+                            text_format(role_chat_style.value,
+                                        char=role_char.value,
+                                        user=role_usr.value))
+        setting_n_keep.value += len(tmp)
+        # ========== 加载角色卡-第一条消息 ==========
+        chatbot = []
+        for one in cfg["role_char_first"]:
+            one['name'] = text_format(one['name'],
+                                      char=role_char.value,
+                                      user=role_usr.value)
+            one['value'] = text_format(one['value'],
+                                       char=role_char.value,
+                                       user=role_usr.value)
+            if one['name'] == role_char.value:
+                chatbot.append((None, chat_display_format(one['value'])))
+            print(one)
+    else:
+        # ========== 加载角色卡-角色描述 ==========
+        tmp = chat_template('system',
+                            text_format(role_char_d.value,
+                                        char=role_char.value,
+                                        user=role_usr.value))
+        setting_n_keep.value = model.eval_t(tmp)  # 此内容永久存在
+
+        # ========== 加载角色卡-回复示例 ==========
+        tmp = chat_template(role_char.value,
+                            text_format(role_chat_style.value,
+                                        char=role_char.value,
+                                        user=role_usr.value))
+        setting_n_keep.value = model.eval_t(tmp)  # 此内容永久存在
+
+        # ========== 加载角色卡-第一条消息 ==========
+        chatbot = []
+        for one in cfg["role_char_first"]:
+            one['name'] = text_format(one['name'],
+                                      char=role_char.value,
+                                      user=role_usr.value)
+            one['value'] = text_format(one['value'],
+                                       char=role_char.value,
+                                       user=role_usr.value)
+            if one['name'] == role_char.value:
+                chatbot.append((None, chat_display_format(one['value'])))
+            print(one)
+            tmp = chat_template(one['name'], one['value'])
+            model.eval_t(tmp)  # 此内容随上下文增加将被丢弃
+
+        # ========== 保存角色卡-缓存 ==========
+        with open(setting_cache_path.value, 'wb') as f:
+            pass
+        tmp = model.save_session(setting_cache_path.value)
+        print(f'save cache {tmp}')
+
+
+# ========== 流式输出函数 ==========
+def btn_submit_com(_n_keep, _n_discard,
+                   _temperature, _repeat_penalty, _frequency_penalty,
+                   _presence_penalty, _repeat_last_n, _top_k,
+                   _top_p, _min_p, _typical_p,
+                   _tfs_z, _mirostat_mode, _mirostat_eta,
+                   _mirostat_tau, _role, _max_tokens):
+    # ========== 初始化输出模版 ==========
+    t_bot = chat_template(_role)
+    completion_tokens = []  # 有可能多个 tokens 才能构成一个 utf-8 编码的文字
+    history = ''
+    # ========== 流式输出 ==========
+    for token in model.generate_t(
+            tokens=t_bot,
+            n_keep=_n_keep,
+            n_discard=_n_discard,
+            im_start=chat_template.im_start_token,
+            top_k=_top_k,
+            top_p=_top_p,
+            min_p=_min_p,
+            typical_p=_typical_p,
+            temp=_temperature,
+            repeat_penalty=_repeat_penalty,
+            repeat_last_n=_repeat_last_n,
+            frequency_penalty=_frequency_penalty,
+            presence_penalty=_presence_penalty,
+            tfs_z=_tfs_z,
+            mirostat_mode=_mirostat_mode,
+            mirostat_tau=_mirostat_tau,
+            mirostat_eta=_mirostat_eta,
+    ):
+        if token in chat_template.eos or token == chat_template.nlnl:
+            t_bot.extend(completion_tokens)
+            print('token in eos', token)
+            break
+        completion_tokens.append(token)
+        all_text = model.str_detokenize(completion_tokens)
+        if not all_text:
+            continue
+        t_bot.extend(completion_tokens)
+        history += all_text
+        yield history
+        if token in chat_template.onenl:
+            # ========== 移除末尾的换行符 ==========
+            if t_bot[-2] in chat_template.onenl:
+                model.venv_pop_token()
+                break
+            if t_bot[-2] in chat_template.onerl and t_bot[-3] in chat_template.onenl:
+                model.venv_pop_token()
+                break
+        if history[-2:] == '\n\n':  # 各种 'x\n\n' 的token，比如'。\n\n'
+            print('t_bot[-4:]', t_bot[-4:], repr(model.str_detokenize(t_bot[-4:])),
+                  repr(model.str_detokenize(t_bot[-1:])))
+            break
+        if len(t_bot) > _max_tokens:
+            break
+        completion_tokens = []
+    # ========== 查看末尾的换行符 ==========
+    print('history', repr(history))
+    # ========== 给 kv_cache 加上输出结束符 ==========
+    model.eval_t(chat_template.im_end_nl, _n_keep, _n_discard)
+    t_bot.extend(chat_template.im_end_nl)
+
+
+# ========== 显示用户消息 ==========
+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)
+
+
+# ========== 模型流式响应 ==========
+def btn_submit_bot(history, _n_keep, _n_discard,
+                   _temperature, _repeat_penalty, _frequency_penalty,
+                   _presence_penalty, _repeat_last_n, _top_k,
+                   _top_p, _min_p, _typical_p,
+                   _tfs_z, _mirostat_mode, _mirostat_eta,
+                   _mirostat_tau, _usr, _char,
+                   _rag, _max_tokens):
+    # ========== 需要临时注入的内容 ==========
+    rag_idx = None
+    if len(_rag) > 0:
+        rag_idx = model.venv_create()  # 记录 venv_idx
+        t_rag = chat_template('system', _rag)
+        model.eval_t(t_rag, _n_keep, _n_discard)
+    model.venv_create()  # 与 t_rag 隔离
+    # ========== 用户输入 ==========
+    t_msg = history[-1][0]
+    t_msg = chat_template(_usr, t_msg)
+    model.eval_t(t_msg, _n_keep, _n_discard)
+    # ========== 模型输出 ==========
+    _tmp = btn_submit_com(_n_keep, _n_discard,
+                          _temperature, _repeat_penalty, _frequency_penalty,
+                          _presence_penalty, _repeat_last_n, _top_k,
+                          _top_p, _min_p, _typical_p,
+                          _tfs_z, _mirostat_mode, _mirostat_eta,
+                          _mirostat_tau, _char, _max_tokens)
+    for _h in _tmp:
+        history[-1][1] = _h
+        yield history, str((model.n_tokens, model.venv))
+    # ========== 输出完毕后格式化输出 ==========
+    history[-1][1] = chat_display_format(history[-1][1])
+    yield history, str((model.n_tokens, model.venv))
+    # ========== 及时清理上一次生成的旁白 ==========
+    if vo_idx > 0:
+        print('vo_idx', vo_idx, model.venv)
+        model.venv_remove(vo_idx)
+        print('vo_idx', vo_idx, model.venv)
+        if rag_idx and vo_idx < rag_idx:
+            rag_idx -= 1
+    # ========== 响应完毕后清除注入的内容 ==========
+    if rag_idx is not None:
+        model.venv_remove(rag_idx)  # 销毁对应的 venv
+    model.venv_disband()  # 退出隔离环境
+    yield history, str((model.n_tokens, model.venv))
+    print('venv_disband', vo_idx, model.venv)
+
+
+# ========== 待实现 ==========
+def btn_rag_(_rag, _msg):
+    retn = ''
+    return retn
+
+
+vo_idx = 0
+
+
+# ========== 输出一段旁白 ==========
+def btn_submit_vo(_n_keep, _n_discard,
+                  _temperature, _repeat_penalty, _frequency_penalty,
+                  _presence_penalty, _repeat_last_n, _top_k,
+                  _top_p, _min_p, _typical_p,
+                  _tfs_z, _mirostat_mode, _mirostat_eta,
+                  _mirostat_tau, _max_tokens):
+    global vo_idx
+    vo_idx = model.venv_create()  # 创建隔离环境
+    # ========== 模型输出旁白 ==========
+    _tmp = btn_submit_com(_n_keep, _n_discard,
+                          _temperature, _repeat_penalty, _frequency_penalty,
+                          _presence_penalty, _repeat_last_n, _top_k,
+                          _top_p, _min_p, _typical_p,
+                          _tfs_z, _mirostat_mode, _mirostat_eta,
+                          _mirostat_tau, '旁白', _max_tokens)
+    for _h in _tmp:
+        yield _h, str((model.n_tokens, model.venv))
+
+
+# ========== 给用户提供默认回复 ==========
+def btn_submit_suggest(_n_keep, _n_discard,
+                       _temperature, _repeat_penalty, _frequency_penalty,
+                       _presence_penalty, _repeat_last_n, _top_k,
+                       _top_p, _min_p, _typical_p,
+                       _tfs_z, _mirostat_mode, _mirostat_eta,
+                       _mirostat_tau, _usr, _max_tokens):
+    model.venv_create()  # 创建隔离环境
+    # ========== 模型输出 ==========
+    _tmp = btn_submit_com(_n_keep, _n_discard,
+                          _temperature, _repeat_penalty, _frequency_penalty,
+                          _presence_penalty, _repeat_last_n, _top_k,
+                          _top_p, _min_p, _typical_p,
+                          _tfs_z, _mirostat_mode, _mirostat_eta,
+                          _mirostat_tau, _usr, _max_tokens)
+    _h = ''
+    for _h in _tmp:
+        yield _h, str((model.n_tokens, model.venv))
+    model.venv_remove()  # 销毁隔离环境
+    yield _h, str((model.n_tokens, model.venv))
+
+
+# ========== 聊天页面 ==========
+with gr.Blocks() as chatting:
+    with gr.Row(equal_height=True):
+        chatbot = gr.Chatbot(height='60vh', scale=2, value=chatbot,
+                             avatar_images=(r'assets/user.png', r'assets/chatbot.webp'))
+        with gr.Column(scale=1, elem_id="area"):
+            rag = gr.Textbox(label='RAG', show_copy_button=True, elem_id="RAG-area")
+            vo = gr.Textbox(label='VO', show_copy_button=True, elem_id="VO-area")
+            s_info = gr.Textbox(value=str((model.n_tokens, model.venv)), max_lines=1, label='info', interactive=False)
+    msg = gr.Textbox(label='Prompt', lines=2, max_lines=2, elem_id='prompt', autofocus=True, **cfg['msg'])
+    with gr.Row():
+        btn_rag = gr.Button("RAG")
+        btn_submit = gr.Button("Submit")
+        btn_retry = gr.Button("Retry")
+        btn_com1 = gr.Button("自定义1")
+        btn_com2 = gr.Button("自定义2")
+        btn_com3 = gr.Button("自定义3")
+
+    btn_rag.click(fn=btn_rag_, outputs=rag,
+                  inputs=[rag, msg])
+
+    btn_submit.click(
+        fn=btn_submit_usr, api_name="submit",
+        inputs=[msg, chatbot],
+        outputs=[msg, chatbot, btn_submit]
+    ).then(
+        fn=btn_submit_bot,
+        inputs=[chatbot, setting_n_keep, setting_n_discard,
+                setting_temperature, setting_repeat_penalty, setting_frequency_penalty,
+                setting_presence_penalty, setting_repeat_last_n, setting_top_k,
+                setting_top_p, setting_min_p, setting_typical_p,
+                setting_tfs_z, setting_mirostat_mode, setting_mirostat_eta,
+                setting_mirostat_tau, role_usr, role_char,
+                rag, setting_max_tokens],
+        outputs=[chatbot, s_info]
+    ).then(
+        fn=btn_submit_vo,
+        inputs=[setting_n_keep, setting_n_discard,
+                setting_temperature, setting_repeat_penalty, setting_frequency_penalty,
+                setting_presence_penalty, setting_repeat_last_n, setting_top_k,
+                setting_top_p, setting_min_p, setting_typical_p,
+                setting_tfs_z, setting_mirostat_mode, setting_mirostat_eta,
+                setting_mirostat_tau, setting_max_tokens],
+        outputs=[vo, s_info]
+    ).then(
+        fn=btn_submit_suggest,
+        inputs=[setting_n_keep, setting_n_discard,
+                setting_temperature, setting_repeat_penalty, setting_frequency_penalty,
+                setting_presence_penalty, setting_repeat_last_n, setting_top_k,
+                setting_top_p, setting_min_p, setting_typical_p,
+                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),
+        outputs=btn_submit
+    )
+
+    # ========== 用于调试 ==========
+    btn_com1.click(fn=lambda: model.str_detokenize(model._input_ids), outputs=rag)
+
+
+    @btn_com2.click(inputs=setting_cache_path,
+                    outputs=s_info)
+    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))
+
+    # ========== 开始运行 ==========
+demo = gr.TabbedInterface([chatting, setting, role],
+                          ["聊天", "设置", '角色'],
+                          css=custom_css)
+gr.close_all()
+demo.queue().launch(share=False)