making our own reranker for faster result

app.py

@@ -11,19 +11,23 @@ from operator import itemgetter
 from fastapi import FastAPI, Depends, HTTPException, Header
 from fastapi.responses import JSONResponse
 
-# Make sure you have these files in a 'utils' folder
 from utils.DocsLoader import load_and_chunk
 from utils.Schemas import RunRequest, RunResponse
 # from concurrent.futures import ThreadPoolExecutor
 from langchain_community.vectorstores import FAISS
 from langchain.schema import Document
 from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_huggingface import HuggingFaceEmbeddings # Correct new import
+from langchain_huggingface import HuggingFaceEmbeddings 
 # from langchain_chroma import Chroma
 from langchain_community.retrievers import BM25Retriever
-from langchain.retrievers import EnsembleRetriever, ContextualCompressionRetriever
-from langchain.retrievers.document_compressors import CrossEncoderReranker 
-from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+from langchain.retrievers import EnsembleRetriever
+from sklearn.metrics.pairwise import cosine_similarity
+import numpy as np
+
+### to make it faster we are now using our built reranker thats why commenting the imports below
+# from langchain.retrievers import ContextualCompressionRetriever
+# from langchain.retrievers.document_compressors import CrossEncoderReranker 
+# from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 from langchain.prompts import PromptTemplate
 
 
@@ -48,9 +52,10 @@ async def lifespan(app: FastAPI):
     
         # Load models into the shared dictionary
         ml_models["embedder"] = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
-        cross_encoder_model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
+        ### to make it faster we are now using our built reranker thats why commenting the code below
+        # cross_encoder_model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
         # cross_encoder_model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-large")
-        ml_models["reranker_compressor"] = CrossEncoderReranker(model=cross_encoder_model, top_n=5)
+        # ml_models["reranker_compressor"] = CrossEncoderReranker(model=cross_encoder_model, top_n=5)
         ml_models["llm"] = ChatGoogleGenerativeAI(model="gemini-2.0-flash", api_key=GOOGLE_API_KEY)
         ml_models["prompt_template"] = PromptTemplate.from_template(
          """
@@ -165,30 +170,50 @@ async def run_hackrx(req: RunRequest):
     keyword_retriever.k = 5
     # dense_retriever = Chroma.from_documents(documents=chunks, embedding=ml_models["embedder"]).as_retriever()
     ensemble_retriever = EnsembleRetriever(retrievers=[keyword_retriever, dense_retriever], weights=[0.4, 0.6])
-    
-    compression_retriever = ContextualCompressionRetriever(
-        base_retriever=ensemble_retriever, base_compressor=ml_models["reranker_compressor"]
-    )
+    ### to make it faster we are now using our built reranker thats why commenting the code below
     # compression_retriever = ContextualCompressionRetriever(
     #     base_retriever=ensemble_retriever, base_compressor=ml_models["reranker_compressor"]
     # )
 
-    # Define the RAG chain using pre-loaded components
-    hybrid_rag_chain = (
-        {"context": itemgetter("full_query") | compression_retriever, "full_query": itemgetter("full_query")}
-        | ml_models["prompt_template"]
-        | ml_models["llm"]
-    )
     
-    # answers = []
-    # for q in req.questions:
-    #     try:
-    #         result = await hybrid_rag_chain.ainvoke({"full_query": q})
-    #         parsed = parse_llm_response(result.content)
-    #         answers.append(parsed)
-    #     except Exception as e:
-    #         return JSONResponse({"error": str(e)}, status_code=500)
-    tasks = [hybrid_rag_chain.ainvoke({"full_query": q}) for q in req.questions]
+    # Define the RAG chain using pre-loaded components
+    # hybrid_rag_chain = (
+    #     {"context": itemgetter("full_query") | compression_retriever, "full_query": itemgetter("full_query")}
+    #     | ml_models["prompt_template"]
+    #     | ml_models["llm"]
+    # )
+
+    ######## OUR SELF RERANKER ######################################################################
+    #Embed all questions at once
+    question_embeddings = ml_models["embedder"].embed_documents(req.questions)
+
+    # For each question, retrieve and rerank with cosine
+    retrieved_chunks_all = []
+    for i, question in enumerate(req.questions):
+        docs = ensemble_retriever.get_relevant_documents(question)
+        doc_texts = [doc.page_content for doc in docs]
+
+        doc_embeddings = ml_models["embedder"].embed_documents(doc_texts)
+        sims = cosine_similarity([question_embeddings[i]], doc_embeddings)[0]
+
+        top_k = 5
+        top_indices = np.argsort(sims)[-top_k:][::-1]
+        top_chunks = [doc_texts[j] for j in top_indices]
+
+        # Join for context
+        joined_context = "\n\n".join(top_chunks)
+        retrieved_chunks_all.append(joined_context)
+
+    ####################################################################################################################
+
+    tasks = []
+    for i in range(len(req.questions)):
+        prompt_input = {
+            "full_query": req.questions[i],
+            "context": retrieved_chunks_all[i]
+        }
+        tasks.append(ml_models["llm"].ainvoke(ml_models["prompt_template"].format_prompt(**prompt_input)))
+    # tasks = [hybrid_rag_chain.ainvoke({"full_query": q}) for q in req.questions]
     results = await asyncio.gather(*tasks)
     
     # Extract the content from each result and parse it