webui.py 6.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184
  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_file_list():
  13. if not os.path.exists("content"):
  14. return []
  15. return [f for f in os.listdir("content")]
  16. file_list = get_file_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 upload_file(file):
  21. if not os.path.exists("content"):
  22. os.mkdir("content")
  23. filename = os.path.basename(file.name)
  24. shutil.move(file.name, "content/" + filename)
  25. # file_list首位插入新上传的文件
  26. file_list.insert(0, filename)
  27. return gr.Dropdown.update(choices=file_list, value=filename)
  28. def get_answer(query, vs_path, history):
  29. if vs_path:
  30. resp, history = local_doc_qa.get_knowledge_based_answer(
  31. query=query, vs_path=vs_path, chat_history=history)
  32. else:
  33. history = history + [[None, "请先加载文件后,再进行提问。"]]
  34. return history, ""
  35. def update_status(history, status):
  36. history = history + [[None, status]]
  37. print(status)
  38. return history
  39. def init_model():
  40. try:
  41. local_doc_qa.init_cfg()
  42. return """模型已成功加载,请选择文件后点击"加载文件"按钮"""
  43. except Exception as e:
  44. print(e)
  45. return """模型未成功加载,请重新选择后点击"加载模型"按钮"""
  46. def reinit_model(llm_model, embedding_model, llm_history_len, top_k, history):
  47. try:
  48. local_doc_qa.init_cfg(llm_model=llm_model,
  49. embedding_model=embedding_model,
  50. llm_history_len=llm_history_len,
  51. top_k=top_k)
  52. model_status = """模型已成功重新加载,请选择文件后点击"加载文件"按钮"""
  53. except Exception as e:
  54. print(e)
  55. model_status = """模型未成功重新加载,请重新选择后点击"加载模型"按钮"""
  56. return history + [[None, model_status]]
  57. def get_vector_store(filepath, history):
  58. if local_doc_qa.llm and local_doc_qa.embeddings:
  59. vs_path = local_doc_qa.init_knowledge_vector_store(["content/" + filepath])
  60. if vs_path:
  61. file_status = "文件已成功加载,请开始提问"
  62. else:
  63. file_status = "文件未成功加载,请重新上传文件"
  64. else:
  65. file_status = "模型未完成加载,请先在加载模型后再导入文件"
  66. vs_path = None
  67. return vs_path, history + [[None, file_status]]
  68. block_css = """.importantButton {
  69. background: linear-gradient(45deg, #7e0570,#5d1c99, #6e00ff) !important;
  70. border: none !important;
  71. }
  72. .importantButton:hover {
  73. background: linear-gradient(45deg, #ff00e0,#8500ff, #6e00ff) !important;
  74. border: none !important;
  75. }"""
  76. webui_title = """
  77. # 🎉langchain-ChatGLM WebUI🎉
  78. 👍 [https://github.com/imClumsyPanda/langchain-ChatGLM](https://github.com/imClumsyPanda/langchain-ChatGLM)
  79. """
  80. init_message = """欢迎使用 langchain-ChatGLM Web UI,开始提问前,请依次如下 3 个步骤:
  81. 1. 选择语言模型、Embedding 模型及相关参数后点击"重新加载模型",并等待加载完成提示
  82. 2. 上传或选择已有文件作为本地知识文档输入后点击"重新加载文档",并等待加载完成提示
  83. 3. 输入要提交的问题后,点击回车提交 """
  84. model_status = init_model()
  85. with gr.Blocks(css=block_css) as demo:
  86. vs_path, file_status, model_status = gr.State(""), gr.State(""), gr.State(model_status)
  87. gr.Markdown(webui_title)
  88. with gr.Row():
  89. with gr.Column(scale=2):
  90. chatbot = gr.Chatbot([[None, init_message], [None, model_status.value]],
  91. elem_id="chat-box",
  92. show_label=False).style(height=750)
  93. query = gr.Textbox(show_label=False,
  94. placeholder="请输入提问内容,按回车进行提交",
  95. ).style(container=False)
  96. with gr.Column(scale=1):
  97. llm_model = gr.Radio(llm_model_dict_list,
  98. label="LLM 模型",
  99. value=LLM_MODEL,
  100. interactive=True)
  101. llm_history_len = gr.Slider(0,
  102. 10,
  103. value=LLM_HISTORY_LEN,
  104. step=1,
  105. label="LLM history len",
  106. interactive=True)
  107. embedding_model = gr.Radio(embedding_model_dict_list,
  108. label="Embedding 模型",
  109. value=EMBEDDING_MODEL,
  110. interactive=True)
  111. top_k = gr.Slider(1,
  112. 20,
  113. value=VECTOR_SEARCH_TOP_K,
  114. step=1,
  115. label="向量匹配 top k",
  116. interactive=True)
  117. load_model_button = gr.Button("重新加载模型")
  118. # with gr.Column():
  119. with gr.Tab("select"):
  120. selectFile = gr.Dropdown(file_list,
  121. label="content file",
  122. interactive=True,
  123. value=file_list[0] if len(file_list) > 0 else None)
  124. with gr.Tab("upload"):
  125. file = gr.File(label="content file",
  126. file_types=['.txt', '.md', '.docx', '.pdf']
  127. ) # .style(height=100)
  128. load_file_button = gr.Button("加载文件")
  129. load_model_button.click(reinit_model,
  130. show_progress=True,
  131. inputs=[llm_model, embedding_model, llm_history_len, top_k, chatbot],
  132. outputs=chatbot
  133. )
  134. # 将上传的文件保存到content文件夹下,并更新下拉框
  135. file.upload(upload_file,
  136. inputs=file,
  137. outputs=selectFile)
  138. load_file_button.click(get_vector_store,
  139. show_progress=True,
  140. inputs=[selectFile, chatbot],
  141. outputs=[vs_path, chatbot],
  142. )
  143. query.submit(get_answer,
  144. [query, vs_path, chatbot],
  145. [chatbot, query],
  146. )
  147. demo.queue(concurrency_count=3).launch(
  148. server_name='0.0.0.0', share=False, inbrowser=False)