app.py

# -*- coding: utf-8 -*-
"""
@author:XuMing(xuming624@qq.com)
@description:
modified from https://github.com/imClumsyPanda/langchain-ChatGLM/blob/master/webui.py
"""
import gradio as gr
import os
import shutil
from loguru import logger
from chatpdf import ChatPDF

pwd_path = os.path.abspath(os.path.dirname(__file__))

CONTENT_DIR = os.path.join(pwd_path, "content")
logger.info(f"CONTENT_DIR: {CONTENT_DIR}")
VECTOR_SEARCH_TOP_K = 3
MAX_INPUT_LEN = 512

embedding_model_dict = {
    "sentence-transformers": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
    "ernie-tiny": "nghuyong/ernie-3.0-nano-zh",
    "ernie-base": "nghuyong/ernie-3.0-base-zh",
    "text2vec": "shibing624/text2vec-base-chinese",
}

# supported LLM models
llm_model_dict = {
    "chatglm-6b-int4": "THUDM/chatglm-6b-int4",
    "chatglm-6b-int4-qe": "THUDM/chatglm-6b-int4-qe",
    "chatglm-6b": "THUDM/chatglm-6b",
    "llama-7b": "decapoda-research/llama-7b-hf",
    "llama-13b": "decapoda-research/llama-13b-hf",
}

llm_model_dict_list = list(llm_model_dict.keys())
embedding_model_dict_list = list(embedding_model_dict.keys())

model = ChatPDF(
    sim_model_name_or_path="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
    gen_model_type="chatglm",
    gen_model_name_or_path="THUDM/chatglm-6b-int4",
    lora_model_name_or_path=None,
    max_input_size=MAX_INPUT_LEN,
)


def get_file_list():
    if not os.path.exists("content"):
        return []
    return [f for f in os.listdir("content") if
            f.endswith(".txt") or f.endswith(".pdf") or f.endswith(".docx") or f.endswith(".md")]


file_list = get_file_list()


def upload_file(file):
    if not os.path.exists(CONTENT_DIR):
        os.mkdir(CONTENT_DIR)
    filename = os.path.basename(file.name)
    shutil.move(file.name, os.path.join(CONTENT_DIR, filename))
    # file_list首位插入新上传的文件
    file_list.insert(0, filename)
    return gr.Dropdown.update(choices=file_list, value=filename)


def get_answer(query, index_path, history):
    if index_path:
        if not model.sim_model.corpus_embeddings:
            model.load_index(index_path)
        response, empty_history = model.query(query, topn=VECTOR_SEARCH_TOP_K)
        history = history + [[query, response]]
    else:
        # history = history + [[None, "请先加载文件后，再进行提问。"]]
        # 未加载文件，仅返回生成模型结果
        response, empty_history = model.gen_model.chat(query)
        history = history + [[query, response]]
    logger.debug(f"query: {query}, response: {response}")
    return history, ""


def update_status(history, status):
    history = history + [[None, status]]
    logger.info(status)
    return history


def reinit_model(llm_model, embedding_model, history):
    try:
        global model
        del model
        model = ChatPDF(
            sim_model_name_or_path=embedding_model_dict.get(
                embedding_model,
                "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"),
            gen_model_type=llm_model.split('-')[0],
            gen_model_name_or_path=llm_model_dict.get(llm_model, "THUDM/chatglm-6b-int4"),
            lora_model_name_or_path=None,
            max_input_size=MAX_INPUT_LEN,
        )
        model_status = """模型已成功重新加载，请选择文件后点击"加载文件"按钮"""
    except Exception as e:
        model = None
        logger.error(e)
        model_status = """模型未成功重新加载，请重新选择后点击"加载模型"按钮"""
    return history + [[None, model_status]]


def get_vector_store(filepath, history):
    logger.info(filepath, history)
    index_path = None
    file_status = ''
    if model is not None:
        local_file_path = os.path.join(CONTENT_DIR, filepath)
        local_index_path = os.path.join(CONTENT_DIR, filepath + ".index.json")
        if os.path.exists(local_file_path):
            model.load_pdf_file(local_file_path)
            model.save_index(local_index_path)
            index_path = local_index_path
            if index_path:
                file_status = "文件已成功加载，请开始提问"
            else:
                file_status = "文件未成功加载，请重新上传文件"
    else:
        file_status = "模型未完成加载，请先在加载模型后再导入文件"

    return index_path, history + [[None, file_status]]


def reset_chat(chatbot, state):
    return None, None


block_css = """.importantButton {
    background: linear-gradient(45deg, #7e0570,#5d1c99, #6e00ff) !important;
    border: none !important;
}
.importantButton:hover {
    background: linear-gradient(45deg, #ff00e0,#8500ff, #6e00ff) !important;
    border: none !important;
}"""

webui_title = """
# 🎉ChatPDF WebUI🎉
Link in: [https://github.com/shibing624/ChatPDF](https://github.com/shibing624/ChatPDF)  PS: 2核CPU 16G内存机器，约2min一条😭
"""

init_message = """欢迎使用 ChatPDF Web UI，可以直接提问或上传文件后提问 """

with gr.Blocks(css=block_css) as demo:
    index_path, file_status, model_status = gr.State(""), gr.State(""), gr.State("")
    gr.Markdown(webui_title)
    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot([[None, init_message], [None, None]],
                                 elem_id="chat-box",
                                 show_label=False).style(height=700)
            query = gr.Textbox(show_label=False,
                               placeholder="请输入提问内容，按回车进行提交",
                               ).style(container=False)
            clear_btn = gr.Button('🔄Clear!', elem_id='clear').style(full_width=True)
        with gr.Column(scale=1):
            llm_model = gr.Radio(llm_model_dict_list,
                                 label="LLM 模型",
                                 value=list(llm_model_dict.keys())[0],
                                 interactive=True)
            embedding_model = gr.Radio(embedding_model_dict_list,
                                       label="Embedding 模型",
                                       value=embedding_model_dict_list[0],
                                       interactive=True)
            load_model_button = gr.Button("重新加载模型")

            with gr.Tab("select"):
                selectFile = gr.Dropdown(
                    file_list,
                    label="content file",
                    interactive=True,
                    value=file_list[0] if len(file_list) > 0 else None
                )
            with gr.Tab("upload"):
                file = gr.File(
                    label="content file",
                    file_types=['.txt', '.md', '.docx', '.pdf']
                )
            load_file_button = gr.Button("加载文件")
    load_model_button.click(
        reinit_model,
        show_progress=True,
        inputs=[llm_model, embedding_model, chatbot],
        outputs=chatbot
    )
    # 将上传的文件保存到content文件夹下,并更新下拉框
    file.upload(upload_file, inputs=file, outputs=selectFile)
    load_file_button.click(
        get_vector_store,
        show_progress=True,
        inputs=[selectFile, chatbot],
        outputs=[index_path, chatbot],
    )
    query.submit(
        get_answer,
        [query, index_path, chatbot],
        [chatbot, query],
    )
    clear_btn.click(reset_chat, [chatbot, query], [chatbot, query])

demo.queue(concurrency_count=3).launch(
    server_name='0.0.0.0', share=False, inbrowser=False
)