docling_rag/eval/
mod.rs

1//! Evaluation harness: sweep a matrix of `{chunk_size, overlap, retrieval mode}`
2//! over a labelled dataset and rank the configurations by retrieval quality.
3//!
4//! Each `(chunk_size, overlap)` builds a fresh in-memory index (chunk → embed →
5//! store) from the dataset's Markdown documents; then every retrieval mode is run
6//! against every query and scored with [`metrics`]. Runs fully offline with the
7//! hashing embedder and no LLM (the LLM-backed modes are included only when a
8//! [`ChatModel`] is supplied).
9
10pub mod metrics;
11
12use crate::chunk::Chunker;
13use crate::config::ChunkUnit;
14use crate::embed::Embedder;
15use crate::llm::ChatModel;
16use crate::model::{Chunk, Document, RetrievalMode};
17use crate::retrieve::Retriever;
18use crate::store::memory::MemoryStore;
19use crate::store::VectorStore;
20use crate::Result;
21use serde::{Deserialize, Serialize};
22use std::sync::Arc;
23use std::time::Instant;
24
25/// A dataset document, supplied as already-converted Markdown for determinism.
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct EvalDoc {
28    pub name: String,
29    pub markdown: String,
30}
31
32/// A labelled query. A retrieved chunk is relevant if it contains any `relevant`
33/// substring (case-insensitive).
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct QueryCase {
36    pub query: String,
37    pub relevant: Vec<String>,
38}
39
40/// A question as loaded from an external questions file. Accepts this crate's
41/// `{query, relevant}` shape, the QA-benchmark `{text, kind}` shape, and the
42/// output of the `answers` subcommand (`{question, answer, …}` — extra fields
43/// are ignored, so `answers --json` output round-trips as a questions file).
44/// Only entries with non-empty `relevant` labels can score retrieval.
45#[derive(Debug, Clone, Serialize, Deserialize)]
46pub struct Question {
47    /// The question text (`text` and `question` accepted as aliases).
48    #[serde(alias = "text", alias = "question")]
49    pub query: String,
50    /// Expected answer kind (`boolean`, `number`, `name`, …) — informational.
51    #[serde(default, skip_serializing_if = "Option::is_none")]
52    pub kind: Option<String>,
53    /// Ground-truth substrings for retrieval scoring (may be empty).
54    #[serde(default)]
55    pub relevant: Vec<String>,
56}
57
58/// Load a questions file (a JSON array of [`Question`]s in either shape).
59pub fn load_questions(path: &std::path::Path) -> Result<Vec<Question>> {
60    let raw = std::fs::read_to_string(path)?;
61    Ok(serde_json::from_str(&raw)?)
62}
63
64/// Build eval documents from a directory tree of Markdown files — typically the
65/// `RAG_DOCUMENTS_OUTPUT` mirror produced by `ingest`. Non-`.md` files are
66/// skipped; `name` is the path relative to `dir`.
67pub fn documents_from_md_dir(dir: &std::path::Path) -> Result<Vec<EvalDoc>> {
68    fn walk(dir: &std::path::Path, root: &std::path::Path, out: &mut Vec<EvalDoc>) {
69        let Ok(entries) = std::fs::read_dir(dir) else {
70            return;
71        };
72        let mut paths: Vec<_> = entries.filter_map(|e| e.ok().map(|e| e.path())).collect();
73        paths.sort();
74        for path in paths {
75            if path.is_dir() {
76                walk(&path, root, out);
77            } else if path.extension().is_some_and(|e| e == "md") {
78                if let Ok(markdown) = std::fs::read_to_string(&path) {
79                    out.push(EvalDoc {
80                        name: path
81                            .strip_prefix(root)
82                            .unwrap_or(&path)
83                            .to_string_lossy()
84                            .into_owned(),
85                        markdown,
86                    });
87                }
88            }
89        }
90    }
91    let mut docs = Vec::new();
92    walk(dir, dir, &mut docs);
93    if docs.is_empty() {
94        return Err(crate::RagError::config(format!(
95            "no .md files found under {}",
96            dir.display()
97        )));
98    }
99    Ok(docs)
100}
101
102/// A full evaluation dataset.
103#[derive(Debug, Clone, Serialize, Deserialize)]
104pub struct EvalDataset {
105    pub documents: Vec<EvalDoc>,
106    pub queries: Vec<QueryCase>,
107}
108
109/// One cell of the sweep: a chunking config paired with a retrieval mode.
110#[derive(Debug, Clone, Serialize)]
111pub struct EvalRow {
112    pub chunk_size: usize,
113    pub overlap: f32,
114    pub mode: String,
115    pub embedder: String,
116    pub recall: f32,
117    pub mrr: f32,
118    pub ndcg: f32,
119    pub avg_latency_ms: f64,
120    pub queries: usize,
121}
122
123/// The ranked results of a sweep.
124#[derive(Debug, Clone, Serialize)]
125pub struct EvalReport {
126    pub rows: Vec<EvalRow>,
127}
128
129/// Runs sweeps against a dataset using a fixed embedder (+ optional LLM).
130pub struct Harness {
131    embedder: Arc<dyn Embedder>,
132    chat: Option<Arc<dyn ChatModel>>,
133}
134
135impl Harness {
136    /// Build a harness with the embedder used for every config in the sweep.
137    pub fn new(embedder: Arc<dyn Embedder>, chat: Option<Arc<dyn ChatModel>>) -> Self {
138        Harness { embedder, chat }
139    }
140
141    /// Run the sweep. `modes` that need an LLM are skipped when no [`ChatModel`]
142    /// was provided.
143    pub async fn run(
144        &self,
145        dataset: &EvalDataset,
146        chunk_configs: &[(usize, f32)],
147        modes: &[RetrievalMode],
148        top_k: usize,
149    ) -> Result<EvalReport> {
150        let mut rows = Vec::new();
151        for &(size, overlap) in chunk_configs {
152            let store = self.build_index(dataset, size, overlap).await?;
153            let retriever = Retriever::new(store, self.embedder.clone(), self.chat.clone());
154            for &mode in modes {
155                if mode.needs_llm() && self.chat.is_none() {
156                    continue;
157                }
158                rows.push(
159                    self.score_mode(&retriever, dataset, mode, size, overlap, top_k)
160                        .await?,
161                );
162            }
163        }
164        // Rank best-first by nDCG, then recall, then MRR.
165        rows.sort_by(|a, b| {
166            b.ndcg
167                .partial_cmp(&a.ndcg)
168                .unwrap_or(std::cmp::Ordering::Equal)
169                .then(
170                    b.recall
171                        .partial_cmp(&a.recall)
172                        .unwrap_or(std::cmp::Ordering::Equal),
173                )
174                .then(
175                    b.mrr
176                        .partial_cmp(&a.mrr)
177                        .unwrap_or(std::cmp::Ordering::Equal),
178                )
179        });
180        Ok(EvalReport { rows })
181    }
182
183    /// Chunk + embed + store every dataset document under one chunking config.
184    async fn build_index(
185        &self,
186        dataset: &EvalDataset,
187        size: usize,
188        overlap: f32,
189    ) -> Result<Arc<dyn VectorStore>> {
190        let store: Arc<dyn VectorStore> = Arc::new(MemoryStore::new());
191        let chunker = Chunker {
192            size,
193            overlap,
194            unit: ChunkUnit::Word,
195        };
196        for d in &dataset.documents {
197            let doc = Document::new(format!("eval://{}", d.name), &d.name, "");
198            store.upsert_document(&doc).await?;
199            let mut chunks: Vec<Chunk> = chunker.chunk(&doc.id, &d.markdown);
200            if chunks.is_empty() {
201                continue;
202            }
203            let texts: Vec<String> = chunks.iter().map(|c| c.text.clone()).collect();
204            let embeddings = self.embedder.embed(&texts).await?;
205            for (c, e) in chunks.iter_mut().zip(embeddings) {
206                c.embedding = Some(e);
207            }
208            store.insert_chunks(&chunks).await?;
209        }
210        Ok(store)
211    }
212
213    /// Average metrics over all queries for one mode.
214    async fn score_mode(
215        &self,
216        retriever: &Retriever,
217        dataset: &EvalDataset,
218        mode: RetrievalMode,
219        size: usize,
220        overlap: f32,
221        top_k: usize,
222    ) -> Result<EvalRow> {
223        let (mut recall, mut mrr, mut ndcg, mut latency) = (0.0, 0.0, 0.0, 0.0);
224        let n = dataset.queries.len().max(1) as f32;
225        for q in &dataset.queries {
226            let start = Instant::now();
227            let hits = retriever.retrieve(mode, &q.query, top_k).await?;
228            latency += start.elapsed().as_secs_f64() * 1000.0;
229            let m = metrics::evaluate(&hits, &q.relevant, top_k);
230            recall += m.recall;
231            mrr += m.mrr;
232            ndcg += m.ndcg;
233        }
234        Ok(EvalRow {
235            chunk_size: size,
236            overlap,
237            mode: mode.to_string(),
238            embedder: self.embedder.id().to_string(),
239            recall: recall / n,
240            mrr: mrr / n,
241            ndcg: ndcg / n,
242            avg_latency_ms: latency / n as f64,
243            queries: dataset.queries.len(),
244        })
245    }
246}
247
248impl EvalReport {
249    /// Render the report as a Markdown table, best config first.
250    pub fn to_markdown(&self) -> String {
251        let mut out = String::from(
252            "| chunk | overlap | mode | embedder | recall | MRR | nDCG | ms/query |\n\
253             |------:|--------:|------|----------|-------:|----:|-----:|---------:|\n",
254        );
255        for r in &self.rows {
256            out.push_str(&format!(
257                "| {} | {:.2} | {} | {} | {:.3} | {:.3} | {:.3} | {:.2} |\n",
258                r.chunk_size,
259                r.overlap,
260                r.mode,
261                r.embedder,
262                r.recall,
263                r.mrr,
264                r.ndcg,
265                r.avg_latency_ms
266            ));
267        }
268        out
269    }
270
271    /// Render the report as pretty JSON.
272    pub fn to_json(&self) -> String {
273        serde_json::to_string_pretty(self).unwrap_or_else(|_| "{}".to_string())
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280    use crate::embed::HashEmbedder;
281
282    fn dataset() -> EvalDataset {
283        EvalDataset {
284            documents: vec![
285                EvalDoc {
286                    name: "chunking".into(),
287                    markdown: "# Chunking\n\nDocuments are split into overlapping chunks of a \
288                               configurable size before embedding for semantic search."
289                        .into(),
290                },
291                EvalDoc {
292                    name: "cooking".into(),
293                    markdown: "# Cooking\n\nBlend banana yogurt and honey to make a smoothie."
294                        .into(),
295                },
296            ],
297            queries: vec![QueryCase {
298                query: "how are documents split for embedding".into(),
299                relevant: vec!["overlapping chunks".into()],
300            }],
301        }
302    }
303
304    #[test]
305    fn loads_questions_in_both_shapes() {
306        let dir = std::env::temp_dir().join(format!("rag-q-{}", crate::model::new_id()));
307        std::fs::create_dir_all(&dir).unwrap();
308        let path = dir.join("questions.json");
309        std::fs::write(
310            &path,
311            r#"[
312                {"text": "Did X mention mergers?", "kind": "boolean"},
313                {"query": "chunk size default", "relevant": ["300 words"]},
314                {"question": "answers output row?", "answer": "yes", "sources": 10, "ms": 12.5, "mode": "hybrid"}
315            ]"#,
316        )
317        .unwrap();
318        let qs = load_questions(&path).unwrap();
319        assert_eq!(qs.len(), 3);
320        assert_eq!(qs[0].query, "Did X mention mergers?");
321        assert_eq!(qs[0].kind.as_deref(), Some("boolean"));
322        assert!(qs[0].relevant.is_empty());
323        assert_eq!(qs[1].relevant, vec!["300 words"]);
324        // The `answers --json` output round-trips: `question` alias, extra
325        // fields ignored.
326        assert_eq!(qs[2].query, "answers output row?");
327        std::fs::remove_dir_all(&dir).ok();
328    }
329
330    #[test]
331    fn builds_documents_from_md_dir() {
332        let dir = std::env::temp_dir().join(format!("rag-mdd-{}", crate::model::new_id()));
333        std::fs::create_dir_all(dir.join("sub")).unwrap();
334        std::fs::write(dir.join("a.pdf.md"), "# A\n\nalpha").unwrap();
335        std::fs::write(dir.join("sub/b.md"), "# B\n\nbeta").unwrap();
336        std::fs::write(dir.join("ignore.txt"), "not markdown").unwrap();
337        let docs = documents_from_md_dir(&dir).unwrap();
338        assert_eq!(docs.len(), 2);
339        assert!(docs.iter().any(|d| d.name == "a.pdf.md"));
340        assert!(docs
341            .iter()
342            .any(|d| d.name.ends_with("b.md") && d.markdown.contains("beta")));
343        std::fs::remove_dir_all(&dir).ok();
344    }
345
346    #[tokio::test]
347    async fn sweep_produces_ranked_rows() {
348        let harness = Harness::new(Arc::new(HashEmbedder::new(512)), None);
349        let report = harness
350            .run(
351                &dataset(),
352                &[(20, 0.0), (40, 0.1)],
353                &RetrievalMode::OFFLINE,
354                3,
355            )
356            .await
357            .unwrap();
358        // 2 chunk configs x 3 offline modes = 6 rows.
359        assert_eq!(report.rows.len(), 6);
360        // At least one config retrieves the relevant chunk.
361        assert!(report.rows.iter().any(|r| r.recall > 0.0));
362        // Markdown/JSON render without panicking.
363        assert!(report.to_markdown().contains("nDCG"));
364        assert!(report.to_json().contains("recall"));
365    }
366}
docling_rag/eval/mod.rs

docling_rag/eval/
mod.rs
