benchmark/metrics.py

"""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),
    }