quant_eval/
rag.rs

1//! Local fixture evaluation for RAG-style retrieval experiments.
2
3use std::collections::HashSet;
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct RagQueryFixture {
7    pub query_id: String,
8    pub query: String,
9    pub relevant_doc_ids: Vec<String>,
10}
11
12#[derive(Debug, Clone, PartialEq)]
13pub struct RagRetrievedDoc {
14    pub doc_id: String,
15    pub score: f32,
16}
17
18#[derive(Debug, Clone, PartialEq)]
19pub struct RagEvalResult {
20    pub recall_at_k: f32,
21    pub ndcg_at_k: f32,
22    pub exact_rerank_recovery: f32,
23}
24
25pub fn evaluate_rag_fixture(
26    fixture: &RagQueryFixture,
27    retrieved: &[RagRetrievedDoc],
28    k: usize,
29) -> RagEvalResult {
30    let relevant_doc_ids: HashSet<&str> = fixture
31        .relevant_doc_ids
32        .iter()
33        .map(String::as_str)
34        .collect();
35
36    if relevant_doc_ids.is_empty() {
37        return RagEvalResult {
38            recall_at_k: 0.0,
39            ndcg_at_k: 0.0,
40            exact_rerank_recovery: 0.0,
41        };
42    }
43
44    let exact_rerank_recovery = retrieved
45        .first()
46        .filter(|doc| relevant_doc_ids.contains(doc.doc_id.as_str()))
47        .map(|_| 1.0)
48        .unwrap_or(0.0);
49
50    let mut seen_doc_ids = HashSet::new();
51    let mut relevant_found = HashSet::new();
52    let mut dcg = 0.0f32;
53
54    for (index, doc) in retrieved.iter().take(k).enumerate() {
55        if !seen_doc_ids.insert(doc.doc_id.as_str()) {
56            continue;
57        }
58
59        if relevant_doc_ids.contains(doc.doc_id.as_str()) {
60            relevant_found.insert(doc.doc_id.as_str());
61            dcg += discounted_gain(index + 1);
62        }
63    }
64
65    let ideal_len = relevant_doc_ids.len().min(k);
66    let idcg = (1..=ideal_len).map(discounted_gain).sum::<f32>();
67    let ndcg_at_k = if idcg > 0.0 { dcg / idcg } else { 0.0 };
68
69    RagEvalResult {
70        recall_at_k: relevant_found.len() as f32 / relevant_doc_ids.len() as f32,
71        ndcg_at_k,
72        exact_rerank_recovery,
73    }
74}
75
76fn discounted_gain(rank: usize) -> f32 {
77    if rank == 1 {
78        1.0
79    } else {
80        1.0 / ((rank + 1) as f32).log2()
81    }
82}
quant_eval/rag.rs

quant_eval/
rag.rs
