Merge pull request #4 from almutareb/reranking

core-langchain-rag.py

@@ -214,7 +214,7 @@ def generate_qa_retriever(history: dict, question: str, llm_model:HuggingFaceEnd
     template = """
     You are a friendly insurance product advisor, your task is to help customers find the best products from Württembergische GmbH.\
     You help the user find the answers to all his questions. Answer in short and simple terms and offer to explain the product and terms to the user.\
-    Respond only using the provided context (delimited by <ctx></ctx>) and only in German or Englisch, depending on the question's language.
+    Respond only using the provided context (delimited by <ctx></ctx>) and only in German or English, depending on the question's language.
     Use the chat history (delimited by <hs></hs>) to help find the best product for the user:
     ------
     <ctx>

rag_app/get_db_retriever.py

@@ -26,4 +26,5 @@ def get_db_retriever(vector_db:str=None):
 
     retriever = db.as_retriever()
 
-    return retriever
+    return retriever
+

rag_app/loading_data/load_S3_vector_stores.py

@@ -10,6 +10,7 @@ from dotenv import load_dotenv
 import os
 import sys
 import logging
+from pathlib import Path
 
 # Load environment variables from a .env file
 config = load_dotenv(".env")
@@ -38,6 +39,7 @@ def get_faiss_vs():
 
     # Define the destination for the downloaded file
     VS_DESTINATION = FAISS_INDEX_PATH + ".zip"
+    
     try:
         # Download the pre-prepared vectorized index from the S3 bucket
         print("Downloading the pre-prepared FAISS vectorized index from S3...")
@@ -51,7 +53,32 @@ def get_faiss_vs():
         
     except Exception as e:
         print(f"Error during downloading or extracting from S3: {e}", file=sys.stderr)
-    #faissdb = FAISS.load_local(FAISS_INDEX_PATH, embeddings)
+    # faissdb = FAISS.load_local(FAISS_INDEX_PATH, embeddings)
+
+
+def get_faiss_vs_from_s3(s3_loc:str, 
+                         s3_vs_name:str,
+                         vs_dir:str='vectorstore') -> None:
+    """ Download the FAISS vector store from S3 bucket
+
+        Args:
+            s3_loc (str): Name of the S3 bucket
+            s3_vs_name (str): Name of the file to be downloaded
+            vs_dir (str): The name of the directory where the file is to be saved
+    """
+    # Initialize an S3 client with unsigned configuration for public access
+    s3 = boto3.client('s3', config=Config(signature_version=UNSIGNED))
+    # Destination folder
+    vs_dir_path = Path("..") / vs_dir
+    assert vs_dir_path.is_dir(), "Cannot find vs_dir folder"
+    try:
+        vs_destination = Path("..") / vs_dir / "faiss-insurance-agent-500.zip"
+        s3.download_file(s3_loc, s3_vs_name, vs_destination)
+        # Extract the downloaded zip file
+        with zipfile.ZipFile(file=vs_destination, mode='r') as zip_ref:
+            zip_ref.extractall(path=vs_dir_path.as_posix())
+    except Exception as e:
+        print(f"Error during downloading or extracting from S3: {e}", file=sys.stderr)
 
 
 ## Chroma DB
@@ -70,4 +97,10 @@ def get_chroma_vs():
         chromadb = Chroma(persist_directory=CHROMA_DIRECTORY, embedding_function=embeddings)
         chromadb.get()
     except Exception as e:
-        print(f"Error during downloading or extracting from S3: {e}", file=sys.stderr)
+        print(f"Error during downloading or extracting from S3: {e}", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    # get_faiss_vs_from_s3(s3_loc=S3_LOCATION, s3_vs_name=FAISS_VS_NAME)
+    pass
+    

rag_app/reranking.py

@@ -0,0 +1,80 @@
+# from get_db_retriever import get_db_retriever
+from pathlib import Path
+from langchain_community.vectorstores import FAISS
+from dotenv import load_dotenv
+import os
+from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
+import requests
+
+load_dotenv()
+
+
+def get_reranked_docs(query:str, 
+                      path_to_db:str, 
+                      embedding_model:str,
+                      hf_api_key:str, 
+                      num_docs:int=5) -> list:
+    """ Re-ranks the similarity search results and returns top-k highest ranked docs
+
+        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: A list of documents with the highest rank
+    """
+    assert num_docs <= 10, "num_docs should be less than similarity search results"
+    
+    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)
+    
+    # Get 10 documents based on similarity search
+    docs =  db.similarity_search(query=query, k=10)
+
+    # Add the page_content, description and title together
+    passages = [doc.page_content + "\n" + doc.metadata.get('title', "") +"\n"+ doc.metadata.get('description', "") 
+                for doc in docs]
+    
+    # Prepare the payload
+    inputs = [{"text": query, "text_pair": passage} for passage in passages]
+
+    API_URL = "https://api-inference.huggingface.co/models/deepset/gbert-base-germandpr-reranking"
+    headers = {"Authorization": f"Bearer {hf_api_key}"}
+
+    response = requests.post(API_URL, headers=headers, json=inputs)
+    scores = response.json()
+    
+    try:
+        relevance_scores = [item[1]['score'] for item in scores]
+    except ValueError as e:
+        print('Could not get the relevance_scores -> something might be wrong with the json output')
+        return 
+    
+    if relevance_scores:
+        ranked_results = sorted(zip(docs, passages, relevance_scores), key=lambda x: x[2], reverse=True)
+        top_k_results = ranked_results[:num_docs]
+        return [doc for doc, _, _ in top_k_results]
+
+ 
+if __name__ == "__main__":
+    
+    HUGGINGFACEHUB_API_TOKEN = os.getenv('HUGGINGFACEHUB_API_TOKEN')
+    EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL")
+    
+    path_to_vector_db = Path("..")/'vectorstore/faiss-insurance-agent-500'
+
+    query = "Ich möchte wissen, ob ich meine geriatrische Haustier-Eidechse versichern kann"
+    
+    top_5_docs = get_reranked_docs(query=query,
+                                   path_to_db=path_to_vector_db,
+                                   embedding_model=EMBEDDING_MODEL,
+                                   hf_api_key=HUGGINGFACEHUB_API_TOKEN,
+                                   num_docs=5)
+    
+    for i, doc in enumerate(top_5_docs):
+        print(f"{i}: {doc}\n")