提交 e9a5db1a 作者: imClumsyPanda

update webui.py

上级 7d4560e5
...@@ -14,19 +14,12 @@ VECTOR_SEARCH_TOP_K = 6 ...@@ -14,19 +14,12 @@ VECTOR_SEARCH_TOP_K = 6
LLM_HISTORY_LEN = 3 LLM_HISTORY_LEN = 3
def get_file_list():
if not os.path.exists("content"):
return []
return [f for f in os.listdir("content")]
def get_vs_list(): def get_vs_list():
if not os.path.exists("vector_store"): if not os.path.exists(VS_ROOT_PATH):
return [] return []
return ["新建知识库"] + os.listdir("vector_store") return ["新建知识库"] + os.listdir(VS_ROOT_PATH)
file_list = get_file_list()
vs_list = get_vs_list() vs_list = get_vs_list()
embedding_model_dict_list = list(embedding_model_dict.keys()) embedding_model_dict_list = list(embedding_model_dict.keys())
...@@ -36,19 +29,8 @@ llm_model_dict_list = list(llm_model_dict.keys()) ...@@ -36,19 +29,8 @@ llm_model_dict_list = list(llm_model_dict.keys())
local_doc_qa = LocalDocQA() local_doc_qa = LocalDocQA()
def upload_file(file, chatbot): def get_answer(query, vs_path, history, mode):
if not os.path.exists("content"): if vs_path and mode == "知识库问答":
os.mkdir("content")
filename = os.path.basename(file.name)
shutil.move(file.name, "content/" + filename)
# file_list首位插入新上传的文件
file_list.insert(0, filename)
status = "已将xx上传至xxx"
return chatbot + [None, status]
def get_answer(query, vs_path, history):
if vs_path:
resp, history = local_doc_qa.get_knowledge_based_answer( resp, history = local_doc_qa.get_knowledge_based_answer(
query=query, vs_path=vs_path, chat_history=history) query=query, vs_path=vs_path, chat_history=history)
source = "".join([f"""<details> <summary>出处 {i + 1}</summary> source = "".join([f"""<details> <summary>出处 {i + 1}</summary>
...@@ -93,24 +75,30 @@ def reinit_model(llm_model, embedding_model, llm_history_len, use_ptuning_v2, to ...@@ -93,24 +75,30 @@ def reinit_model(llm_model, embedding_model, llm_history_len, use_ptuning_v2, to
return history + [[None, model_status]] return history + [[None, model_status]]
def get_vector_store(filepath, history): def get_vector_store(vs_id, files, history):
vs_path = VS_ROOT_PATH + vs_id
filelist = []
for file in files:
filename = os.path.split(file.name)[-1]
shutil.move(file.name, UPLOAD_ROOT_PATH + filename)
filelist.append(UPLOAD_ROOT_PATH + filename)
if local_doc_qa.llm and local_doc_qa.embeddings: if local_doc_qa.llm and local_doc_qa.embeddings:
vs_path = local_doc_qa.init_knowledge_vector_store(["content/" + filepath]) vs_path, loaded_files = local_doc_qa.init_knowledge_vector_store(filelist, vs_path)
if vs_path: if len(loaded_files):
file_status = "文件已成功加载,请开始提问" file_status = f"已上传 {'、'.join([os.path.split(i)[-1] for i in loaded_files])} 至知识库,并已加载知识库,请开始提问"
else: else:
file_status = "文件未成功加载,请重新上传文件" file_status = "文件未成功加载,请重新上传文件"
else: else:
file_status = "模型未完成加载,请先在加载模型后再导入文件" file_status = "模型未完成加载,请先在加载模型后再导入文件"
vs_path = None vs_path = None
return vs_path, history + [[None, file_status]] return vs_path, None, history + [[None, file_status]]
def change_vs_name_input(vs): def change_vs_name_input(vs):
if vs == "新建知识库": if vs == "新建知识库":
return gr.update(visible=True), gr.update(visible=True) return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)
else: else:
return gr.update(visible=False), gr.update(visible=False) return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
def change_mode(mode): def change_mode(mode):
...@@ -119,15 +107,15 @@ def change_mode(mode): ...@@ -119,15 +107,15 @@ def change_mode(mode):
else: else:
return gr.update(visible=False) return gr.update(visible=False)
def add_vs_name(vs_name, vs_list, chatbot): def add_vs_name(vs_name, vs_list, chatbot):
if vs_name in vs_list: if vs_name in vs_list:
chatbot = chatbot+[None, "与已有知识库名称冲突,请重新选择其他名称后提交"] chatbot = chatbot + [[None, "与已有知识库名称冲突,请重新选择其他名称后提交"]]
return gr.update(visible=True),vs_list, chatbot return gr.update(visible=True), vs_list, chatbot
else: else:
chatbot = chatbot + [None, f"""已新增知识库"{vs_name}" """] chatbot = chatbot + [
vs_list = vs_list+[vs_name] [None, f"""已新增知识库"{vs_name}",将在上传文件并载入成功后进行存储。请在开始对话前,先完成文件上传。 """]]
return gr.update(visible=True),vs_list, chatbot return gr.update(visible=True, choices=vs_list + [vs_name], value=vs_name), vs_list + [vs_name], chatbot
block_css = """.importantButton { block_css = """.importantButton {
...@@ -155,7 +143,7 @@ init_message = """欢迎使用 langchain-ChatGLM Web UI,开始提问前,请 ...@@ -155,7 +143,7 @@ init_message = """欢迎使用 langchain-ChatGLM Web UI,开始提问前,请
model_status = init_model() model_status = init_model()
with gr.Blocks(css=block_css) as demo: with gr.Blocks(css=block_css) as demo:
vs_path, file_status, model_status = gr.State(""), gr.State(""), gr.State(model_status) vs_path, file_status, model_status, vs_list = gr.State(""), gr.State(""), gr.State(model_status), gr.State(vs_list)
gr.Markdown(webui_title) gr.Markdown(webui_title)
with gr.Tab("对话"): with gr.Tab("对话"):
with gr.Row(): with gr.Row():
...@@ -169,16 +157,17 @@ with gr.Blocks(css=block_css) as demo: ...@@ -169,16 +157,17 @@ with gr.Blocks(css=block_css) as demo:
with gr.Column(scale=5): with gr.Column(scale=5):
mode = gr.Radio(["LLM 对话", "知识库问答"], mode = gr.Radio(["LLM 对话", "知识库问答"],
label="请选择使用模式", label="请选择使用模式",
value="知识库问答",) value="知识库问答", )
vs_setting = gr.Accordion("配置知识库") vs_setting = gr.Accordion("配置知识库")
mode.change(fn=change_mode, mode.change(fn=change_mode,
inputs=mode, inputs=mode,
outputs=vs_setting) outputs=vs_setting)
with vs_setting: with vs_setting:
select_vs = gr.Dropdown(vs_list, select_vs = gr.Dropdown(vs_list.value,
label="请选择要加载的知识库", label="请选择要加载的知识库",
interactive=True, interactive=True,
value=vs_list[0] if len(vs_list) > 0 else None) value=vs_list.value[0] if len(vs_list.value) > 0 else None
)
vs_name = gr.Textbox(label="请输入新建知识库名称", vs_name = gr.Textbox(label="请输入新建知识库名称",
lines=1, lines=1,
interactive=True) interactive=True)
...@@ -186,24 +175,42 @@ with gr.Blocks(css=block_css) as demo: ...@@ -186,24 +175,42 @@ with gr.Blocks(css=block_css) as demo:
vs_add.click(fn=add_vs_name, vs_add.click(fn=add_vs_name,
inputs=[vs_name, vs_list, chatbot], inputs=[vs_name, vs_list, chatbot],
outputs=[select_vs, vs_list, chatbot]) outputs=[select_vs, vs_list, chatbot])
file2vs = gr.Box(visible=False)
with file2vs:
gr.Markdown("向知识库中添加文件")
with gr.Tab("上传文件"):
files = gr.File(label="添加文件",
file_types=['.txt', '.md', '.docx', '.pdf'],
file_count="multiple",
show_label=False
)
load_file_button = gr.Button("上传文件")
with gr.Tab("上传文件夹"):
folder_files = gr.File(label="添加文件",
# file_types=['.txt', '.md', '.docx', '.pdf'],
file_count="directory",
show_label=False
)
load_folder_button = gr.Button("上传文件夹")
select_vs.change(fn=change_vs_name_input, select_vs.change(fn=change_vs_name_input,
inputs=select_vs, inputs=select_vs,
outputs=[vs_name, vs_add]) outputs=[vs_name, vs_add, file2vs])
gr.Markdown("向知识库中添加文件") # 将上传的文件保存到content文件夹下,并更新下拉框
with gr.Tab("上传文件"): load_file_button.click(get_vector_store,
files = gr.File(label="添加文件", show_progress=True,
file_types=['.txt', '.md', '.docx', '.pdf'], inputs=[select_vs, files, chatbot],
file_count="multiple", outputs=[vs_path, files, chatbot],
show_label=False )
) load_folder_button.click(get_vector_store,
load_file_button = gr.Button("上传文件") show_progress=True,
with gr.Tab("上传文件夹"): inputs=[select_vs, folder_files, chatbot],
folder_files = gr.File(label="添加文件", outputs=[vs_path, folder_files, chatbot],
file_types=['.txt', '.md', '.docx', '.pdf'], )
file_count="directory", query.submit(get_answer,
show_label=False [query, vs_path, chatbot, mode],
) [chatbot, query],
load_folder_button = gr.Button("上传文件夹") )
with gr.Tab("模型配置"): with gr.Tab("模型配置"):
llm_model = gr.Radio(llm_model_dict_list, llm_model = gr.Radio(llm_model_dict_list,
label="LLM 模型", label="LLM 模型",
...@@ -213,7 +220,7 @@ with gr.Blocks(css=block_css) as demo: ...@@ -213,7 +220,7 @@ with gr.Blocks(css=block_css) as demo:
10, 10,
value=LLM_HISTORY_LEN, value=LLM_HISTORY_LEN,
step=1, step=1,
label="LLM history len", label="LLM 对话轮数",
interactive=True) interactive=True)
use_ptuning_v2 = gr.Checkbox(USE_PTUNING_V2, use_ptuning_v2 = gr.Checkbox(USE_PTUNING_V2,
label="使用p-tuning-v2微调过的模型", label="使用p-tuning-v2微调过的模型",
...@@ -234,19 +241,6 @@ with gr.Blocks(css=block_css) as demo: ...@@ -234,19 +241,6 @@ with gr.Blocks(css=block_css) as demo:
inputs=[llm_model, embedding_model, llm_history_len, use_ptuning_v2, top_k, chatbot], inputs=[llm_model, embedding_model, llm_history_len, use_ptuning_v2, top_k, chatbot],
outputs=chatbot outputs=chatbot
) )
# 将上传的文件保存到content文件夹下,并更新下拉框
files.upload(upload_file,
inputs=[files, chatbot],
outputs=chatbot)
load_file_button.click(get_vector_store,
show_progress=True,
inputs=[select_vs, chatbot],
outputs=[vs_path, chatbot],
)
query.submit(get_answer,
[query, vs_path, chatbot],
[chatbot, query],
)
demo.queue(concurrency_count=3 demo.queue(concurrency_count=3
).launch(server_name='0.0.0.0', ).launch(server_name='0.0.0.0',
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论