120 lines
4.5 KiB
Rust
120 lines
4.5 KiB
Rust
use anyhow::Result;
|
|
use searchi_semantic::{Device, GemmaEmbedder};
|
|
|
|
fn cosine(a: &[f32], b: &[f32]) -> f64 {
|
|
a.iter()
|
|
.zip(b)
|
|
.map(|(a, b)| f64::from(*a) * f64::from(*b))
|
|
.sum()
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Requires verified Google model and current CPU-capable native worker"]
|
|
fn query_context_preserves_real_vectors_ranking_and_token_boundary() -> 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 full = GemmaEmbedder::load_with_device(&model, &worker, 1, Device::Cpu)?;
|
|
let mut query = GemmaEmbedder::load_query_with_device(&model, &worker, 1, Device::Cpu)?;
|
|
assert_eq!(full.info.context_tokens, 2048);
|
|
assert_eq!(query.info.context_tokens, 512);
|
|
assert_eq!(query.info.role, "query");
|
|
// This describes the binary's media capability, not this session's role.
|
|
assert!(query.info.media_supported);
|
|
let documents = full.embed_documents(
|
|
&[
|
|
"The supplier agreement may be terminated with sixty days written notice.",
|
|
"Business train tickets and hotel stays are reimbursed after submitting receipts.",
|
|
"Die Gartenplanung sieht ein Hochbeet und eine Regentonne vor.",
|
|
],
|
|
"none",
|
|
)?;
|
|
for (text, 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 before = full.embed_query(text)?;
|
|
let after = query.embed_query(text)?;
|
|
let parity = cosine(&before, &after);
|
|
assert!(parity >= 0.99999, "Context changed vector: {parity}");
|
|
let rank = |vector: &[f32]| {
|
|
documents
|
|
.iter()
|
|
.enumerate()
|
|
.max_by(|a, b| cosine(a.1, vector).total_cmp(&cosine(b.1, vector)))
|
|
.unwrap()
|
|
.0
|
|
};
|
|
assert_eq!(rank(&before), expected);
|
|
assert_eq!(rank(&after), expected);
|
|
println!("Query role parity {parity:.8}; expected document {expected}");
|
|
}
|
|
for text in [
|
|
"retry failed requests with exponential backoff",
|
|
"Zugriffstoken bei Ablauf erneuern",
|
|
"متى تغلق المكتبة",
|
|
"ইলিশ মাছ রান্না",
|
|
] {
|
|
assert!(
|
|
cosine(
|
|
&full.embed_code_query(text)?,
|
|
&query.embed_code_query(text)?
|
|
) >= 0.99999
|
|
);
|
|
}
|
|
// Actual token counts determine the exact boundary, rather than assuming
|
|
// words equal tokens. Each repeated plain token contributes one token.
|
|
let mut repetitions = 490usize;
|
|
let mut boundary = String::new();
|
|
for _ in 0..8 {
|
|
let text = " cat".repeat(repetitions);
|
|
full.embed_query_debug(&text)?;
|
|
let tokens = full.query_diagnostics()["worker_token_count"] as usize;
|
|
if tokens == 512 {
|
|
boundary = text;
|
|
break;
|
|
}
|
|
assert!(tokens < 512, "Unexpected fixture token count {tokens}");
|
|
repetitions += 512 - tokens;
|
|
}
|
|
assert!(
|
|
!boundary.is_empty(),
|
|
"Could not construct exact512token query"
|
|
);
|
|
let before = full.embed_query(&boundary)?;
|
|
let after = query.embed_query_debug(&boundary)?;
|
|
assert_eq!(query.query_diagnostics()["worker_token_count"], 512);
|
|
assert!(cosine(&before, &after) >= 0.99999);
|
|
assert!(query.embed_query(&(boundary + " cat")).is_err());
|
|
assert!(query
|
|
.embed_query("supplier cancellation")?
|
|
.iter()
|
|
.all(|v| v.is_finite()));
|
|
assert!(query
|
|
.embed_documents(&["document"], "none")
|
|
.unwrap_err()
|
|
.to_string()
|
|
.contains("Foreground query worker"));
|
|
assert!(query.split_windows("document", "none").is_err());
|
|
assert!(query
|
|
.embed_audio_pcm(&[0.0; 16])
|
|
.unwrap_err()
|
|
.to_string()
|
|
.contains("Foreground query worker"));
|
|
assert!(query
|
|
.embed_images(&[std::path::Path::new("missing.png")])
|
|
.unwrap_err()
|
|
.to_string()
|
|
.contains("Foreground query worker"));
|
|
// Full2048-token indexing survives, including long document extraction.
|
|
assert!(
|
|
full.split_windows(&"ordinary maintenance log ".repeat(1000), "log")?
|
|
.len()
|
|
> 1
|
|
);
|
|
Ok(())
|
|
}
|