Skip to main content

docling_rag/eval/
metrics.rs

1//! Retrieval quality metrics: recall@k, MRR, and nDCG@k.
2//!
3//! Relevance is judged by substring match: a retrieved chunk is relevant to a
4//! query if its text contains one of the query's expected substrings
5//! (case-insensitive). This keeps eval datasets portable — they need no knowledge
6//! of chunk ids, which change with every chunking config.
7
8use crate::model::Scored;
9
10/// Per-query scores.
11#[derive(Debug, Clone, Copy, Default)]
12pub struct QueryMetrics {
13    /// Fraction of the expected substrings matched by some top-`k` chunk.
14    pub recall: f32,
15    /// Reciprocal rank of the first relevant chunk (0 if none).
16    pub mrr: f32,
17    /// Normalized discounted cumulative gain over the top `k`.
18    pub ndcg: f32,
19}
20
21/// Whether `text` satisfies any expected substring (case-insensitive).
22pub fn is_relevant(text: &str, expected: &[String]) -> bool {
23    let hay = text.to_ascii_lowercase();
24    expected
25        .iter()
26        .any(|e| hay.contains(&e.to_ascii_lowercase()))
27}
28
29/// Compute recall@k / MRR / nDCG@k for one query's ranked results.
30pub fn evaluate(hits: &[Scored], expected: &[String], k: usize) -> QueryMetrics {
31    if expected.is_empty() {
32        return QueryMetrics::default();
33    }
34    let top = &hits[..hits.len().min(k)];
35
36    // Recall: how many distinct expected substrings were surfaced.
37    let matched = expected
38        .iter()
39        .filter(|e| {
40            let el = e.to_ascii_lowercase();
41            top.iter()
42                .any(|h| h.chunk.text.to_ascii_lowercase().contains(&el))
43        })
44        .count();
45    let recall = matched as f32 / expected.len() as f32;
46
47    // MRR: reciprocal rank of the first relevant hit.
48    let mut mrr = 0.0;
49    for (i, h) in top.iter().enumerate() {
50        if is_relevant(&h.chunk.text, expected) {
51            mrr = 1.0 / (i as f32 + 1.0);
52            break;
53        }
54    }
55
56    // nDCG with binary gains; ideal ranking places all found relevants first.
57    let mut dcg = 0.0;
58    let mut relevant_found = 0usize;
59    for (i, h) in top.iter().enumerate() {
60        if is_relevant(&h.chunk.text, expected) {
61            dcg += 1.0 / ((i as f32 + 2.0).log2());
62            relevant_found += 1;
63        }
64    }
65    let idcg: f32 = (0..relevant_found)
66        .map(|i| 1.0 / ((i as f32 + 2.0).log2()))
67        .sum();
68    let ndcg = if idcg > 0.0 { dcg / idcg } else { 0.0 };
69
70    QueryMetrics { recall, mrr, ndcg }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use crate::model::Chunk;
77
78    fn hit(text: &str) -> Scored {
79        Scored::new(Chunk::new("d", 0, text, 0), 1.0)
80    }
81
82    #[test]
83    fn perfect_ranking_scores_one() {
84        let hits = vec![hit("the answer is chunking"), hit("unrelated")];
85        let m = evaluate(&hits, &["chunking".into()], 5);
86        assert!((m.recall - 1.0).abs() < 1e-6);
87        assert!((m.mrr - 1.0).abs() < 1e-6);
88        assert!((m.ndcg - 1.0).abs() < 1e-6);
89    }
90
91    #[test]
92    fn relevant_lower_down_lowers_mrr() {
93        let hits = vec![hit("noise"), hit("noise"), hit("real chunking answer")];
94        let m = evaluate(&hits, &["chunking".into()], 5);
95        assert!((m.mrr - (1.0 / 3.0)).abs() < 1e-6);
96        assert!((m.recall - 1.0).abs() < 1e-6); // still found within k
97    }
98
99    #[test]
100    fn no_match_scores_zero() {
101        let hits = vec![hit("a"), hit("b")];
102        let m = evaluate(&hits, &["missing".into()], 5);
103        assert_eq!(m.recall, 0.0);
104        assert_eq!(m.mrr, 0.0);
105        assert_eq!(m.ndcg, 0.0);
106    }
107}