Allow loading of pickle

open-cogsci · Apr 18, 2024 · 4a2776f · 4a2776f
1 parent aa92457
commit 4a2776f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/heymans/documentation.py b/heymans/documentation.py
@@ -104,7 +104,8 @@ def __init__(self, heymans):
         self._embeddings_model = OpenAIEmbeddings(
             openai_api_key=config.openai_api_key)
         logger.info('reading FAISS documentation cache')
-        self._db = FAISS.load_local(Path('.db.cache'), self._embeddings_model)
+        self._db = FAISS.load_local(Path('.db.cache'), self._embeddings_model,
+                                    allow_dangerous_deserialization=True)
         self._retriever = self._db.as_retriever(
             search_kwargs={'k': config.search_docs_per_query,
                            'metric': config.search_metric})

diff --git a/heymans/library.py b/heymans/library.py
@@ -15,7 +15,8 @@ def load_library(force_reindex=False, cache_folder=config.db_cache,
     embeddings_model = OpenAIEmbeddings(openai_api_key=config.openai_api_key)
     if not force_reindex and db_cache.exists():
         logger.info('loading library from cache')
-        db = FAISS.load_local(db_cache, embeddings_model)
+        db = FAISS.load_local(db_cache, embeddings_model,
+                              allow_dangerous_deserialization=True)
     else:
         from langchain_community.document_loaders import TextLoader, \
             PyPDFLoader, JSONLoader