Limour/llama-python-streamingllm

1import gradio as gr2import threading3from llama_cpp_python_streamingllm import StreamingLLM4from mods.read_cfg import cfg5 6from mods.text_display import init as text_display_init7 8from mods.btn_rag import init as btn_rag_init9 10# ========== 按钮中用到的共同的函数 ==========11from mods.btn_com import init as btn_com_init12 13# ========== 输出一段回答 ==========14from mods.btn_submit import init as btn_submit_init15 16# ========== 输出一段旁白 ==========17from mods.btn_vo import init as btn_vo_init18 19# ========== 重新输出一段回答 ==========20from mods.btn_retry import init as btn_retry_init21 22# ========== 给用户提供默认回复的建议 ==========23from mods.btn_suggest import init as btn_suggest_init24 25# ========== 融合功能的按钮 ==========26from mods.btn_submit_vo_suggest import init as btn_submit_vo_suggest_init27 28# ========== 更新状态栏的按钮 ==========29from mods.btn_status_bar import init as btn_status_bar_init30 31# ========== 重置按钮 ==========32from mods.btn_reset import init as btn_reset_init33 34# ========== 聊天的模版 默认 chatml ==========35from chat_template import ChatTemplate36 37# ========== 加载角色卡-缓存 ==========38from mods.load_cache import init as load_cache_init39 40#  ========== 全局锁，确保只能进行一个会话 ==========41cfg['session_lock'] = threading.Lock()42cfg['session_active'] = False43 44#  ========== 温度、采样之类的设置 ==========45with gr.Blocks() as setting:46    with gr.Row():47        cfg['setting_path'] = gr.Textbox(label="模型路径", max_lines=1, scale=2, **cfg['setting_path'])48        cfg['setting_cache_path'] = gr.Textbox(label="缓存路径", max_lines=1, scale=2, **cfg['setting_cache_path'])49        cfg['setting_seed'] = gr.Number(label="随机种子", scale=1, **cfg['setting_seed'])50        cfg['setting_n_gpu_layers'] = gr.Number(label="n_gpu_layers", scale=1, **cfg['setting_n_gpu_layers'])51    with gr.Row(elem_classes='setting'):52        cfg['setting_ctx'] = gr.Number(label="上下文大小（Tokens）", **cfg['setting_ctx'])53        cfg['setting_max_tokens'] = gr.Number(label="最大响应长度（Tokens）", interactive=True,54                                              **cfg['setting_max_tokens'])55        cfg['setting_n_keep'] = gr.Number(value=10, label="n_keep", interactive=False)56        cfg['setting_n_discard'] = gr.Number(label="n_discard", interactive=True, **cfg['setting_n_discard'])57    with gr.Row(elem_classes='setting'):58        cfg['setting_temperature'] = gr.Number(label="温度", interactive=True, **cfg['setting_temperature'])59        cfg['setting_repeat_penalty'] = gr.Number(label="重复惩罚", interactive=True, **cfg['setting_repeat_penalty'])60        cfg['setting_frequency_penalty'] = gr.Number(label="频率惩罚", interactive=True,61                                                     **cfg['setting_frequency_penalty'])62        cfg['setting_presence_penalty'] = gr.Number(label="存在惩罚", interactive=True,63                                                    **cfg['setting_presence_penalty'])64        cfg['setting_repeat_last_n'] = gr.Number(label="惩罚范围", interactive=True, **cfg['setting_repeat_last_n'])65    with gr.Row(elem_classes='setting'):66        cfg['setting_top_k'] = gr.Number(label="Top-K", interactive=True, **cfg['setting_top_k'])67        cfg['setting_top_p'] = gr.Number(label="Top P", interactive=True, **cfg['setting_top_p'])68        cfg['setting_min_p'] = gr.Number(label="Min P", interactive=True, **cfg['setting_min_p'])69        cfg['setting_typical_p'] = gr.Number(label="Typical", interactive=True, **cfg['setting_typical_p'])70        cfg['setting_tfs_z'] = gr.Number(label="TFS", interactive=True, **cfg['setting_tfs_z'])71    with gr.Row(elem_classes='setting'):72        cfg['setting_mirostat_mode'] = gr.Number(label="Mirostat 模式", **cfg['setting_mirostat_mode'])73        cfg['setting_mirostat_eta'] = gr.Number(label="Mirostat 学习率", interactive=True,74                                                **cfg['setting_mirostat_eta'])75        cfg['setting_mirostat_tau'] = gr.Number(label="Mirostat 目标熵", interactive=True,76                                                **cfg['setting_mirostat_tau'])77    with gr.Row(elem_classes='setting'):78        cfg['setting_btn_vo_keep_last'] = gr.Number(label="旁白限制", interactive=True,79                                                    **cfg['setting_btn_vo_keep_last'])80 81    #  ========== 加载模型 ==========82    cfg['model'] = StreamingLLM(model_path=cfg['setting_path'].value,83                                seed=cfg['setting_seed'].value,84                                n_gpu_layers=cfg['setting_n_gpu_layers'].value,85                                n_ctx=cfg['setting_ctx'].value)86    cfg['chat_template'] = ChatTemplate(cfg['model'])87    cfg['setting_ctx'].value = cfg['model'].n_ctx()88 89# ========== 展示角色卡 ==========90with gr.Blocks() as role:91    with gr.Row():92        cfg['role_usr'] = gr.Textbox(label="用户名称", max_lines=1, interactive=False, **cfg['role_usr'])93        cfg['role_char'] = gr.Textbox(label="角色名称", max_lines=1, interactive=False, **cfg['role_char'])94 95    cfg['role_char_d'] = gr.Textbox(lines=10, label="故事描述", **cfg['role_char_d'])96    cfg['role_chat_style'] = gr.Textbox(lines=10, label="回复示例", **cfg['role_chat_style'])97 98    # ========== 加载角色卡-缓存 ==========99    text_display_init(cfg)100    load_cache_init(cfg)101 102# ========== 聊天页面 ==========103with gr.Blocks() as chatting:104    with gr.Row(equal_height=True):105        cfg['chatbot'] = gr.Chatbot(height='60vh', scale=2, value=cfg['chatbot'],106                                    avatar_images=(r'assets/user.png', r'assets/chatbot.webp'))107        with gr.Column(scale=1):108            with gr.Tab(label='Main', elem_id='area'):109                cfg['rag'] = gr.Textbox(label='RAG', lines=2, show_copy_button=True, elem_classes="area")110                cfg['vo'] = gr.Textbox(label='VO', lines=2, show_copy_button=True, elem_classes="area")111                cfg['s_info'] = gr.Textbox(value=cfg['model'].venv_info, max_lines=1, label='info', interactive=False)112            with gr.Tab(label='状态栏', elem_id='area'):113                cfg['status_bar'] = gr.Dataframe(114                    headers=['属性', '值'],115                    type="array",116                    elem_id='StatusBar'117                )118    cfg['msg'] = gr.Textbox(label='Prompt', lines=2, max_lines=2, elem_id='prompt', autofocus=True, **cfg['msg'])119 120    cfg['gr'] = gr121    btn_com_init(cfg)122 123    btn_rag_init(cfg)124 125    btn_submit_init(cfg)126 127    btn_vo_init(cfg)128 129    btn_suggest_init(cfg)130 131    btn_retry_init(cfg)132 133    btn_submit_vo_suggest_init(cfg)134 135    btn_status_bar_init(cfg)136 137    # ========== 用于调试 ==========138    btn_reset_init(cfg)139 140#  ========== 让聊天界面的文本框等高 ==========141custom_css = r'''142#area > div > div {143    height: 53vh;144}145.area {146    flex-grow: 1;147}148.area > label {149    height: 100%;150    display: flex;151    flex-direction: column;152    max-height: 16vh;153}154.area > label > textarea {155    flex-grow: 1;156}157#prompt > label > textarea {158    max-height: 63px;159}160.setting label {161    display: flex;162    flex-direction: column;163    height: 100%;164}165.setting input {166    margin-top: auto;167}168#StatusBar {169    max-height: 53vh;170}171'''172 173# ========== 开始运行 ==========174demo = gr.TabbedInterface([chatting, setting, role],175                          ["聊天", "设置", '角色'],176                          css=custom_css)177gr.close_all()178demo.queue(api_open=False, max_size=1).launch(share=False, show_error=True, show_api=False)179