Update app.py

app.py

@@ -203,6 +203,35 @@ 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)
+    
+        # 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]
+    
+        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]