Spaces:
Runtime error
Runtime error
File size: 3,009 Bytes
361de4d bf736ad 361de4d bf736ad 361de4d bf736ad ae43f82 bf736ad |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import gradio as gr
from datasets import load_dataset
from sentence_transformers import SentenceTransformer, util
import numpy as np
# Load SBERT model
model = SentenceTransformer("sentence-transformers/all-mpnet-base-v2")
# ✅ 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),
}
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__":
iface.launch()
|