From a10fbed2488504512e9fc8d82955d9f0df698d9c Mon Sep 17 00:00:00 2001
From: yanqiangmiffy <1185918903@qq.com>
Date: Tue, 18 Apr 2023 23:45:16 +0800
Subject: [PATCH] =?UTF-8?q?feature@=E6=B7=BB=E5=8A=A0=E7=9F=A5=E8=AF=86?=
=?UTF-8?q?=E5=BA=93=E9=80=89=E6=8B=A9=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
---
README.md | 10 +++++++++-
clc/langchain_application.py | 13 +++++++++++--
clc/source_service.py | 8 ++++++--
create_knowledge.py | 34 +++++++++++++++++++++++-----------
main.py | 34 ++++++++++++++++++++++++++++++----
tests/test_langchain.py | 1 -
tests/test_vector_store.py | 11 +++++++++++
7 files changed, 90 insertions(+), 21 deletions(-)
create mode 100644 tests/test_vector_store.py
diff --git a/README.md b/README.md
index 85439a3..0b2929c 100644
--- a/README.md
+++ b/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
diff --git a/clc/langchain_application.py b/clc/langchain_application.py
index 4ca5483..ca4ab3f 100644
--- a/clc/langchain_application.py
+++ b/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}")
diff --git a/clc/source_service.py b/clc/source_service.py
index bc83764..c30cb40 100644
--- a/clc/source_service.py
+++ b/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__':
diff --git a/create_knowledge.py b/create_knowledge.py
index 1ae89fa..248b828 100644
--- a/create_knowledge.py
+++ b/create_knowledge.py
@@ -7,24 +7,36 @@
@time: 2023/04/18
@contact: yanqiangmiffy@gamil.com
@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 = []
+
+# 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 = []
-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))
-
+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')
diff --git a/main.py b/main.py
index f74c406..0bd68a2 100644
--- a/main.py
+++ b/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("""提醒:
+ [Chinese-LangChain](https://github.com/yanqiangmiffy/Chinese-LangChain)
有任何使用问题[Github Issue区](https://github.com/yanqiangmiffy/Chinese-LangChain)进行反馈.
""")
demo.queue(concurrency_count=2).launch(
diff --git a/tests/test_langchain.py b/tests/test_langchain.py
index fcf30d4..57fca86 100644
--- a/tests/test_langchain.py
+++ b/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'})]
diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py
new file mode 100644
index 0000000..146bc7b
--- /dev/null
+++ b/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)