searchi/crates/searchi-semantic/tests/gemma_cpu.rs

use anyhow::Result;
use searchi_semantic::{Device, GemmaEmbedder, GEMMA_DIMENSIONS};

#[test]
#[ignore = "Requires explicitly provisioned Google model and CPU worker"]
fn google_cpu_real_embeddings_and_verified_tail_windows() -> Result<()> {
    let model = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set SEARCHI_GEMMA_MODEL_DIR");
    let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set SEARCHI_GEMMA_WORKER");
    let mut embedder = GemmaEmbedder::load_with_device(model, worker, 1, Device::Cpu)?;
    assert_eq!(embedder.info.device, "cpu");
    assert_eq!(embedder.info.dimensions, GEMMA_DIMENSIONS);
    let reference: serde_json::Value =
        serde_json::from_str(include_str!("fixtures/digup-text-reference.json"))?;
    let long_reference=(1..=60).map(|i|format!("Note {i}: the review moved to room {} because the projector in the main hall failed again, and the team agreed to send the slides the night before.",i%7+2)).collect::<Vec<_>>().join(" ");
    let mut made = Vec::new();
    for (name, query) in [
        ("query", "payment declined error"),
        ("query-bn", "ইলিশ মাছ রান্না"),
        ("query-ar", "متى تغلق المكتبة"),
    ] {
        made.push((name, embedder.embed_query(query)?));
    }
    made.push(("document",embedder.embed_documents(&["Flight to Tokyo Haneda on November 3. Then the train to Kamakura to see the Great Buddha, and back for dinner in Shibuya."],"Tokyo itinerary")?.remove(0)));
    made.push((
        "document-long",
        embedder.embed_documents(&[&long_reference], "")?.remove(0),
    ));
    for (name, vector) in made {
        let expected = reference["vectors"][name].as_array().unwrap();
        let cosine: f64 = vector
            .iter()
            .zip(expected)
            .map(|(a, b)| f64::from(*a) * b.as_f64().unwrap())
            .sum();
        println!("DigUp text reference {name}: cosine {cosine:.8}");
        assert!(
            cosine >= 0.999,
            "Google CPU reference parity failed for {name}: {cosine}"
        );
    }
    let documents = [
        "The supplier agreement may be ended with sixty days written notice.",
        "Business train tickets and hotel stays can be reimbursed by submitting receipts.",
        "Die Gartenplanung sieht ein Hochbeet für Gemüse und eine Regentonne vor.",
    ];
    let vectors = embedder.embed_documents(&documents, "none")?;
    for (query, expected) in [
        ("Wie beenden wir die Zusammenarbeit mit dem Lieferanten?", 0),
        ("How do I claim my accommodation and railway expenses?", 1),
        (
            "Where is the plan for the vegetable bed and rainwater barrel?",
            2,
        ),
    ] {
        let query = embedder.embed_query(query)?;
        let scores: Vec<f32> = vectors
            .iter()
            .map(|vector| vector.iter().zip(&query).map(|(a, b)| a * b).sum())
            .collect();
        let best = scores
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .unwrap()
            .0;
        assert_eq!(
            best, expected,
            "Semantic cross-language ranking failed: {scores:?}"
        );
    }
    let long = format!("{}\nHotel accommodation on business trips is reimbursed after submitting the overnight stay receipt.",
        "The routine maintenance log records ordinary uptime checks and cable inspections. ".repeat(350));
    let windows = embedder.split_windows(&long, "maintenance-log.txt")?;
    assert!(windows.len() > 1);
    assert_eq!(windows.last().unwrap().end_byte, long.len());
    assert!(windows
        .last()
        .unwrap()
        .text
        .contains("overnight stay receipt"));
    let texts: Vec<_> = windows.iter().map(|window| window.text.as_str()).collect();
    let vectors = embedder.embed_documents(&texts, "maintenance-log.txt")?;
    let query = embedder.embed_query(
        "Where does the policy say that lodging costs during work travel can be claimed?",
    )?;
    let scores: Vec<f32> = vectors
        .iter()
        .map(|vector| vector.iter().zip(&query).map(|(a, b)| a * b).sum())
        .collect();
    let best = scores
        .iter()
        .enumerate()
        .max_by(|a, b| a.1.total_cmp(b.1))
        .unwrap()
        .0;
    assert_eq!(
        best,
        windows.len() - 1,
        "Tail retrieval did not find final window: {scores:?}"
    );
    Ok(())
}