提交 7d4560e5 作者: imClumsyPanda

update local_doc_qa.py

上级 25ecc7f8
...@@ -54,10 +54,11 @@ async def upload_file(UserFile: UploadFile=File(...)): ...@@ -54,10 +54,11 @@ async def upload_file(UserFile: UploadFile=File(...)):
# print(UserFile.filename) # print(UserFile.filename)
with open(filepath, 'wb') as f: with open(filepath, 'wb') as f:
f.write(content) f.write(content)
vs_path = local_doc_qa.init_knowledge_vector_store(filepath) vs_path, files = local_doc_qa.init_knowledge_vector_store(filepath)
response = { response = {
'msg': 'seccessful', 'msg': 'seccess' if len(files)>0 else 'fail',
'status': 1 'status': 1 if len(files)>0 else 0,
'loaded_files': files
} }
except Exception as err: except Exception as err:
......
...@@ -53,7 +53,9 @@ class LocalDocQA: ...@@ -53,7 +53,9 @@ class LocalDocQA:
self.top_k = top_k self.top_k = top_k
def init_knowledge_vector_store(self, def init_knowledge_vector_store(self,
filepath: str or List[str]): filepath: str or List[str],
vs_path: str or os.PathLike = None):
loaded_files = []
if isinstance(filepath, str): if isinstance(filepath, str):
if not os.path.exists(filepath): if not os.path.exists(filepath):
print("路径不存在") print("路径不存在")
...@@ -63,6 +65,7 @@ class LocalDocQA: ...@@ -63,6 +65,7 @@ class LocalDocQA:
try: try:
docs = load_file(filepath) docs = load_file(filepath)
print(f"{file} 已成功加载") print(f"{file} 已成功加载")
loaded_files.append(filepath)
except Exception as e: except Exception as e:
print(e) print(e)
print(f"{file} 未能成功加载") print(f"{file} 未能成功加载")
...@@ -74,6 +77,7 @@ class LocalDocQA: ...@@ -74,6 +77,7 @@ class LocalDocQA:
try: try:
docs += load_file(fullfilepath) docs += load_file(fullfilepath)
print(f"{file} 已成功加载") print(f"{file} 已成功加载")
loaded_files.append(fullfilepath)
except Exception as e: except Exception as e:
print(e) print(e)
print(f"{file} 未能成功加载") print(f"{file} 未能成功加载")
...@@ -83,14 +87,21 @@ class LocalDocQA: ...@@ -83,14 +87,21 @@ class LocalDocQA:
try: try:
docs += load_file(file) docs += load_file(file)
print(f"{file} 已成功加载") print(f"{file} 已成功加载")
loaded_files.append(file)
except Exception as e: except Exception as e:
print(e) print(e)
print(f"{file} 未能成功加载") print(f"{file} 未能成功加载")
if vs_path and os.path.isdir(vs_path):
vector_store = FAISS.load_local(vs_path, self.embeddings)
vector_store.add_documents(docs)
else:
if not vs_path:
vs_path = f"""{VS_ROOT_PATH}{os.path.splitext(file)[0]}_FAISS_{datetime.datetime.now().strftime("%Y%m%d_%H%M%S")}"""
vector_store = FAISS.from_documents(docs, self.embeddings) vector_store = FAISS.from_documents(docs, self.embeddings)
vs_path = f"""./vector_store/{os.path.splitext(file)[0]}_FAISS_{datetime.datetime.now().strftime("%Y%m%d_%H%M%S")}"""
vector_store.save_local(vs_path) vector_store.save_local(vs_path)
return vs_path if len(docs) > 0 else None return vs_path if len(docs) > 0 else None, loaded_files
def get_knowledge_based_answer(self, def get_knowledge_based_answer(self,
query, query,
......
...@@ -24,7 +24,7 @@ if __name__ == "__main__": ...@@ -24,7 +24,7 @@ if __name__ == "__main__":
vs_path = None vs_path = None
while not vs_path: while not vs_path:
filepath = input("Input your local knowledge file path 请输入本地知识文件路径:") filepath = input("Input your local knowledge file path 请输入本地知识文件路径:")
vs_path = local_doc_qa.init_knowledge_vector_store(filepath) vs_path, _ = local_doc_qa.init_knowledge_vector_store(filepath)
history = [] history = []
while True: while True:
query = input("Input your question 请输入问题:") query = input("Input your question 请输入问题:")
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论