notescraft/crates/knowledge/tests/eval.rs

//! Evaluation on a synthetic vault: recall@6 and MRR of the hybrid ranker, and proof
//! that plain BM25 does worse on questions that use other words than the note.

#![allow(clippy::cast_precision_loss)]

use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};

use notescraft_knowledge::{Index, MapSignals, Options, Vocab, retrieve};

const RECALL_AT: usize = 6;
const MRR_DEPTH: usize = 10;
const MIN_RECALL: f64 = 0.80;
const MIN_NOTES: usize = 30;
const MIN_CASES: usize = 25;

struct Case {
    question: String,
    path: String,
    heading: String,
    gap: bool,
}

struct Vault {
    index: Index,
    vocab: Vocab,
    signals: MapSignals,
    files: HashMap<String, String>,
}

fn root() -> PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/eval")
}

fn walk(dir: &Path, base: &Path, out: &mut Vec<(String, String)>) {
    let mut entries: Vec<_> = fs::read_dir(dir).unwrap().map(|e| e.unwrap().path()).collect();
    entries.sort();
    for p in entries {
        if p.is_dir() {
            walk(&p, base, out);
        } else if p.extension().is_some_and(|e| e == "md") {
            let rel = p.strip_prefix(base).unwrap().to_string_lossy().replace('\\', "/");
            out.push((rel, fs::read_to_string(&p).unwrap()));
        }
    }
}

/// Tags come from a `tags: [a, b]` frontmatter line, links from `[[stem]]`.
fn load() -> Vault {
    let base = root().join("vault");
    let mut notes = Vec::new();
    walk(&base, &base, &mut notes);
    let by_stem: HashMap<String, String> = notes.iter().map(|(p, _)| (Path::new(p).file_stem().unwrap().to_string_lossy().to_string(), p.clone())).collect();
    let (mut index, mut vocab, mut signals, mut files) = (Index::default(), Vocab::default(), MapSignals::default(), HashMap::new());
    for (path, text) in notes {
        index.update(&path, &text).unwrap();
        vocab.update(&path, &text).unwrap();
        let tags = text
            .lines()
            .find_map(|l| l.strip_prefix("tags:"))
            .map(|t| t.trim().trim_matches(['[', ']']).split(',').map(|s| s.trim().to_string()).collect())
            .unwrap_or_default();
        signals.tags.insert(path.clone(), tags);
        let links = text.split("[[").skip(1).filter_map(|r| r.split("]]").next()).filter_map(|s| by_stem.get(s)).cloned().collect();
        signals.links.insert(path.clone(), links);
        files.insert(path, text);
    }
    Vault { index, vocab, signals, files }
}

fn cases() -> Vec<Case> {
    let text = fs::read_to_string(root().join("cases.tsv")).unwrap();
    text.lines()
        .filter(|l| !l.is_empty() && !l.starts_with('#'))
        .map(|l| {
            let c: Vec<&str> = l.split('\t').collect();
            Case { question: c[0].into(), path: c[1].into(), heading: c[2].into(), gap: c.get(3) == Some(&"gap") }
        })
        .collect()
}

/// 1-based rank of the expected passage in the top `MRR_DEPTH`, if present.
fn rank_of(v: &Vault, c: &Case, options: &Options) -> Option<usize> {
    let options = Options { limit: MRR_DEPTH, ..*options };
    let r = retrieve(&v.index, &v.vocab, &v.signals, |p| v.files.get(p).cloned(), &c.question, &options).unwrap();
    r.passages.iter().position(|p| p.path == c.path && p.heading == c.heading).map(|i| i + 1)
}

struct Score {
    recall: f64,
    mrr: f64,
}

fn score(v: &Vault, cases: &[&Case], options: &Options, label: &str) -> Score {
    let (mut hit, mut rr) = (0.0, 0.0);
    for c in cases {
        let rank = rank_of(v, c, options);
        println!("  [{label}] {:>8} {}", rank.map_or("miss".to_string(), |r| format!("rank {r}")), c.question);
        hit += f64::from(u8::from(rank.is_some_and(|r| r <= RECALL_AT)));
        rr += rank.map_or(0.0, |r| 1.0 / r as f64);
    }
    let n = cases.len() as f64;
    Score { recall: hit / n, mrr: rr / n }
}

const BASELINE: Options = Options { limit: RECALL_AT, follow_links: false, use_tags: false, learn_vocab: false, prefer_recent: false };

#[test]
fn eval_set_is_big_enough() {
    let v = load();
    assert!(v.index.note_count() >= MIN_NOTES, "{} notes", v.index.note_count());
    assert!(cases().len() >= MIN_CASES);
    assert!(v.signals.links.values().any(|l| !l.is_empty()) && v.signals.tags.values().any(|t| !t.is_empty()));
}

#[test]
fn hybrid_reaches_recall_at_six() {
    let v = load();
    let all = cases();
    let refs: Vec<&Case> = all.iter().collect();
    let hybrid = score(&v, &refs, &Options::default(), "hybrid");
    let base = score(&v, &refs, &BASELINE, "bm25");
    println!("hybrid: recall@{RECALL_AT} = {:.3}, MRR = {:.3}", hybrid.recall, hybrid.mrr);
    println!("bm25  : recall@{RECALL_AT} = {:.3}, MRR = {:.3}", base.recall, base.mrr);
    assert!(hybrid.recall >= MIN_RECALL, "recall@{RECALL_AT} {:.3} < {MIN_RECALL}", hybrid.recall);
    assert!(hybrid.recall >= base.recall);
}

#[test]
fn plain_bm25_is_worse_on_vocabulary_gaps() {
    let v = load();
    let all = cases();
    let gaps: Vec<&Case> = all.iter().filter(|c| c.gap).collect();
    assert!(gaps.len() >= 6);
    let hybrid = score(&v, &gaps, &Options::default(), "hybrid gap");
    let base = score(&v, &gaps, &BASELINE, "bm25 gap");
    println!("gap cases hybrid: recall {:.3} MRR {:.3}; bm25: recall {:.3} MRR {:.3}", hybrid.recall, hybrid.mrr, base.recall, base.mrr);
    assert!(hybrid.recall > base.recall && hybrid.mrr > base.mrr);
}