webui.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278
  1. import gradio as gr
  2. import os
  3. import shutil
  4. from chains.local_doc_qa import LocalDocQA
  5. from configs.model_config import *
  6. import nltk
  7. nltk.data.path = [os.path.join(os.path.dirname(__file__), "nltk_data")] + nltk.data.path
  8. # return top-k text chunk from vector store
  9. VECTOR_SEARCH_TOP_K = 6
  10. # LLM input history length
  11. LLM_HISTORY_LEN = 3
  12. def get_vs_list():
  13. if not os.path.exists(VS_ROOT_PATH):
  14. return []
  15. return os.listdir(VS_ROOT_PATH)
  16. vs_list = ["新建知识库"] + get_vs_list()
  17. embedding_model_dict_list = list(embedding_model_dict.keys())
  18. llm_model_dict_list = list(llm_model_dict.keys())
  19. local_doc_qa = LocalDocQA()
  20. def get_answer(query, vs_path, history, mode,
  21. streaming: bool = STREAMING):
  22. if mode == "知识库问答" and vs_path:
  23. for resp, history in local_doc_qa.get_knowledge_based_answer(
  24. query=query,
  25. vs_path=vs_path,
  26. chat_history=history,
  27. streaming=streaming):
  28. source = "\n\n"
  29. source += "".join(
  30. [f"""<details> <summary>出处 [{i + 1}] {os.path.split(doc.metadata["source"])[-1]}</summary>\n"""
  31. f"""{doc.page_content}\n"""
  32. f"""</details>"""
  33. for i, doc in
  34. enumerate(resp["source_documents"])])
  35. history[-1][-1] += source
  36. yield history, ""
  37. else:
  38. for resp, history in local_doc_qa.llm._call(query, history,
  39. streaming=streaming):
  40. history[-1][-1] = resp + (
  41. "\n\n当前知识库为空,如需基于知识库进行问答,请先加载知识库后,再进行提问。" if mode == "知识库问答" else "")
  42. yield history, ""
  43. def update_status(history, status):
  44. history = history + [[None, status]]
  45. print(status)
  46. return history
  47. def init_model():
  48. try:
  49. local_doc_qa.init_cfg()
  50. local_doc_qa.llm._call("你好")
  51. reply = """模型已成功加载,可以开始对话,或从右侧选择模式后开始对话"""
  52. print(reply)
  53. return reply
  54. except Exception as e:
  55. print(e)
  56. reply = """模型未成功加载,请到页面左上角"模型配置"选项卡中重新选择后点击"加载模型"按钮"""
  57. if str(e) == "Unknown platform: darwin":
  58. print("该报错可能因为您使用的是 macOS 操作系统,需先下载模型至本地后执行 Web UI,具体方法请参考项目 README 中本地部署方法及常见问题:"
  59. " https://github.com/imClumsyPanda/langchain-ChatGLM")
  60. else:
  61. print(reply)
  62. return reply
  63. def reinit_model(llm_model, embedding_model, llm_history_len, use_ptuning_v2, top_k, history):
  64. try:
  65. local_doc_qa.init_cfg(llm_model=llm_model,
  66. embedding_model=embedding_model,
  67. llm_history_len=llm_history_len,
  68. use_ptuning_v2=use_ptuning_v2,
  69. top_k=top_k,)
  70. model_status = """模型已成功重新加载,可以开始对话,或从右侧选择模式后开始对话"""
  71. print(model_status)
  72. except Exception as e:
  73. print(e)
  74. model_status = """模型未成功重新加载,请到页面左上角"模型配置"选项卡中重新选择后点击"加载模型"按钮"""
  75. print(model_status)
  76. return history + [[None, model_status]]
  77. def get_vector_store(vs_id, files, history):
  78. vs_path = VS_ROOT_PATH + vs_id
  79. filelist = []
  80. for file in files:
  81. filename = os.path.split(file.name)[-1]
  82. shutil.move(file.name, UPLOAD_ROOT_PATH + filename)
  83. filelist.append(UPLOAD_ROOT_PATH + filename)
  84. if local_doc_qa.llm and local_doc_qa.embeddings:
  85. vs_path, loaded_files = local_doc_qa.init_knowledge_vector_store(filelist, vs_path)
  86. if len(loaded_files):
  87. file_status = f"已上传 {'、'.join([os.path.split(i)[-1] for i in loaded_files])} 至知识库,并已加载知识库,请开始提问"
  88. else:
  89. file_status = "文件未成功加载,请重新上传文件"
  90. else:
  91. file_status = "模型未完成加载,请先在加载模型后再导入文件"
  92. vs_path = None
  93. print(file_status)
  94. return vs_path, None, history + [[None, file_status]]
  95. def change_vs_name_input(vs_id):
  96. if vs_id == "新建知识库":
  97. return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), None
  98. else:
  99. return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), VS_ROOT_PATH + vs_id
  100. def change_mode(mode):
  101. if mode == "知识库问答":
  102. return gr.update(visible=True)
  103. else:
  104. return gr.update(visible=False)
  105. def add_vs_name(vs_name, vs_list, chatbot):
  106. if vs_name in vs_list:
  107. vs_status = "与已有知识库名称冲突,请重新选择其他名称后提交"
  108. chatbot = chatbot + [[None, vs_status]]
  109. return gr.update(visible=True), vs_list, chatbot
  110. else:
  111. vs_status = f"""已新增知识库"{vs_name}",将在上传文件并载入成功后进行存储。请在开始对话前,先完成文件上传。 """
  112. chatbot = chatbot + [[None, vs_status]]
  113. return gr.update(visible=True, choices=vs_list + [vs_name], value=vs_name), vs_list + [vs_name], chatbot
  114. block_css = """.importantButton {
  115. background: linear-gradient(45deg, #7e0570,#5d1c99, #6e00ff) !important;
  116. border: none !important;
  117. }
  118. .importantButton:hover {
  119. background: linear-gradient(45deg, #ff00e0,#8500ff, #6e00ff) !important;
  120. border: none !important;
  121. }"""
  122. webui_title = """
  123. # 🎉langchain-ChatGLM WebUI🎉
  124. 👍 [https://github.com/imClumsyPanda/langchain-ChatGLM](https://github.com/imClumsyPanda/langchain-ChatGLM)
  125. """
  126. init_message = """欢迎使用 langchain-ChatGLM Web UI!
  127. 请在右侧切换模式,目前支持直接与 LLM 模型对话或基于本地知识库问答。
  128. 知识库问答模式中,选择知识库名称后,即可开始问答,如有需要可以在选择知识库名称后上传文件/文件夹至知识库。
  129. 知识库暂不支持文件删除,该功能将在后续版本中推出。
  130. """
  131. model_status = init_model()
  132. with gr.Blocks(css=block_css) as demo:
  133. vs_path, file_status, model_status, vs_list = gr.State(""), gr.State(""), gr.State(model_status), gr.State(vs_list)
  134. gr.Markdown(webui_title)
  135. with gr.Tab("对话"):
  136. with gr.Row():
  137. with gr.Column(scale=10):
  138. chatbot = gr.Chatbot([[None, init_message], [None, model_status.value]],
  139. elem_id="chat-box",
  140. show_label=False).style(height=750)
  141. query = gr.Textbox(show_label=False,
  142. placeholder="请输入提问内容,按回车进行提交",
  143. ).style(container=False)
  144. with gr.Column(scale=5):
  145. mode = gr.Radio(["LLM 对话", "知识库问答"],
  146. label="请选择使用模式",
  147. value="知识库问答", )
  148. vs_setting = gr.Accordion("配置知识库")
  149. mode.change(fn=change_mode,
  150. inputs=mode,
  151. outputs=vs_setting)
  152. with vs_setting:
  153. select_vs = gr.Dropdown(vs_list.value,
  154. label="请选择要加载的知识库",
  155. interactive=True,
  156. value=vs_list.value[0] if len(vs_list.value) > 0 else None
  157. )
  158. vs_name = gr.Textbox(label="请输入新建知识库名称",
  159. lines=1,
  160. interactive=True)
  161. vs_add = gr.Button(value="添加至知识库选项")
  162. vs_add.click(fn=add_vs_name,
  163. inputs=[vs_name, vs_list, chatbot],
  164. outputs=[select_vs, vs_list, chatbot])
  165. file2vs = gr.Column(visible=False)
  166. with file2vs:
  167. # load_vs = gr.Button("加载知识库")
  168. gr.Markdown("向知识库中添加文件")
  169. with gr.Tab("上传文件"):
  170. files = gr.File(label="添加文件",
  171. file_types=['.txt', '.md', '.docx', '.pdf'],
  172. file_count="multiple",
  173. show_label=False
  174. )
  175. load_file_button = gr.Button("上传文件并加载知识库")
  176. with gr.Tab("上传文件夹"):
  177. folder_files = gr.File(label="添加文件",
  178. # file_types=['.txt', '.md', '.docx', '.pdf'],
  179. file_count="directory",
  180. show_label=False
  181. )
  182. load_folder_button = gr.Button("上传文件夹并加载知识库")
  183. # load_vs.click(fn=)
  184. select_vs.change(fn=change_vs_name_input,
  185. inputs=select_vs,
  186. outputs=[vs_name, vs_add, file2vs, vs_path])
  187. # 将上传的文件保存到content文件夹下,并更新下拉框
  188. load_file_button.click(get_vector_store,
  189. show_progress=True,
  190. inputs=[select_vs, files, chatbot],
  191. outputs=[vs_path, files, chatbot],
  192. )
  193. load_folder_button.click(get_vector_store,
  194. show_progress=True,
  195. inputs=[select_vs, folder_files, chatbot],
  196. outputs=[vs_path, folder_files, chatbot],
  197. )
  198. query.submit(get_answer,
  199. [query, vs_path, chatbot, mode],
  200. [chatbot, query],
  201. )
  202. with gr.Tab("模型配置"):
  203. llm_model = gr.Radio(llm_model_dict_list,
  204. label="LLM 模型",
  205. value=LLM_MODEL,
  206. interactive=True)
  207. llm_history_len = gr.Slider(0,
  208. 10,
  209. value=LLM_HISTORY_LEN,
  210. step=1,
  211. label="LLM 对话轮数",
  212. interactive=True)
  213. use_ptuning_v2 = gr.Checkbox(USE_PTUNING_V2,
  214. label="使用p-tuning-v2微调过的模型",
  215. interactive=True)
  216. embedding_model = gr.Radio(embedding_model_dict_list,
  217. label="Embedding 模型",
  218. value=EMBEDDING_MODEL,
  219. interactive=True)
  220. top_k = gr.Slider(1,
  221. 20,
  222. value=VECTOR_SEARCH_TOP_K,
  223. step=1,
  224. label="向量匹配 top k",
  225. interactive=True)
  226. load_model_button = gr.Button("重新加载模型")
  227. load_model_button.click(reinit_model,
  228. show_progress=True,
  229. inputs=[llm_model, embedding_model, llm_history_len, use_ptuning_v2, top_k, chatbot],
  230. outputs=chatbot
  231. )
  232. (demo
  233. .queue(concurrency_count=3)
  234. .launch(server_name='0.0.0.0',
  235. server_port=7860,
  236. show_api=False,
  237. share=False,
  238. inbrowser=False))