Update app.py

app.py

@@ -25,9 +25,9 @@ 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.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import CrossEncoderReranker 
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 from langchain.prompts import PromptTemplate
 
 
@@ -53,9 +53,9 @@ async def lifespan(app: FastAPI):
         # Load models into the shared dictionary
         ml_models["embedder"] = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
         ### 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-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("""
 You are an expert decision maker Assistant in the domain such as insurance, legal compliance, human resources, and contract management.
@@ -172,21 +172,21 @@ async def run_hackrx(req: RunRequest):
     # 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])
     ### 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"]
-    # )
+    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"]
-    # )
+    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)
+    # question_embeddings = ml_models["embedder"].embed_documents(req.questions)
 
     # For each question, retrieve and rerank with cosine
     # retrieved_chunks_all = []
@@ -204,66 +204,66 @@ async def run_hackrx(req: RunRequest):
     #     # Join for context
     #     joined_context = "\n\n".join(top_chunks)
     #     retrieved_chunks_all.append(joined_context)
-    def mmr_select(query_embedding, doc_embeddings, k=6, lambda_mult=0.6):
-        selected = []
-        candidates = list(range(len(doc_embeddings)))
-        doc_embeddings = np.array(doc_embeddings)
+    # def mmr_select(query_embedding, doc_embeddings, k=6, lambda_mult=0.6):
+    #     selected = []
+    #     candidates = list(range(len(doc_embeddings)))
+    #     doc_embeddings = np.array(doc_embeddings)
     
-        # Convert query_embedding to 2D
-        query_embedding = np.array(query_embedding).reshape(1, -1)
+    #     # Convert query_embedding to 2D
+    #     query_embedding = np.array(query_embedding).reshape(1, -1)
     
-        # Compute similarity between query and all documents
-        query_doc_sims = cosine_similarity(query_embedding, doc_embeddings)[0]
+    #     # Compute similarity between query and all documents
+    #     query_doc_sims = cosine_similarity(query_embedding, doc_embeddings)[0]
     
-        for _ in range(k):
-            mmr_score = []
-            for idx in candidates:
-                if not selected:
-                    diversity = 0
-                else:
-                    selected_embeddings = doc_embeddings[selected]
-                    diversity = max(cosine_similarity(
-                        doc_embeddings[idx].reshape(1, -1),
-                        selected_embeddings
-                    )[0])
-                score = lambda_mult * query_doc_sims[idx] - (1 - lambda_mult) * diversity
-                mmr_score.append(score)
-            selected_idx = candidates[np.argmax(mmr_score)]
-            selected.append(selected_idx)
-            candidates.remove(selected_idx)
+    #     for _ in range(k):
+    #         mmr_score = []
+    #         for idx in candidates:
+    #             if not selected:
+    #                 diversity = 0
+    #             else:
+    #                 selected_embeddings = doc_embeddings[selected]
+    #                 diversity = max(cosine_similarity(
+    #                     doc_embeddings[idx].reshape(1, -1),
+    #                     selected_embeddings
+    #                 )[0])
+    #             score = lambda_mult * query_doc_sims[idx] - (1 - lambda_mult) * diversity
+    #             mmr_score.append(score)
+    #         selected_idx = candidates[np.argmax(mmr_score)]
+    #         selected.append(selected_idx)
+    #         candidates.remove(selected_idx)
     
-        return selected
-    async def async_retrieve_and_rerank(question: str, q_idx: int):
-        docs = await ensemble_retriever.ainvoke(question)
-        doc_texts = [doc.page_content for doc in docs]
-        doc_embeddings = ml_models["embedder"].embed_documents(doc_texts)
-        # sims = cosine_similarity([question_embeddings[q_idx]], doc_embeddings)[0]
-        query_embedding = question_embeddings[q_idx]
-        selected_indices = mmr_select(
-        query_embedding=query_embedding,
-        doc_embeddings=doc_embeddings,
-        k=6,
-        lambda_mult=0.6,
-    )
-        # top_indices = np.argsort(sims)[-top_k:][::-1]
-        top_chunks = [doc_texts[j] for j in selected_indices]
-        return "\n\n".join(top_chunks)
-    # Retrieve and rerank all in parallel
-    retrieved_chunks_all = await asyncio.gather(
-        *[async_retrieve_and_rerank(q, i) for i, q in enumerate(req.questions)]
-    ) 
+    #     return selected
+    # async def async_retrieve_and_rerank(question: str, q_idx: int):
+    #     docs = await ensemble_retriever.ainvoke(question)
+    #     doc_texts = [doc.page_content for doc in docs]
+    #     doc_embeddings = ml_models["embedder"].embed_documents(doc_texts)
+    #     # sims = cosine_similarity([question_embeddings[q_idx]], doc_embeddings)[0]
+    #     query_embedding = question_embeddings[q_idx]
+    #     selected_indices = mmr_select(
+    #     query_embedding=query_embedding,
+    #     doc_embeddings=doc_embeddings,
+    #     k=6,
+    #     lambda_mult=0.6,
+    # )
+    #     # top_indices = np.argsort(sims)[-top_k:][::-1]
+    #     top_chunks = [doc_texts[j] for j in selected_indices]
+    #     return "\n\n".join(top_chunks)
+    # # Retrieve and rerank all in parallel
+    # retrieved_chunks_all = await asyncio.gather(
+    #     *[async_retrieve_and_rerank(q, i) for i, q in enumerate(req.questions)]
+    # ) 
 
     ####################################################################################################################
 
-    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)
+    # 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)
     # answers = []
 
     # for msg in results:
@@ -271,6 +271,8 @@ async def run_hackrx(req: RunRequest):
     #     if hasattr(msg, "content"):
     #         answers.append(msg.content.strip())
     #     # Extract the content from each result and parse it
+    tasks = [hybrid_rag_chain.ainvoke({"full_query": q}) for q in req.questions]
+    results = await asyncio.gather(*tasks)
     answers = [parse_llm_response(result.content) for result in results]
 
     return JSONResponse({"answers": answers}, status_code=200)