Skip to main content

entrenar/eval/generative/
retrieval.rs

1//! Information retrieval evaluation metrics
2//!
3//! Provides NDCG@k (Normalized Discounted Cumulative Gain) for evaluating
4//! ranking quality in search and retrieval systems.
5
6/// Compute NDCG@k (Normalized Discounted Cumulative Gain).
7///
8/// Measures ranking quality by comparing the actual ranking against the
9/// ideal ranking. Returns a value in [0, 1] where 1.0 indicates a
10/// perfect ranking.
11///
12/// # Arguments
13/// * `relevance_scores` - Relevance scores in the order returned by the system
14/// * `k` - Number of top results to consider
15///
16/// # Returns
17/// 0.0 if there are no relevant documents or k is 0.
18pub fn ndcg_at_k(relevance_scores: &[f64], k: usize) -> f64 {
19    contract_pre_ndcg_at_k!();
20    if k == 0 || relevance_scores.is_empty() {
21        return 0.0;
22    }
23
24    let k = k.min(relevance_scores.len());
25
26    let actual_dcg = dcg(&relevance_scores[..k]);
27
28    // Ideal DCG: sort scores descending and compute DCG
29    let mut ideal_scores: Vec<f64> = relevance_scores.to_vec();
30    ideal_scores.sort_by(|a, b| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));
31    let idcg = dcg(&ideal_scores[..k]);
32
33    if idcg == 0.0 {
34        return 0.0;
35    }
36
37    actual_dcg / idcg
38}
39
40/// Compute Discounted Cumulative Gain.
41///
42/// DCG = sum_{i=1}^{k} (2^rel_i - 1) / log2(i + 1)
43fn dcg(scores: &[f64]) -> f64 {
44    scores
45        .iter()
46        .enumerate()
47        .map(|(i, &rel)| {
48            let gain = (2.0f64).powf(rel) - 1.0;
49            let discount = ((i + 2) as f64).max(f64::MIN_POSITIVE).log2(); // log2(i+1+1) since i is 0-indexed
50            gain / discount
51        })
52        .sum()
53}