Normalized discounted cumulative gain for a single query.
(predicted_ids: list[str], relevant: dict[str, int], k: int = 10)
| 40 | |
| 41 | |
| 42 | def ndcg_at_k(predicted_ids: list[str], relevant: dict[str, int], k: int = 10) -> float: |
| 43 | """Normalized discounted cumulative gain for a single query.""" |
| 44 | dcg = sum( |
| 45 | relevant.get(doc_id, 0) / math.log2(rank + 2) |
| 46 | for rank, doc_id in enumerate(predicted_ids[:k]) |
| 47 | ) |
| 48 | ideal_rels = sorted(relevant.values(), reverse=True)[:k] |
| 49 | idcg = sum(rel / math.log2(rank + 2) for rank, rel in enumerate(ideal_rels)) |
| 50 | return dcg / idcg if idcg > 0 else 0.0 |
| 51 | |
| 52 | |
| 53 | # -------------------------------------------------------------- |