Code Changes - Upload Results

app.py

@@ -384,7 +384,7 @@ def _evaluate_using_groq(context_docs, question, generated_answer):
 
 
 def run_rag_pipeline(subset, question, custom_question, chunking, embed_model, retriever, chunk_count, retriever_type,
-                     reranking, evaluator):
+                     reranking, reranking_dropdown, evaluator):
     final_question = custom_question if custom_question.strip() else question
     print(f"The query is {final_question}")
 
@@ -395,13 +395,17 @@ def run_rag_pipeline(subset, question, custom_question, chunking, embed_model, r
     logging.info("Starting RAG Pipeline using logging")
     gr.Info("Starting RAG Pipeline using gradio")  # Shows as a toast notification in UI
 
+    ranking_method = reranking_dropdown if reranking else None
+    print(f"Using reranking: {reranking}, method: {ranking_method}")
+
     rag = RAGSystem(
             subset=subset,
             dataset_type="test",
             strategy=chunking,
             chunks=[],  # Not needed for loading
             generator_model_name=retriever,
-            retriever_model_name=embed_model
+            retriever_model_name=embed_model,
+            reranker_model_name=ranking_method
         )
 
     # 3. Load or use stored vector DB
@@ -621,7 +625,7 @@ with gr.Blocks(
             subset_dropdown, question_dropdown, custom_question_input,
             chunking_dropdown, embed_dropdown, generator_dropdown,
             chunk_count, retriever_type,
-            reranking_checkbox, evaluator_dropdown
+            reranking_checkbox, reranking_dropdown, evaluator_dropdown
         ],
         outputs=[gen_ans_display, y_pred_metrics_display, chunks_retrieved_display, evaluator_json_output, download_file]
     )

ragbench.py

@@ -65,6 +65,7 @@ class RAGSystem:
         chunk_overlap: int = 50,
         generator_model_name: str = "mistralai/Mistral-7B-Instruct-v0.2",
         retriever_model_name: str = "BAAI/bge-large-en-v1.5",
+        reranker_model_name: str = "cross-encoder/ms-marco-MiniLM-L-12-v2",
         hf_api_token: str = None
     ):
         self.subset = subset
@@ -74,6 +75,7 @@ class RAGSystem:
         self.chunk_overlap = chunk_overlap
         self.generator_model_name = generator_model_name
         self.retriever_model_name = retriever_model_name
+        self.reranker_model_name = reranker_model_name
         self.chunks = chunks
         self.hf_api_token = hf_api_token or os.getenv("HF_API_TOKEN")
 
@@ -377,24 +379,58 @@ class RAGSystem:
         except Exception as e:
             print(f"Generation failed: {str(e)}")
             return "I couldn't generate an answer."
+        
+    def _use_reranker(self, docs: List[LangchainDocument], query: str, top_k: int) -> List[LangchainDocument]:
+        """Use the reranker model to re-rank retrieved documents"""
+        if not self.reranker_model_name:
+            return docs
+        
+        sentence_chunks = []
+        for doc in docs:
+            for sentence in doc.page_content.strip().split("."):
+                sentence = sentence.strip()
+                if len(sentence) > 15:
+                    sentence_chunks.append((sentence, doc.metadata))
+
+        pairs = [[query, sent] for sent, _ in sentence_chunks]
+        scores = self.reranker.predict(pairs)
+
+        top_pairs = sorted(zip(sentence_chunks, scores), key=lambda x: x[1], reverse=True)[:top_k]
+
+        top_chunks = []
+        for (sentence, meta), score in top_pairs:
+            top_chunks.append(Chunk(
+                chunk_id=meta.get("chunk_id", ""),
+                text=sentence,
+                doc_id=meta.get("doc_id", ""),
+                source=meta.get("source", ""),
+                chunk_num=meta.get("chunk_num", -1),
+                total_chunks=meta.get("total_chunks", -1),
+                metadata={**meta, "reranker_score": score}
+            ))
+            
+        print(f"Reranked {len(top_chunks)} chunks from {len(docs)} documents")
+        return top_chunks
 
     def retrieve(self, query: str, top_k: int = 10) -> List[Chunk]:
         """Retrieve relevant chunks using HYDE"""
         pseudo_answer = self.generate_hypothetical_answer(query)
         docs = self.hybrid_retriever.invoke(pseudo_answer)
 
-        
-        return [
-            Chunk(
-                chunk_id=doc.metadata.get("chunk_id", ""),
-                text=doc.page_content,
-                doc_id=doc.metadata.get("doc_id", ""),
-                source=doc.metadata.get("source", ""),
-                chunk_num=doc.metadata.get("chunk_num", -1),
-                total_chunks=doc.metadata.get("total_chunks", -1),
-                metadata=doc.metadata
-            ) for doc in docs[:top_k]
-        ]
+        if self.reranker_model_name is not None:
+            return self._use_reranker(docs, query, top_k)
+        else:
+            return [
+                Chunk(
+                    chunk_id=doc.metadata.get("chunk_id", ""),
+                    text=doc.page_content,
+                    doc_id=doc.metadata.get("doc_id", ""),
+                    source=doc.metadata.get("source", ""),
+                    chunk_num=doc.metadata.get("chunk_num", -1),
+                    total_chunks=doc.metadata.get("total_chunks", -1),
+                    metadata=doc.metadata
+                ) for doc in docs[:top_k]
+            ]
 
     def generate(self, question: str, context: List[str] = None) -> str:
         """Generate final answer with RAG context"""