Update app.py

app.py

@@ -1,122 +1,95 @@
 import gradio as gr
-from sentence_transformers import SentenceTransformer
 import faiss
 import numpy as np
 from datasets import load_dataset
-from sklearn.metrics import precision_score, recall_score
-import pandas as pd
-import os
+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")
 
+# ----------------------------
 # Load SBERT model
+# ----------------------------
 model = SentenceTransformer("sentence-transformers/all-mpnet-base-v2")
 
-# Load MS MARCO dataset (10,000 passages for demo)
-dataset = load_dataset("ms_marco", "passage", split="train[:10000]")
-passages = dataset["passage"]
-passage_embeddings = model.encode(passages, convert_to_numpy=True, show_progress_bar=True)
-
+# ----------------------------
 # Build FAISS index
-dimension = passage_embeddings.shape[1]
+# ----------------------------
+embeddings = model.encode(passages, convert_to_numpy=True, show_progress_bar=True)
+dimension = embeddings.shape[1]
 index = faiss.IndexFlatL2(dimension)
-index.add(passage_embeddings)
-
-# CSV log file
-log_file = "results_log.csv"
-if not os.path.exists(log_file):
-    pd.DataFrame(columns=["Query", "Relevant Passage", "Precision@10", "Recall@10", "F1@10", "MRR", "nDCG@10"]).to_csv(log_file, index=False)
+index.add(embeddings)
+print("FAISS index built with", index.ntotal, "passages")
 
+# ----------------------------
 # Search function
-def semantic_search(query, k=10):
+# ----------------------------
+def search(query, k=10):
     query_vec = model.encode([query], convert_to_numpy=True)
     distances, indices = index.search(query_vec, k)
-    results = [passages[i] for i in indices[0]]
-    return "\n\n".join(results)
-
-# Helper functions for metrics
-def mean_reciprocal_rank(y_true):
-    for rank, rel in enumerate(y_true, start=1):
-        if rel == 1:
-            return 1 / rank
-    return 0
-
-def ndcg_at_k(y_true, k=10):
-    y_true = np.array(y_true)[:k]
-    if y_true.sum() == 0:
-        return 0.0
-    dcg = np.sum((2**y_true - 1) / np.log2(np.arange(2, len(y_true) + 2)))
-    ideal = np.sort(y_true)[::-1]
-    idcg = np.sum((2**ideal - 1) / np.log2(np.arange(2, len(ideal) + 2)))
-    return dcg / idcg if idcg > 0 else 0.0
-
-# Evaluation function
-def evaluate(query, relevant_passage, k=10):
-    results_text = semantic_search(query, k)
-    results_list = results_text.split("\n\n")
-
-    y_true = [1 if relevant_passage in r else 0 for r in results_list]
-    y_pred = [1] * len(results_list)
-
-    precision = precision_score(y_true, y_pred, zero_division=0)
-    recall = recall_score(y_true, y_pred, zero_division=0)
-    f1 = 2 * (precision * recall) / (precision + recall + 1e-9)
-    mrr = mean_reciprocal_rank(y_true)
-    ndcg = ndcg_at_k(y_true, k)
-
-    # Log results to CSV
-    df = pd.read_csv(log_file)
-    new_row = {
-        "Query": query,
-        "Relevant Passage": relevant_passage,
+    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@10": round(f1, 3),
+        "F1": round(f1, 3),
         "MRR": round(mrr, 3),
-        "nDCG@10": round(ndcg, 3),
+        "nDCG@10": round(ndcg, 3)
     }
-    df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
-    df.to_csv(log_file, index=False)
-
-    return new_row, f"📊 {len(df)} evaluations logged so far."
-
-# Function to download CSV log
-def download_log():
-    return log_file
-
-# Function to check current log count
-def check_log_count():
-    df = pd.read_csv(log_file)
-    return f"📊 {len(df)} evaluations logged so far."
 
+# ----------------------------
 # Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("## 🔎 Semantic Search with SBERT (MS MARCO Subset)")
-
-    with gr.Row():
-        query_input = gr.Textbox(label="Enter your search query")
-        k_input = gr.Slider(1, 20, value=10, step=1, label="Top-K Results")
-
-    results_output = gr.Textbox(label="Search Results", lines=10)
-    run_btn = gr.Button("Search")
-
-    run_btn.click(fn=semantic_search, inputs=[query_input, k_input], outputs=results_output)
-
-    gr.Markdown("### 📊 Evaluation")
-    with gr.Row():
-        eval_query = gr.Textbox(label="Evaluation Query")
-        relevant_passage = gr.Textbox(label="Known Relevant Passage")
-    eval_btn = gr.Button("Run Evaluation")
-    eval_output = gr.JSON(label="Evaluation Metrics")
-    eval_counter = gr.Label(label="Evaluation Log Count")
-
-    eval_btn.click(fn=evaluate, inputs=[eval_query, relevant_passage, k_input], outputs=[eval_output, eval_counter])
-
-    gr.Markdown("### 📂 Download Logged Results")
-    download_btn = gr.Button("Download CSV")
-    file_output = gr.File()
-    count_btn = gr.Button("Check Log Count")
-    count_output = gr.Label(label="Evaluation Log Count")
-
-    download_btn.click(fn=download_log, outputs=file_output)
-    count_btn.click(fn=check_log_count, outputs=count_output)
-
-demo.launch()
+# ----------------------------
+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."
+)
+
+if __name__ == "__main__":
+    demo.launch()