| # retrievers.py | |
| import pickle | |
| from langchain.vectorstores import FAISS | |
| from langchain.retrievers import EnsembleRetriever | |
| from kiwipiepy import Kiwi | |
| from config import FAISS_DB_INDEX, BM25_INDEX | |
| def load_retrievers(embeddings): | |
| faiss_db = FAISS.load_local( | |
| FAISS_DB_INDEX, embeddings, allow_dangerous_deserialization=True | |
| ) | |
| faiss_retriever = faiss_db.as_retriever(search_type="mmr", search_kwargs={"k": 10}) | |
| with open(BM25_INDEX, "rb") as f: | |
| bm25_retriever = pickle.load(f) | |
| return EnsembleRetriever( | |
| retrievers=[bm25_retriever, faiss_retriever], | |
| weights=[0.7, 0.3], | |
| search_type="mmr", | |
| ) | |