Retrieval Evaluation & Reranking

# script_id: day_092_retrieval_evaluation_reranking/metrics
# An eval set: each query maps to the set of doc IDs a human marked relevant.
# Build this once (10-50 queries is plenty to start) and reuse it on every change.
eval_set = {
    "how do I reset my password": {"doc_12", "doc_47"},
    "what is the refund window": {"doc_03"},
    "do you ship internationally": {"doc_21", "doc_22", "doc_88"},
}

# What your retriever actually returned, in rank order (best first):
retrieved = {
    "how do I reset my password": ["doc_12", "doc_99", "doc_47", "doc_05"],
    "what is the refund window": ["doc_77", "doc_03", "doc_14"],
    "do you ship internationally": ["doc_22", "doc_21", "doc_61", "doc_88"],
}

# script_id: day_092_retrieval_evaluation_reranking/metrics
import math

def precision_at_k(retrieved_ids: list[str], relevant_ids: set[str], k: int) -> float:
    """Of the top-k returned, what fraction are relevant? (less noise = higher)"""
    top_k = retrieved_ids[:k]
    if not top_k:
        return 0.0
    hits = sum(1 for doc_id in top_k if doc_id in relevant_ids)
    return hits / len(top_k)

def recall_at_k(retrieved_ids: list[str], relevant_ids: set[str], k: int) -> float:
    """Of all relevant docs, what fraction did we find in the top-k? (missed = lower)"""
    if not relevant_ids:
        return 0.0
    hits = sum(1 for doc_id in retrieved_ids[:k] if doc_id in relevant_ids)
    return hits / len(relevant_ids)

def mrr(retrieved_ids: list[str], relevant_ids: set[str]) -> float:
    """Mean Reciprocal Rank: 1/rank of the FIRST relevant hit. Rewards getting
    a right answer to the top. 1.0 = first result is relevant; 0.5 = second; etc."""
    for rank, doc_id in enumerate(retrieved_ids, start=1):
        if doc_id in relevant_ids:
            return 1.0 / rank
    return 0.0

def ndcg_at_k(retrieved_ids: list[str], relevant_ids: set[str], k: int) -> float:
    """Normalized Discounted Cumulative Gain: like a graded results page, but the
    further down a relevant hit sits, the less it counts (the 1/log2(rank+1)
    discount). Divided by the best-possible ordering, so it's always 0-to-1."""
    def dcg(ids: list[str]) -> float:
        return sum(
            (1.0 if doc_id in relevant_ids else 0.0) / math.log2(rank + 1)
            for rank, doc_id in enumerate(ids[:k], start=1)
        )
    actual = dcg(retrieved_ids)
    # ideal = all relevant docs ranked first
    ideal_order = list(relevant_ids) + [d for d in retrieved_ids if d not in relevant_ids]
    best = dcg(ideal_order)
    return actual / best if best > 0 else 0.0


# Run the whole eval set and average
def evaluate(eval_set: dict, retrieved: dict, k: int = 5) -> dict:
    queries = list(eval_set)
    return {
        f"precision@{k}": sum(precision_at_k(retrieved[q], eval_set[q], k) for q in queries) / len(queries),
        f"recall@{k}":    sum(recall_at_k(retrieved[q], eval_set[q], k) for q in queries) / len(queries),
        "mrr":            sum(mrr(retrieved[q], eval_set[q]) for q in queries) / len(queries),
        f"ndcg@{k}":      sum(ndcg_at_k(retrieved[q], eval_set[q], k) for q in queries) / len(queries),
    }

print(evaluate(eval_set, retrieved, k=5))
# e.g. {'precision@5': 0.43, 'recall@5': 0.78, 'mrr': 0.61, 'ndcg@5': 0.74}

# script_id: day_092_retrieval_evaluation_reranking/reranking
# pip install sentence-transformers
# `candidates` = your top-20 docs from vector search; each item is (doc_id, text)
from sentence_transformers import CrossEncoder

reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

def rerank(query: str, candidates: list[tuple[str, str]], top_n: int = 5) -> list[str]:
    """Re-score (doc_id, text) candidates against the query; return top_n doc IDs."""
    pairs = [(query, text) for _doc_id, text in candidates]
    scores = reranker.predict(pairs)
    ranked = sorted(zip(scores, candidates), key=lambda x: x[0], reverse=True)
    return [doc_id for _score, (doc_id, _text) in ranked[:top_n]]

# script_id: day_092_retrieval_evaluation_reranking/reranking_managed
# pip install cohere   (SDK v5+ uses ClientV2)
import cohere

co = cohere.ClientV2()  # reads COHERE_API_KEY

def rerank_managed(query: str, candidates: list[tuple[str, str]], top_n: int = 5) -> list[str]:
    res = co.rerank(
        model="rerank-v3.5",
        query=query,
        documents=[text for _doc_id, text in candidates],
        top_n=top_n,
    )
    return [candidates[r.index][0] for r in res.results]

# script_id: day_092_retrieval_evaluation_reranking/reranking
# Compare baseline retrieval vs reranked retrieval on the SAME eval set.
def benchmark_reranking(eval_set, baseline_retrieved, reranked_retrieved, k=5):
    base = evaluate(eval_set, baseline_retrieved, k)
    rer = evaluate(eval_set, reranked_retrieved, k)
    for metric in base:
        delta = rer[metric] - base[metric]
        print(f"{metric:14} {base[metric]:.3f} -> {rer[metric]:.3f}  ({delta:+.3f})")