Update app.py

app.py

@@ -1,95 +1,86 @@
 import gradio as gr
-import faiss
-import numpy as np
 from datasets import load_dataset
-from sentence_transformers import SentenceTransformer
-from sklearn.metrics import ndcg_score
-
-# ----------------------------
-# Load dataset (MS MARCO v1.1)
-# ----------------------------
-dataset = load_dataset("ms_marco", "v1.1", split="train[:10000]")
-passages = [item["passage"] for item in dataset]
-print(f"Loaded {len(passages)} passages")
+from sentence_transformers import SentenceTransformer, util
+import numpy as np
 
-# ----------------------------
 # Load SBERT model
-# ----------------------------
 model = SentenceTransformer("sentence-transformers/all-mpnet-base-v2")
 
-# ----------------------------
-# Build FAISS index
-# ----------------------------
-embeddings = model.encode(passages, convert_to_numpy=True, show_progress_bar=True)
-dimension = embeddings.shape[1]
-index = faiss.IndexFlatL2(dimension)
-index.add(embeddings)
-print("FAISS index built with", index.ntotal, "passages")
-
-# ----------------------------
-# Search function
-# ----------------------------
-def search(query, k=10):
-    query_vec = model.encode([query], convert_to_numpy=True)
-    distances, indices = index.search(query_vec, k)
-    results = [(passages[i], float(dist)) for i, dist in zip(indices[0], distances[0])]
-    return results
-
-# ----------------------------
-# Evaluation metrics
-# ----------------------------
-def evaluate(query, relevant_passages, k=10):
-    """Compute IR metrics for a query given a list of relevant passages (ground truth)."""
-    results = search(query, k)
-    retrieved = [res[0] for res in results]
-    
-    # Binary relevance vector
-    y_true = [1 if p in relevant_passages else 0 for p in retrieved]
-    y_true_full = np.array([[1 if passages[i] in relevant_passages else 0 for i in range(len(passages))]])
-    y_scores_full = np.zeros((1, len(passages)))
-    for idx, (res, dist) in enumerate(results):
-        pos = passages.index(res)
-        y_scores_full[0, pos] = 1.0 - dist  # higher score = more relevant
-    
-    # Metrics
-    precision = sum(y_true) / k
-    recall = sum(y_true) / len(relevant_passages) if relevant_passages else 0
-    f1 = (2 * precision * recall) / (precision + recall) if (precision+recall) > 0 else 0
-    mrr = 1.0 / (y_true.index(1)+1) if 1 in y_true else 0
-    ndcg = ndcg_score(y_true_full, y_scores_full, k=k)
-    
-    return {
-        "Precision@10": round(precision, 3),
-        "Recall@10": round(recall, 3),
-        "F1": round(f1, 3),
-        "MRR": round(mrr, 3),
-        "nDCG@10": round(ndcg, 3)
+# ✅ Load dataset with passages
+dataset = load_dataset("sentence-transformers/msmarco", "v1.1", split="train[:10000]")
+passages = dataset["passage"]
+
+# Encode passages once for efficiency
+passage_embeddings = model.encode(passages, convert_to_tensor=True)
+
+# ---------- Evaluation Metrics ----------
+def precision_at_k(relevant, retrieved, k):
+    return len(set(relevant) & set(retrieved[:k])) / k
+
+def recall_at_k(relevant, retrieved, k):
+    return len(set(relevant) & set(retrieved[:k])) / len(relevant) if relevant else 0
+
+def f1_at_k(relevant, retrieved, k):
+    p = precision_at_k(relevant, retrieved, k)
+    r = recall_at_k(relevant, retrieved, k)
+    return 2*p*r / (p+r) if (p+r) > 0 else 0
+
+def mrr(relevant, retrieved):
+    for i, r in enumerate(retrieved):
+        if r in relevant:
+            return 1 / (i+1)
+    return 0
+
+def ndcg_at_k(relevant, retrieved, k):
+    dcg = 0
+    for i, r in enumerate(retrieved[:k]):
+        if r in relevant:
+            dcg += 1 / np.log2(i+2)
+    ideal_dcg = sum(1 / np.log2(i+2) for i in range(min(len(relevant), k)))
+    return dcg / ideal_dcg if ideal_dcg > 0 else 0
+
+# ---------- Search ----------
+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
+
+# ---------- Interface Logic ----------
+def search_and_evaluate(query):
+    results, retrieved = semantic_search(query, top_k=10)
+
+    # Example: assume top-3 are relevant (for demo purposes)
+    relevant = set(retrieved[:3])
+
+    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),
     }
 
-# ----------------------------
-# Gradio interface
-# ----------------------------
-def gradio_interface(query, relevant_texts):
-    results = search(query, k=10)
-    metrics = {}
-    if relevant_texts.strip():
-        relevant_passages = [t.strip() for t in relevant_texts.split("\n") if t.strip()]
-        metrics = evaluate(query, relevant_passages, k=10)
-    return results, metrics
-
-demo = gr.Interface(
-    fn=gradio_interface,
-    inputs=[
-        gr.Textbox(label="Enter your query"),
-        gr.Textbox(label="Enter relevant passages (ground truth, one per line)", placeholder="Optional")
-    ],
-    outputs=[
-        gr.Dataframe(headers=["Passage", "Distance"], label="Top-10 Results"),
-        gr.Label(label="Evaluation Metrics")
-    ],
-    title="SBERT + FAISS Semantic Search",
-    description="Enter a query to search MS MARCO passages. Optionally provide ground truth passages to compute IR metrics."
+    output_text = "### Search Results:\n"
+    for i, (text, score) in enumerate(results, 1):
+        output_text += f"{i}. {text} (score: {score:.4f})\n\n"
+
+    output_text += "\n### Evaluation Metrics:\n"
+    for k, v in metrics.items():
+        output_text += f"{k}: {v:.4f}\n"
+
+    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."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch()