148 lines
4.6 KiB
Python
148 lines
4.6 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding:utf-8 _*-
|
|
"""
|
|
@author:quincy qiang
|
|
@license: Apache Licence
|
|
@file: main.py
|
|
@time: 2023/04/17
|
|
@contact: yanqiangmiffy@gamil.com
|
|
@software: PyCharm
|
|
@description: coding..
|
|
"""
|
|
|
|
import os
|
|
import shutil
|
|
|
|
import gradio as gr
|
|
|
|
from clc.langchain_application import LangChainApplication
|
|
|
|
|
|
# 修改成自己的配置!!!
|
|
class LangChainCFG:
|
|
llm_model_name = '../../pretrained_models/chatglm-6b' # 本地模型文件 or huggingface远程仓库
|
|
embedding_model_name = '../../pretrained_models/text2vec-large-chinese' # 检索模型文件 or huggingface远程仓库
|
|
vector_store_path = './cache'
|
|
docs_path = './docs'
|
|
|
|
|
|
config = LangChainCFG()
|
|
application = LangChainApplication(config)
|
|
|
|
|
|
def get_file_list():
|
|
if not os.path.exists("docs"):
|
|
return []
|
|
return [f for f in os.listdir("docs")]
|
|
|
|
|
|
file_list = get_file_list()
|
|
|
|
|
|
def upload_file(file):
|
|
if not os.path.exists("docs"):
|
|
os.mkdir("docs")
|
|
filename = os.path.basename(file.name)
|
|
shutil.move(file.name, "docs/" + filename)
|
|
# file_list首位插入新上传的文件
|
|
file_list.insert(0, filename)
|
|
application.source_service.add_document("docs/" + filename)
|
|
return gr.Dropdown.update(choices=file_list, value=filename)
|
|
|
|
|
|
def clear_session():
|
|
return '', None
|
|
|
|
|
|
def predict(input,
|
|
large_language_model,
|
|
embedding_model,
|
|
history=None):
|
|
print(large_language_model, embedding_model)
|
|
if history == None:
|
|
history = []
|
|
resp = application.get_knowledge_based_answer(
|
|
query=input,
|
|
history_len=5,
|
|
temperature=0.1,
|
|
top_p=0.9,
|
|
chat_history=history
|
|
)
|
|
print(resp)
|
|
history.append((input, resp['result']))
|
|
|
|
search_text = ''
|
|
for idx, source in enumerate(resp['source_documents'][:2]):
|
|
search_text += f'【搜索结果{idx}:】{source.page_content}\n\n'
|
|
return '', history, history, search_text
|
|
|
|
|
|
block = gr.Blocks()
|
|
with block as demo:
|
|
gr.Markdown("""<h1><center>Chinese-LangChain</center></h1>
|
|
<center><font size=3>
|
|
</center></font>
|
|
""")
|
|
with gr.Row():
|
|
with gr.Column(scale=1):
|
|
embedding_model = gr.Dropdown([
|
|
"text2vec-base"
|
|
],
|
|
label="Embedding model",
|
|
value="text2vec-base")
|
|
|
|
large_language_model = gr.Dropdown(
|
|
[
|
|
"ChatGLM-6B-int4",
|
|
],
|
|
label="large language model",
|
|
value="ChatGLM-6B-int4")
|
|
|
|
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="请上传知识库文件",
|
|
file_types=['.txt', '.md', '.docx', '.pdf']
|
|
)
|
|
|
|
file.upload(upload_file,
|
|
inputs=file,
|
|
outputs=selectFile)
|
|
with gr.Column(scale=4):
|
|
state = gr.State()
|
|
with gr.Row():
|
|
with gr.Column(scale=4):
|
|
chatbot = gr.Chatbot(label='Chinese-LangChain').style(height=400)
|
|
message = gr.Textbox(label='请输入问题')
|
|
with gr.Row():
|
|
clear_history = gr.Button("🧹 清除历史对话")
|
|
send = gr.Button("🚀 发送")
|
|
with gr.Column(scale=2):
|
|
search = gr.Textbox(label='搜索结果')
|
|
# 发送按钮 提交
|
|
send.click(predict,
|
|
inputs=[
|
|
message, large_language_model,
|
|
embedding_model, state
|
|
],
|
|
outputs=[message, chatbot, state, search])
|
|
|
|
# 清空历史对话按钮 提交
|
|
clear_history.click(fn=clear_session,
|
|
inputs=[],
|
|
outputs=[chatbot, state],
|
|
queue=False)
|
|
|
|
# 输入框 回车
|
|
message.submit(predict,
|
|
inputs=[
|
|
message, large_language_model,
|
|
embedding_model, state
|
|
],
|
|
outputs=[message, chatbot, state, search])
|
|
|
|
demo.queue().launch(server_name='0.0.0.0', server_port=8008, share=False)
|