fix bugs in MyFAISS.delete_doc

This commit is contained in:
imClumsyPanda 2023-06-16 23:51:20 +08:00
parent 1eba834985
commit 2282320d41
1 changed files with 6 additions and 0 deletions

View File

@ -6,6 +6,7 @@ from langchain.docstore.base import Docstore
from langchain.docstore.document import Document from langchain.docstore.document import Document
import numpy as np import numpy as np
import copy import copy
import os
class MyFAISS(FAISS, VectorStore): class MyFAISS(FAISS, VectorStore):
@ -113,8 +114,10 @@ class MyFAISS(FAISS, VectorStore):
try: try:
if isinstance(source, str): if isinstance(source, str):
ids = [k for k, v in self.docstore._dict.items() if v.metadata["source"] == source] ids = [k for k, v in self.docstore._dict.items() if v.metadata["source"] == source]
vs_path = os.path.join(os.path.split(os.path.split(source)[0])[0], "vector_store")
else: else:
ids = [k for k, v in self.docstore._dict.items() if v.metadata["source"] in source] ids = [k for k, v in self.docstore._dict.items() if v.metadata["source"] in source]
vs_path = os.path.join(os.path.split(os.path.split(source[0])[0])[0], "vector_store")
if len(ids) == 0: if len(ids) == 0:
return f"docs delete fail" return f"docs delete fail"
else: else:
@ -122,6 +125,9 @@ class MyFAISS(FAISS, VectorStore):
index = list(self.index_to_docstore_id.keys())[list(self.index_to_docstore_id.values()).index(id)] index = list(self.index_to_docstore_id.keys())[list(self.index_to_docstore_id.values()).index(id)]
self.index_to_docstore_id.pop(index) self.index_to_docstore_id.pop(index)
self.docstore._dict.pop(id) self.docstore._dict.pop(id)
# TODO: 从 self.index 中删除对应id
# self.index.reset()
self.save_local(vs_path)
return f"docs delete success" return f"docs delete success"
except Exception as e: except Exception as e:
print(e) print(e)