feature@添加知识库选择功能

README.md

@@ -7,13 +7,20 @@
 ![](https://github.com/yanqiangmiffy/Chinese-LangChain/blob/master/images/web_demo.png)
 
 ## 🚀 特性
+
 - 🚀 2023/04/18 webui增加知识库选择功能
 - 🚀 2023/04/18 修复推理预测超时5s报错问题
-- 🎉 2023/04/17  支持多种文档上传与内容解析：pdf、docx，ppt等
+- 🎉 2023/04/17 支持多种文档上传与内容解析：pdf、docx，ppt等
 - 🎉 2023/04/17 支持知识增量更新
 
 [//]: # (- 支持检索结果与LLM生成结果对比)
 
+## 🧰 知识库
+
+| 知识库数据  |FAISS向量|
+|--------|----|
+|💹 [大规模金融研报知识图谱](http://openkg.cn/dataset/fr2kg)|链接：https://pan.baidu.com/s/1FcIH5Fi3EfpS346DnDu51Q?pwd=ujjv 提取码：ujjv |
+
 ## 🔨 TODO
 
 * [x] 支持上下文
@@ -26,6 +33,7 @@
 * [ ] 增加非LangChain策略
 
 ## 交流
+
 欢迎多提建议、Bad cases，目前尚不完善，欢迎进群及时交流，也欢迎大家多提PR
 <img src="https://raw.githubusercontent.com/yanqiangmiffy/Chinese-LangChain/master/images/ch.jpg" width="300px">
 <img src="https://raw.githubusercontent.com/yanqiangmiffy/Chinese-LangChain/master/images/chatgroup.jpg" width="300px">

clc/langchain_application.py

@@ -12,6 +12,7 @@
 
 from langchain.chains import RetrievalQA
 from langchain.prompts.prompt import PromptTemplate
+
 from clc.gpt_service import ChatGLMService
 from clc.source_service import SourceService
 
@@ -22,7 +23,15 @@ class LangChainApplication(object):
         self.llm_service = ChatGLMService()
         self.llm_service.load_model(model_name_or_path=self.config.llm_model_name)
         self.source_service = SourceService(config)
-        self.source_service.init_source_vector()
+        if self.config.kg_vector_stores is None:
+            print("init a source vector store")
+            self.source_service.init_source_vector()
+        else:
+            print("load zh_wikipedia source vector store ")
+            try:
+                self.source_service.load_vector_store(self.config.kg_vector_stores['初始化知识库'])
+            except Exception as e:
+                self.source_service.init_source_vector()
 
     def get_knowledge_based_answer(self, query,
                                    history_len=5,
@@ -45,7 +54,7 @@ class LangChainApplication(object):
         knowledge_chain = RetrievalQA.from_llm(
             llm=self.llm_service,
             retriever=self.source_service.vector_store.as_retriever(
-                search_kwargs={"k": 2}),
+                search_kwargs={"k": 4}),
             prompt=prompt)
         knowledge_chain.combine_documents_chain.document_prompt = PromptTemplate(
             input_variables=["page_content"], template="{page_content}")

clc/source_service.py

@@ -19,6 +19,7 @@ from langchain.vectorstores import FAISS
 
 class SourceService(object):
     def __init__(self, config):
+        self.vector_store = None
         self.config = config
         self.embeddings = HuggingFaceEmbeddings(model_name=self.config.embedding_model_name)
         self.docs_path = self.config.docs_path
@@ -45,8 +46,11 @@ class SourceService(object):
         self.vector_store.add_documents(doc)
         self.vector_store.save_local(self.vector_store_path)
 
-    def load_vector_store(self):
-        self.vector_store = FAISS.load_local(self.vector_store_path, self.embeddings)
+    def load_vector_store(self, path):
+        if path is None:
+            self.vector_store = FAISS.load_local(self.vector_store_path, self.embeddings)
+        else:
+            self.vector_store = FAISS.load_local(path, self.embeddings)
         return self.vector_store
 
 # if __name__ == '__main__':

create_knowledge.py

@@ -7,24 +7,36 @@
 @time: 2023/04/18
 @contact: [email protected]
 @software: PyCharm
-@description: coding..
+@description: - emoji：https://emojixd.com/pocket/science
 """
-from langchain.docstore.document import Document
+import os
+
+from langchain.document_loaders import UnstructuredFileLoader
 from langchain.embeddings.huggingface import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 from tqdm import tqdm
-
 # 中文Wikipedia数据导入示例：
-embedding_model_name = '/home/searchgpt/pretrained_models/ernie-gram-zh'
-docs_path = '/home/searchgpt/yq/Knowledge-ChatGLM/docs'
+embedding_model_name = '/root/pretrained_models/text2vec-large-chinese'
+docs_path = '/root/GoMall/Knowledge-ChatGLM/cache/financial_research_reports'
 embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name)
 
-docs = []
+# docs = []
 
-with open('docs/zh_wikipedia/zhwiki.sim.utf8', 'r', encoding='utf-8') as f:
-    for idx, line in tqdm(enumerate(f.readlines())):
-        metadata = {"source": f'doc_id_{idx}'}
-        docs.append(Document(page_content=line.strip(), metadata=metadata))
+# with open('docs/zh_wikipedia/zhwiki.sim.utf8', 'r', encoding='utf-8') as f:
+#     for idx, line in tqdm(enumerate(f.readlines())):
+#         metadata = {"source": f'doc_id_{idx}'}
+#         docs.append(Document(page_content=line.strip(), metadata=metadata))
+#
+# vector_store = FAISS.from_documents(docs, embeddings)
+# vector_store.save_local('cache/zh_wikipedia/')
+
+docs = []
 
+for doc in tqdm(os.listdir(docs_path)):
+    if doc.endswith('.txt'):
+        # print(doc)
+        loader = UnstructuredFileLoader(f'{docs_path}/{doc}', mode="elements")
+        doc = loader.load()
+        docs.extend(doc)
 vector_store = FAISS.from_documents(docs, embeddings)
-vector_store.save_local('cache/zh_wikipedia/')
+vector_store.save_local('cache/financial_research_reports')

main.py

@@ -14,6 +14,12 @@ class LangChainCFG:
     embedding_model_name = '../../pretrained_models/text2vec-large-chinese'  # 检索模型文件 or huggingface远程仓库
     vector_store_path = './cache'
     docs_path = './docs'
+    kg_vector_stores = {
+        '中文维基百科': '/root/GoMall/Knowledge-ChatGLM/cache/zh_wikipedia',
+        '大规模金融研报知识图谱': '/root/GoMall/Knowledge-ChatGLM/cache/financial_research_reports',
+        '初始化知识库': '/root/GoMall/Knowledge-ChatGLM/cache',
+    }  # 可以替换成自己的知识库，如果没有需要设置为None
+    # kg_vector_stores=None
 
 
 config = LangChainCFG()
@@ -40,6 +46,15 @@ def upload_file(file):
     return gr.Dropdown.update(choices=file_list, value=filename)
 
 
+def set_knowledge(kg_name, history):
+    try:
+        application.source_service.load_vector_store(config.kg_vector_stores[kg_name])
+        msg_status = f'{kg_name}知识库已成功加载'
+    except Exception as e:
+        msg_status = f'{kg_name}知识库未成功加载'
+    return history + [[None, msg_status]]
+
+
 def clear_session():
     return '', None
 
@@ -61,8 +76,8 @@ def predict(input,
     )
     history.append((input, resp['result']))
     search_text = ''
-    for idx, source in enumerate(resp['source_documents'][:2]):
-        sep = f'----------【搜索结果{idx}：】---------------\n'
+    for idx, source in enumerate(resp['source_documents'][:4]):
+        sep = f'----------【搜索结果{idx+1}：】---------------\n'
         search_text += f'{sep}\n{source.page_content}\n\n'
     print(search_text)
     return '', history, history, search_text
@@ -97,10 +112,15 @@ with block as demo:
                               step=1,
                               label="向量匹配 top k",
                               interactive=True)
-            kg_name = gr.Radio(['中文维基百科', '百度百科数据', '坦克世界'],
+            kg_name = gr.Radio(['中文维基百科',
+                                '大规模金融研报知识图谱',
+                                '初始化知识库'
+                                ],
                                label="知识库",
                                value='中文维基百科',
                                interactive=True)
+            set_kg_btn = gr.Button("重新加载知识库")
+
             file = gr.File(label="将文件上传到数据库",
                            visible=True,
                            file_types=['.txt', '.md', '.docx', '.pdf']
@@ -119,7 +139,12 @@ with block as demo:
                         send = gr.Button("🚀 发送")
                 with gr.Column(scale=2):
                     search = gr.Textbox(label='搜索结果')
-
+        set_kg_btn.click(
+            set_knowledge,
+            show_progress=True,
+            inputs=[kg_name, chatbot],
+            outputs=chatbot
+        )
         # 发送按钮 提交
         send.click(predict,
                    inputs=[
@@ -142,6 +167,7 @@ with block as demo:
                        ],
                        outputs=[message, chatbot, state, search])
     gr.Markdown("""提醒：<br>
+            [Chinese-LangChain](https://github.com/yanqiangmiffy/Chinese-LangChain) <br>
             有任何使用问题[Github Issue区](https://github.com/yanqiangmiffy/Chinese-LangChain)进行反馈. <br>
             """)
 demo.queue(concurrency_count=2).launch(

tests/test_langchain.py

@@ -29,7 +29,6 @@ print(doc)
 search_result = vector_store.similarity_search_with_score(query='科比·布莱恩特', k=2)
 print(search_result)
 
-
 """
 [(Document(page_content='王治郅，1977年7月8日出生于北京，前中国篮球运动员，司职大前锋/中锋，现已退役。 [1]', metadata={'source': 'docs/王治郅.txt', 'filename': 'docs/王治郅.txt', 'category': 'Title'}), 285.40765), (Document(page_content='王治郅是中国篮球界进入NBA的第一人，被评选为中国篮坛50大杰出人物和中国申办奥运特使。他和姚明、蒙克·巴特尔一起，被称为篮球场上的“移动长城”。 [5]', metadata={'source': 'docs/王治郅.txt', 'filename': 'docs/王治郅.txt', 'category': 'NarrativeText'}), 290.19086)]
 [Document(page_content='科比·布莱恩特（Kobe Bryant，1978年8月23日—2020年1月26日），全名科比·比恩·布莱恩特·考克斯（Kobe Bean Bryant Cox），出生于美国宾夕法尼亚州费城，美国已故篮球运动员，司职得分后卫/小前锋。 [5]  [24]  [84]', metadata={'source': 'docs/added/科比.txt', 'filename': 'docs/added/科比.txt', 'category': 'NarrativeText'}), Document(page_content='1996年NBA选秀，科比于第1轮第13顺位被夏洛特黄蜂队选中并被交易至洛杉矶湖人队，整个NBA生涯都效力于洛杉矶湖人队；共获得5次NBA总冠军、1次NBA常规赛MVP、2次NBA总决赛MVP、4次NBA全明星赛MVP、2次NBA赛季得分王；共入选NBA全明星首发阵容18次、NBA最佳阵容15次（其中一阵11次、二阵2次、三阵2次）、NBA最佳防守阵容12次（其中一阵9次、二阵3次）。 [9]  [24]', metadata={'source': 'docs/added/科比.txt', 'filename': 'docs/added/科比.txt', 'category': 'Title'}), Document(page_content='2007年，科比首次入选美国国家男子篮球队，后帮助美国队夺得2007年美洲男篮锦标赛金牌、2008年北京奥运会男子篮球金牌以及2012年伦敦奥运会男子篮球金牌。 [91]', metadata={'source': 'docs/added/科比.txt', 'filename': 'docs/added/科比.txt', 'category': 'Title'}), Document(page_content='2015年11月30日，科比发文宣布将在赛季结束后退役。 [100]  2017年12月19日，湖人队为科比举行球衣退役仪式。 [22]  2020年4月5日，科比入选奈·史密斯篮球名人纪念堂。 [7]', metadata={'source': 'docs/added/科比.txt', 'filename': 'docs/added/科比.txt', 'category': 'Title'}), Document(page_content='美国时间2020年1月26日（北京时间2020年1月27日），科比因直升机事故遇难，享年41岁。 [23]', metadata={'source': 'docs/added/科比.txt', 'filename': 'docs/added/科比.txt', 'category': 'Title'})]

tests/test_vector_store.py

@@ -0,0 +1,11 @@
+from langchain.embeddings.huggingface import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+
+# 中文Wikipedia数据导入示例：
+embedding_model_name = '/root/pretrained_models/ernie-gram-zh'
+embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name)
+
+vector_store = FAISS.load_local("/root/GoMall/Knowledge-ChatGLM/cache/zh_wikipedia", embeddings)
+print(vector_store)
+res = vector_store.similarity_search_with_score('闫强')
+print(res)