searchi/crates/searchi-service/examples/evaluate.rs

//! Evaluate the immutable seed corpus with the actual configured Google model.
//! Usage: cargo run -p searchi-service --example evaluate -- MODEL_DIR WORKER OUTPUT_JSON
use anyhow::{ensure, Context, Result};
use searchi_core::SearchIndex;
use searchi_service::{HybridSearch, SearchOptions, SearchResult};
use serde_json::{json, Value};
use std::{fs, path::PathBuf, time::Instant};

fn expected_rank(results: &[SearchResult], filename: &str) -> Option<usize> {
    results
        .iter()
        .position(|result| result.filename == filename)
        .map(|rank| rank + 1)
}

fn main() -> Result<()> {
    let args: Vec<_> = std::env::args_os().skip(1).collect();
    ensure!(
        args.len() == 3,
        "Supply MODEL_DIR WORKER OUTPUT_JSON; this command never downloads assets"
    );
    ensure!(
        std::env::var("SEARCHI_FORCE_CPU").is_ok_and(|value| value == "1"),
        "Windows VM evaluation requires SEARCHI_FORCE_CPU=1"
    );
    let fixture: Value =
        serde_json::from_str(include_str!("../../../fixtures/relevance/fixture.json"))?;
    let temp = tempfile::tempdir()?;
    let root = temp.path().join("documents");
    fs::create_dir(&root)?;
    let documents = fixture["documents"]
        .as_array()
        .context("documents missing")?;
    for document in documents {
        let filename = document["filename"].as_str().context("filename missing")?;
        ensure!(
            PathBuf::from(filename).file_name() == Some(std::ffi::OsStr::new(filename)),
            "Fixture filename must be a basename"
        );
        fs::write(
            root.join(filename),
            document["text"].as_str().context("text missing")?,
        )?;
    }
    let db = temp.path().join("index.db");
    let mut lexical = SearchIndex::open(&db)?;
    lexical.index_folder(&root)?;
    let start = Instant::now();
    let mut hybrid = HybridSearch::open(&db, &args[0], &args[1], 1)?;
    let load_ms = start.elapsed().as_secs_f64() * 1000.0;
    let start = Instant::now();
    let indexed = hybrid.index_embeddings()?;
    let indexing_ms = start.elapsed().as_secs_f64() * 1000.0;
    drop(hybrid);
    let start = Instant::now();
    let mut hybrid = HybridSearch::open(&db, &args[0], &args[1], 1)?;
    let startup_init_ms = start.elapsed().as_secs_f64() * 1000.0;
    let start = Instant::now();
    hybrid.warmup()?;
    let background_prewarm_ms = start.elapsed().as_secs_f64() * 1000.0;
    let start = Instant::now();
    let first = hybrid.search(
        "How can we end our relationship with the supplier?",
        20,
        &SearchOptions::default(),
    )?;
    let first_foreground_ms = start.elapsed().as_secs_f64() * 1000.0;
    let first_foreground_top1 = first.first().map(|hit| hit.filename.as_str());
    let mut rankings = Vec::new();
    let mut durations = Vec::new();
    for query in fixture["queries"].as_array().context("queries missing")? {
        let text = query["text"].as_str().context("query text missing")?;
        let expected = query["expected"].as_str().context("expected id missing")?;
        let filename = documents
            .iter()
            .find(|doc| doc["id"] == expected)
            .context("Expected document missing")?["filename"]
            .as_str()
            .unwrap();
        let lexical_hits = lexical.search(text, 20)?;
        let semantic_hits = hybrid.search_semantic(text, 20, &SearchOptions::default())?;
        let start = Instant::now();
        let hybrid_hits = hybrid.search(text, 20, &SearchOptions::default())?;
        durations.push(start.elapsed().as_secs_f64() * 1000.0);
        rankings.push(json!({"query":text,"expected":expected,"language":query["language"],
            "lexical_rank":expected_rank(&lexical_hits,filename),
            "semantic_rank":expected_rank(&semantic_hits,filename),
            "hybrid_rank":expected_rank(&hybrid_hits,filename),
            "hybrid_top3":hybrid_hits.iter().take(3).map(|hit| json!({"filename":hit.filename,"match_kind":hit.match_kind,"source":hit.source})).collect::<Vec<_>>()}));
    }
    durations.sort_by(f64::total_cmp);
    let metrics = |lane: &str| {
        json!({
            "top1_accuracy":rankings.iter().filter(|row| row[lane]==1).count() as f64/rankings.len() as f64,
            "recall_at3":rankings.iter().filter(|row| row[lane].as_u64().is_some_and(|rank| rank<=3)).count() as f64/rankings.len() as f64,
            "recall_at10":rankings.iter().filter(|row| row[lane].as_u64().is_some_and(|rank| rank<=10)).count() as f64/rankings.len() as f64,
        })
    };
    let report = json!({"model":"google/embeddinggemma-2","device":"cpu","threads":1,
        "documents":documents.len(),"queries":rankings.len(),"load_ms":load_ms,"indexing_ms":indexing_ms,
        "startup_init_ms":startup_init_ms,"background_prewarm_ms":background_prewarm_ms,"first_foreground_ms":first_foreground_ms,"first_foreground_top1":first_foreground_top1,
        "hybrid_p50_ms":durations[durations.len()/2],"hybrid_p95_ms":durations[(durations.len()*95/100).min(durations.len()-1)],
        "embedding_report":indexed,"lexical":metrics("lexical_rank"),"semantic":metrics("semantic_rank"),"hybrid":metrics("hybrid_rank"),
        "rankings":rankings,"limitation":"Synthetic seed corpus only; no production quality or large-corpus latency certification"});
    fs::write(&args[2], serde_json::to_string_pretty(&report)?)?;
    println!(
        "{}",
        json!({"lexical":report["lexical"],"semantic":report["semantic"],"hybrid":report["hybrid"],"p95_ms":report["hybrid_p95_ms"]})
    );
    Ok(())
}