66 lines
3.0 KiB
Rust
66 lines
3.0 KiB
Rust
|
use anyhow::Result;
|
||
|
|
use searchi_semantic::{Device, GemmaEmbedder};
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
#[ignore = "Requires explicitly provisioned CPU Google tokenizer and worker"]
|
||
|
|
fn normalized_code_tokenization_preserves_original_source_windows() -> Result<()> {
|
||
|
|
let models = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set model directory");
|
||
|
|
let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set CPU worker");
|
||
|
|
let mut embedder = GemmaEmbedder::load_with_device(models, worker, 1, Device::Cpu)?;
|
||
|
|
// Genuine frozen TypeScript fixture that does not roundtrip through the
|
||
|
|
// tokenizer. Do not replace source bytes with normalized decoded tokens.
|
||
|
|
let code = "export function process(rows: {category: string; amount: number}[]) {\n const totals = new Map<string, number>();\n for (const row of rows) totals.set(row.category, (totals.get(row.category) ?? 0) + row.amount);\n return totals;\n}\n";
|
||
|
|
for text in [
|
||
|
|
code.to_owned(),
|
||
|
|
format!("{code}// Grüße 日本語 🦀\n").repeat(160),
|
||
|
|
] {
|
||
|
|
let windows = embedder.split_windows(&text, "c.ts")?;
|
||
|
|
let mut covered = 0;
|
||
|
|
for window in &windows {
|
||
|
|
assert!(window.start_byte <= covered);
|
||
|
|
assert!(window.end_byte > covered);
|
||
|
|
assert_eq!(window.text, text[window.start_byte..window.end_byte]);
|
||
|
|
covered = window.end_byte;
|
||
|
|
}
|
||
|
|
assert_eq!(covered, text.len());
|
||
|
|
assert!(windows.len() <= 256);
|
||
|
|
if text.len() > 1000 {
|
||
|
|
assert!(windows.len() > 1);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
#[ignore = "Requires explicitly provisioned CPU Google model and rebuilt code-capable worker"]
|
||
|
|
fn actual_code_task_differs_from_search_and_debug_preserves_vectors() -> Result<()> {
|
||
|
|
let models = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set model directory");
|
||
|
|
let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set CPU worker");
|
||
|
|
let mut embedder = GemmaEmbedder::load_with_device(models, worker, 1, Device::Cpu)?;
|
||
|
|
assert!(embedder.info.code_query_supported);
|
||
|
|
assert_eq!(embedder.info.device, "cpu");
|
||
|
|
let query = "parse a JSON response and retry a failed network request";
|
||
|
|
let ordinary = embedder.embed_query(query)?;
|
||
|
|
let code = embedder.embed_code_query(query)?;
|
||
|
|
let cosine = |a: &[f32], b: &[f32]| {
|
||
|
|
a.iter()
|
||
|
|
.zip(b)
|
||
|
|
.map(|(a, b)| f64::from(*a) * f64::from(*b))
|
||
|
|
.sum::<f64>()
|
||
|
|
};
|
||
|
|
assert!(
|
||
|
|
cosine(&ordinary, &code) < 0.999999,
|
||
|
|
"Task-specific prompt must not silently reuse generic search embeddings"
|
||
|
|
);
|
||
|
|
let debug = embedder.embed_code_query_debug(query)?;
|
||
|
|
assert!(cosine(&code, &debug) > 0.999999);
|
||
|
|
assert!(embedder
|
||
|
|
.query_diagnostics()
|
||
|
|
.contains_key("worker_token_count"));
|
||
|
|
let ordinary_again = embedder.embed_query(query)?;
|
||
|
|
assert!(
|
||
|
|
cosine(&ordinary, &ordinary_again) > 0.999999,
|
||
|
|
"Previous code query must not contaminate a retained encoder context"
|
||
|
|
);
|
||
|
|
Ok(())
|
||
|
|
}
|