File size: 1,678 Bytes
6ced533 | 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 | """Retrieval quality metrics for graded relevance judgments."""
import math
def recall_at_k(retrieved_ids, relevance, k):
"""Fraction of all known-relevant items that appear in the top-k results."""
relevant = {doc_id for doc_id, rel in relevance.items() if rel > 0}
if not relevant:
return None
top_k = set(retrieved_ids[:k])
return len(top_k & relevant) / len(relevant)
def reciprocal_rank(retrieved_ids, relevance):
"""1 / rank of the first retrieved item with relevance > 0, else 0."""
relevant = {doc_id for doc_id, rel in relevance.items() if rel > 0}
for rank, doc_id in enumerate(retrieved_ids, start=1):
if doc_id in relevant:
return 1.0 / rank
return 0.0
def _dcg(gains):
return sum(gain / math.log2(pos + 1) for pos, gain in enumerate(gains, start=1))
def ndcg_at_k(retrieved_ids, relevance, k):
"""Normalized DCG@k using the standard 2^rel - 1 gain function."""
gains = [
(2 ** relevance.get(doc_id, 0)) - 1
for doc_id in retrieved_ids[:k]
]
dcg = _dcg(gains)
ideal_gains = sorted(
((2 ** rel) - 1 for rel in relevance.values()),
reverse=True,
)[:k]
idcg = _dcg(ideal_gains)
if idcg == 0:
return 0.0
return dcg / idcg
def evaluate_ranking(retrieved_ids, relevance):
"""Compute all retrieval-quality metrics for a single query's ranking."""
return {
"recall@10": recall_at_k(retrieved_ids, relevance, 10),
"recall@50": recall_at_k(retrieved_ids, relevance, 50),
"mrr": reciprocal_rank(retrieved_ids, relevance),
"ndcg@10": ndcg_at_k(retrieved_ids, relevance, 10),
}
|