Merge pull request #6 from almutareb/hybrid_search

rag_app/hybrid_search.py

@@ -0,0 +1,63 @@
+from pathlib import Path
+from langchain_community.vectorstores import FAISS
+from dotenv import load_dotenv
+import os
+from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
+from langchain.retrievers import EnsembleRetriever
+from langchain_community.retrievers import BM25Retriever
+
+
+def get_hybrid_search_results(query:str,
+                              path_to_db:str,
+                              embedding_model:str,
+                              hf_api_key:str,
+                              num_docs:int=5) -> list:
+    """ Uses an ensemble retriever of BM25 and FAISS to return k num documents
+
+        Args:
+            query (str): The search query
+            path_to_db (str): Path to the vectorstore database
+            embedding_model (str): Embedding model used in the vector store
+            num_docs (int): Number of documents to return
+        
+        Returns
+            List of documents
+    
+    """
+    
+    embeddings = HuggingFaceInferenceAPIEmbeddings(api_key=hf_api_key,
+                                                   model_name=embedding_model)
+    # Load the vectorstore database
+    db = FAISS.load_local(folder_path=path_to_db,
+                          embeddings=embeddings,
+                          allow_dangerous_deserialization=True)
+
+    all_docs = db.similarity_search("", k=db.index.ntotal)
+
+    bm25_retriever = BM25Retriever.from_documents(all_docs)
+    bm25_retriever.k = num_docs  # How many results you want
+
+    faiss_retriever = db.as_retriever(search_kwargs={'k': num_docs})
+
+    ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever, faiss_retriever],
+                                           weights=[0.5,0.5])
+    
+    results = ensemble_retriever.invoke(input=query) 
+    return results
+
+
+if __name__ == "__main__":
+    query = "Haustierversicherung"
+    HUGGINGFACEHUB_API_TOKEN = os.getenv('HUGGINGFACEHUB_API_TOKEN')
+    EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL")
+    
+    path_to_vector_db = Path("..")/'vectorstore/faiss-insurance-agent-500'
+
+    results = get_hybrid_search_results(query=query, 
+                                    path_to_db=path_to_vector_db, 
+                                    embedding_model=EMBEDDING_MODEL, 
+                                    hf_api_key=HUGGINGFACEHUB_API_TOKEN)
+    
+    for doc in results:
+        print(doc)
+        print()