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}