Update app.py

app.py

@@ -3,17 +3,28 @@ from datasets import load_dataset
 from sentence_transformers import SentenceTransformer, util
 import numpy as np
 
-# Load SBERT model
+# ---------- Load model ----------
 model = SentenceTransformer("sentence-transformers/all-mpnet-base-v2")
 
-# ✅ Load dataset with passages
+# ---------- Load MS MARCO dataset ----------
+# 10k sample passages
 dataset = load_dataset("sentence-transformers/msmarco", "v1.1", split="train[:10000]")
 passages = dataset["passage"]
 
-# Encode passages once for efficiency
+# Precompute embeddings
 passage_embeddings = model.encode(passages, convert_to_tensor=True)
 
-# ---------- Evaluation Metrics ----------
+# Map index -> passage
+id_to_passage = {i: passages[i] for i in range(len(passages))}
+
+# ---------- Load queries and qrels ----------
+queries_dataset = load_dataset("sentence-transformers/msmarco", "v1.1", split="validation[:500]")  # small sample
+qrels_dataset = load_dataset("ms_marco", "v1.1", split="validation[:500]")  # contains relevant passage ids
+
+query_id_to_text = {i: q["query"] for i, q in enumerate(queries_dataset)}
+query_id_to_relevant = {i: set(q["positive_passages"]) for i, q in enumerate(qrels_dataset)}
+
+# ---------- Evaluation metrics ----------
 def precision_at_k(relevant, retrieved, k):
     return len(set(relevant) & set(retrieved[:k])) / k
 
@@ -44,23 +55,27 @@ def semantic_search(query, top_k=10):
     query_embedding = model.encode(query, convert_to_tensor=True)
     scores = util.cos_sim(query_embedding, passage_embeddings)[0]
     top_results = scores.topk(k=top_k)
-    retrieved = [int(idx) for idx in top_results[1]]
-    results = [(passages[idx], float(scores[idx])) for idx in retrieved]
-    return results, retrieved
+    retrieved_indices = [int(idx) for idx in top_results[1]]
+    results = [(id_to_passage[idx], float(scores[idx])) for idx in retrieved_indices]
+    return results, retrieved_indices
 
-# ---------- Interface Logic ----------
+# ---------- Gradio interface ----------
 def search_and_evaluate(query):
-    results, retrieved = semantic_search(query, top_k=10)
+    results, retrieved_indices = semantic_search(query, top_k=10)
 
-    # Example: assume top-3 are relevant (for demo purposes)
-    relevant = set(retrieved[:3])
+    # Match against actual relevant passages if available
+    relevant_indices = set()
+    for i, q in query_id_to_text.items():
+        if q.strip().lower() == query.strip().lower():
+            relevant_indices = query_id_to_relevant[i]
+            break
 
     metrics = {
-        "Precision@10": precision_at_k(relevant, retrieved, 10),
-        "Recall@10": recall_at_k(relevant, retrieved, 10),
-        "F1@10": f1_at_k(relevant, retrieved, 10),
-        "MRR": mrr(relevant, retrieved),
-        "nDCG@10": ndcg_at_k(relevant, retrieved, 10),
+        "Precision@10": precision_at_k(relevant_indices, retrieved_indices, 10),
+        "Recall@10": recall_at_k(relevant_indices, retrieved_indices, 10),
+        "F1@10": f1_at_k(relevant_indices, retrieved_indices, 10),
+        "MRR": mrr(relevant_indices, retrieved_indices),
+        "nDCG@10": ndcg_at_k(relevant_indices, retrieved_indices, 10)
     }
 
     output_text = "### Search Results:\n"
@@ -73,13 +88,12 @@ def search_and_evaluate(query):
 
     return output_text
 
-# ---------- Gradio App ----------
 iface = gr.Interface(
     fn=search_and_evaluate,
     inputs=gr.Textbox(label="Enter your query"),
     outputs=gr.Textbox(label="Results + Metrics"),
     title="SBERT Semantic Search + Evaluation Metrics",
-    description="Semantic search on MS MARCO (10,000 sample passages) using all-mpnet-base-v2. Includes Precision@10, Recall@10, F1, MRR, nDCG@10."
+    description="Semantic search on MS MARCO (10,000 sample passages) using all-mpnet-base-v2 with true evaluation metrics."
 )
 
 if __name__ == "__main__":