Skip to content

Commit

Permalink
Fix allow_dangerous_deserialization bug (#18)
Browse files Browse the repository at this point in the history
Co-authored-by: Nekxis <nexispolska16@gmail.com>
  • Loading branch information
latekvo and Nekxis authored Apr 15, 2024
1 parent a73fd8d commit 884d7a1
Showing 1 changed file with 24 additions and 6 deletions.
30 changes: 24 additions & 6 deletions core/tools/dbops.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,32 @@


def create_db_if_not_exists(db_name: str, embeddings: Embeddings):
if not exists('store/vector/' + db_name + '.faiss'):
print("Creating new database:", db_name + '.faiss')
tmp_db = FAISS.from_texts(['You are a large language model, intended for research purposes.'], embeddings)
tmp_db.save_local(folder_path='store/vector', index_name=db_name)
if not exists("store/vector/" + db_name + ".faiss"):
print("Creating new database:", db_name + ".faiss")
tmp_db = FAISS.from_texts(
["You are a large language model, intended for research purposes."],
embeddings,
)
tmp_db.save_local(folder_path="store/vector", index_name=db_name)
else:
print("Already exists:", db_name + '.faiss')
print("Already exists:", db_name + ".faiss")


def get_db_by_name(db_name: str, embeddings: Embeddings) -> FAISS:
create_db_if_not_exists(db_name, embeddings)
return FAISS.load_local(folder_path='store/vector', embeddings=embeddings, index_name=db_name, allow_dangerous_deserialization=True)

try:
# windows
db_connection = FAISS.load_local(
folder_path="store/vector",
embeddings=embeddings,
index_name=db_name,
allow_dangerous_deserialization=True,
)
except Exception:
# linux & mac
db_connection = FAISS.load_local(
folder_path="store/vector", embeddings=embeddings, index_name=db_name
)

return db_connection

0 comments on commit 884d7a1

Please sign in to comment.