163 lines
5.9 KiB
Python
163 lines
5.9 KiB
Python
import gradio as gr
|
||
import os
|
||
import shutil
|
||
from chains.local_doc_qa import LocalDocQA
|
||
from configs.model_config import *
|
||
|
||
|
||
def get_file_list():
|
||
if not os.path.exists("content"):
|
||
return []
|
||
return [f for f in os.listdir("content")]
|
||
|
||
|
||
file_list = get_file_list()
|
||
|
||
embedding_model_dict_list = list(embedding_model_dict.keys())
|
||
|
||
llm_model_dict_list = list(llm_model_dict.keys())
|
||
|
||
local_doc_qa = LocalDocQA()
|
||
|
||
|
||
def upload_file(file):
|
||
if not os.path.exists("content"):
|
||
os.mkdir("content")
|
||
filename = os.path.basename(file.name)
|
||
shutil.move(file.name, "content/" + filename)
|
||
# file_list首位插入新上传的文件
|
||
file_list.insert(0, filename)
|
||
return gr.Dropdown.update(choices=file_list, value=filename)
|
||
|
||
|
||
def get_answer(query, vs_path, history):
|
||
resp, history = local_doc_qa.get_knowledge_based_answer(
|
||
query=query, vs_path=vs_path, chat_history=history)
|
||
return history, history
|
||
|
||
|
||
def get_model_status(history):
|
||
return history + [[None, "模型已完成加载，请选择要加载的文档"]]
|
||
|
||
|
||
def get_file_status(history):
|
||
return history + [[None, "文档已完成加载，请开始提问"]]
|
||
|
||
|
||
def init_model():
|
||
try:
|
||
local_doc_qa.init_cfg()
|
||
return """模型已成功加载，请选择文件后点击"加载文件"按钮"""
|
||
except:
|
||
return """模型未成功加载，请重新选择后点击"加载模型"按钮"""
|
||
|
||
|
||
def reinit_model(llm_model, embedding_model, llm_history_len, top_k):
|
||
local_doc_qa.init_cfg(llm_model=llm_model,
|
||
embedding_model=embedding_model,
|
||
llm_history_len=llm_history_len,
|
||
top_k=top_k),
|
||
|
||
|
||
model_status = gr.State()
|
||
history = gr.State([])
|
||
vs_path = gr.State()
|
||
model_status = init_model()
|
||
with gr.Blocks(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;
|
||
}
|
||
|
||
""") as demo:
|
||
gr.Markdown(
|
||
f"""
|
||
# 🎉langchain-ChatGLM WebUI🎉
|
||
|
||
👍 [https://github.com/imClumsyPanda/langchain-ChatGLM](https://github.com/imClumsyPanda/langchain-ChatGLM)
|
||
|
||
""")
|
||
with gr.Row():
|
||
with gr.Column(scale=2):
|
||
chatbot = gr.Chatbot([[None, """欢迎使用 langchain-ChatGLM Web UI，开始提问前，请依次如下 3 个步骤：
|
||
1. 选择语言模型、Embedding 模型及相关参数后点击"重新加载模型"，并等待加载完成提示
|
||
2. 上传或选择已有文件作为本地知识文档输入后点击"重新加载文档"，并等待加载完成提示
|
||
3. 输入要提交的问题后，点击回车提交 """], [None, str(model_status)]],
|
||
elem_id="chat-box",
|
||
show_label=False).style(height=600)
|
||
query = gr.Textbox(show_label=False,
|
||
placeholder="请提问",
|
||
lines=1,
|
||
value="用200字总结一下"
|
||
).style(container=False)
|
||
|
||
with gr.Column(scale=1):
|
||
llm_model = gr.Radio(llm_model_dict_list,
|
||
label="LLM 模型",
|
||
value="chatglm-6b",
|
||
interactive=True)
|
||
llm_history_len = gr.Slider(0,
|
||
10,
|
||
value=3,
|
||
step=1,
|
||
label="LLM history len",
|
||
interactive=True)
|
||
embedding_model = gr.Radio(embedding_model_dict_list,
|
||
label="Embedding 模型",
|
||
value="text2vec",
|
||
interactive=True)
|
||
top_k = gr.Slider(1,
|
||
20,
|
||
value=6,
|
||
step=1,
|
||
label="向量匹配 top k",
|
||
interactive=True)
|
||
load_model_button = gr.Button("重新加载模型")
|
||
|
||
# with gr.Column():
|
||
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']
|
||
) # .style(height=100)
|
||
load_button = gr.Button("重新加载文件")
|
||
load_model_button.click(reinit_model,
|
||
show_progress=True,
|
||
api_name="init_cfg",
|
||
inputs=[llm_model, embedding_model, llm_history_len, top_k]
|
||
).then(
|
||
get_model_status, chatbot, chatbot
|
||
)
|
||
# 将上传的文件保存到content文件夹下,并更新下拉框
|
||
file.upload(upload_file,
|
||
inputs=file,
|
||
outputs=selectFile)
|
||
# load_button.click(local_doc_qa.init_knowledge_vector_store,
|
||
# show_progress=True,
|
||
# api_name="init_knowledge_vector_store",
|
||
# inputs=selectFile,
|
||
# outputs=vs_path
|
||
# ).then(
|
||
# get_file_status,
|
||
# chatbot,
|
||
# chatbot,
|
||
# show_progress=True,
|
||
# )
|
||
# query.submit(get_answer,
|
||
# [query, vs_path, chatbot],
|
||
# [chatbot, history],
|
||
# api_name="get_knowledge_based_answer"
|
||
# )
|
||
|
||
demo.queue(concurrency_count=3).launch(
|
||
server_name='0.0.0.0', share=False, inbrowser=False)
|