# -*- 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 import hashlib 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 = 2048 embedding_model_dict = { "text2vec-large": "GanymedeNil/text2vec-large-chinese", "text2vec-base": "shibing624/text2vec-base-chinese", "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", } # 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 = None 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 parse_text(text): """copy from https://github.com/GaiZhenbiao/ChuanhuChatGPT/""" lines = text.split("\n") lines = [line for line in lines if line != ""] count = 0 for i, line in enumerate(lines): if "```" in line: count += 1 items = line.split('`') if count % 2 == 1: lines[i] = f'
'
            else:
                lines[i] = f'
' else: if i > 0: if count % 2 == 1: line = line.replace("`", "\`") line = line.replace("<", "<") line = line.replace(">", ">") line = line.replace(" ", " ") line = line.replace("*", "*") line = line.replace("_", "_") line = line.replace("-", "-") line = line.replace(".", ".") line = line.replace("!", "!") line = line.replace("(", "(") line = line.replace(")", ")") line = line.replace("$", "$") lines[i] = "
" + line text = "".join(lines) return text def get_answer(query, index_path, history, topn=VECTOR_SEARCH_TOP_K, max_input_size=1024, only_chat=False): if model is None: return [None, "模型还未加载"], query if index_path and not only_chat: if not model.sim_model.corpus_embeddings: model.load_index(index_path) response, empty_history, reference_results = model.query(query=query, topn=topn, max_input_size=max_input_size) logger.debug(f"query: {query}, response with content: {response}") for i in range(len(reference_results)): r = reference_results[i] response += f"\n{r.strip()}" response = parse_text(response) history = history + [[query, response]] else: # 未加载文件，仅返回生成模型结果 response, empty_history = model.gen_model.chat(query) response = parse_text(response) 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 if model is not None: 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, ) model_status = """模型已成功重新加载，请选择文件后点击"加载文件"按钮""" except Exception as e: model = None logger.error(e) model_status = """模型未成功重新加载，请重新选择后点击"加载模型"按钮""" return history + [[None, model_status]] def get_file_hash(fpath): return hashlib.md5(open(fpath, 'rb').read()).hexdigest() def get_vector_store(filepath, history, embedding_model): logger.info(filepath, history) index_path = None file_status = '' if model is not None: local_file_path = os.path.join(CONTENT_DIR, filepath) local_file_hash = get_file_hash(local_file_path) index_file_name = f"{filepath}.{embedding_model}.{local_file_hash}.index.json" local_index_path = os.path.join(CONTENT_DIR, index_file_name) if os.path.exists(local_index_path): model.load_index(local_index_path) index_path = local_index_path file_status = "文件已成功加载，请开始提问" elif 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 def change_max_input_size(input_size): if model is not None: model.max_input_size = input_size return 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.Row(): only_chat = gr.Checkbox(False, label="不加载文件(纯聊天)") with gr.Row(): topn = gr.Slider(1, 100, 20, step=1, label="最大搜索数量") max_input_size = gr.Slider(512, 4096, MAX_INPUT_LEN, step=10, label="摘要最大长度") 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("加载文件") max_input_size.change( change_max_input_size, inputs=max_input_size ) 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, embedding_model], outputs=[index_path, chatbot], ) query.submit( get_answer, [query, index_path, chatbot, topn, max_input_size, only_chat], [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 )