added rerank model options

app.py

@@ -19,7 +19,10 @@ from huggingface_hub import dataset_info
 
 # Load embedding model
 QUERY_EMBEDDING_MODEL = SentenceTransformer('all-MiniLM-L6-v2')
-RERANKING_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
+RERANKING_MODELS = {
+    "MS MARCO MiniLM": "cross-encoder/ms-marco-MiniLM-L-6-v2",
+    "MonoT5 Base": "castorini/monot5-base-msmarco",
+}
 PROMPT_MODEL = "llama-3.3-70b-specdec"
 EVAL_MODEL = "llama-3.3-70b-specdec"
 WINDOW_SIZE = 5
@@ -107,6 +110,12 @@ with gr.Blocks() as iface:
             label="Select a Model"
         )
 
+        reranker_dropdown = gr.Dropdown(
+            list(RERANKING_MODELS.keys()),
+            value="MS MARCO MiniLM",
+            label="Select Reranking Model"
+        )
+
     submit_button = gr.Button("Evaluate Model")
 
     with gr.Row():

finetuneresults.py

@@ -1,5 +1,62 @@
 from sentence_transformers import CrossEncoder
-
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
+import numpy as np
+from typing import List, Tuple
+
+class MonoT5Reranker:
+    def __init__(self, model_name: str):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"Using device: {self.device}")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
+        self.model.to(self.device)
+        self.model.eval()
+        
+    def predict(self, query_doc_pairs: List[Tuple[str, str]]) -> np.ndarray:
+        scores = []
+        batch_size = 8  # Adjust based on your GPU/CPU memory
+        
+        for i in range(0, len(query_doc_pairs), batch_size):
+            batch_pairs = query_doc_pairs[i:i + batch_size]
+            
+            # Format input as per MonoT5 requirements
+            inputs = [f"Query: {query} Document: {doc}" for query, doc in batch_pairs]
+            
+            # Tokenize
+            encoded = self.tokenizer(
+                inputs,
+                padding=True,
+                truncation=True,
+                max_length=512,
+                return_tensors="pt"
+            ).to(self.device)
+            
+            # Get predictions
+            with torch.no_grad():
+                outputs = self.model(**encoded)
+                batch_scores = outputs.logits.squeeze(-1).cpu().numpy()
+                scores.extend(batch_scores.tolist())
+        
+        return np.array(scores)
+
+class MSMARCOReranker:
+    def __init__(self, model_name: str):
+        self.model = CrossEncoder(model_name)
+        
+    def predict(self, query_doc_pairs: List[Tuple[str, str]]) -> np.ndarray:
+        return self.model.predict(query_doc_pairs)
+
+
+def get_reranker(model_name: str):
+    """Factory function to get appropriate reranker based on model name."""
+    if "monot5" in model_name.lower():
+        print(f"Using MonoT5 reranker: {model_name}")
+        return MonoT5Reranker(model_name)
+    else:
+        print(f"Using MS MARCO reranker: {model_name}")
+        return MSMARCOReranker(model_name)
+    
 """
     Retrieves unique full documents based on the top-ranked document IDs.
 
@@ -37,25 +94,35 @@ Returns:
 """
 
 def rerank_documents(query, retrieved_docs_df, model_name="cross-encoder/ms-marco-MiniLM-L-6-v2"):
+    """Reranks documents using the specified reranking model."""
+    try:
+        # Load Cross-Encoder model
+        model = get_reranker(model_name)
 
-    # Load Cross-Encoder model
-    model = CrossEncoder(model_name)
-
-    # Prepare query-document pairs
-    query_doc_pairs = [(query, " ".join(doc)) for doc in retrieved_docs_df["document"]]
+        # Prepare query-document pairs
+        query_doc_pairs = [(query, " ".join(doc)) for doc in retrieved_docs_df["document"]]
 
-    # Compute relevance scores
-    scores = model.predict(query_doc_pairs)
+        # Compute relevance scores
+        scores = model.predict(query_doc_pairs)
 
-    # Add scores to the DataFrame
-    retrieved_docs_df["relevance_score"] = scores
+        # Add scores to the DataFrame
+        retrieved_docs_df["relevance_score"] = scores
 
-    # Sort by score in descending order (higher score = more relevant)
-    reranked_docs_df = retrieved_docs_df.sort_values(by="relevance_score", ascending=False).reset_index(drop=True)
+        # Sort by score in descending order (higher score = more relevant)
+        reranked_docs_df = retrieved_docs_df.sort_values(by="relevance_score", ascending=False).reset_index(drop=True)
 
-    return reranked_docs_df
+        return reranked_docs_df
+    except Exception as e:
+            print(f"Error in reranking: {e}")
+            # Return original order if reranking fails
+            retrieved_docs_df["relevance_score"] = 1.0
+            return retrieved_docs_df
 
 def FineTuneAndRerankSearchResults(top_10_chunk_results, rag_extarcted_data, question, reranking_model):
-    unique_docs= retrieve_full_documents(top_10_chunk_results, rag_extarcted_data)
-    reranked_results = rerank_documents(question, unique_docs, reranking_model)
-    return reranked_results
+    try:
+        unique_docs= retrieve_full_documents(top_10_chunk_results, rag_extarcted_data)
+        reranked_results = rerank_documents(question, unique_docs, reranking_model)
+        return reranked_results
+    except Exception as e:
+        print(f"Error in FineTuneAndRerankSearchResults: {e}")
+        return None