diff --git a/app/src-tauri/src/lib.rs b/app/src-tauri/src/lib.rs index 801f0ff..5de41ab 100644 --- a/app/src-tauri/src/lib.rs +++ b/app/src-tauri/src/lib.rs @@ -4424,20 +4424,20 @@ fn apply_embedding_model(app: &AppHandle, id: Option) -> Result<(), Stri return Err(format!("embedding model not downloaded: {model_id}")); } - load_embedder(app, model_id)?; - let _ = fs::write(&state.embed_model_path, model_id); - - // Re-embed existing notes with the new model so semantic search covers them too (and so a - // model switch never mixes vectors of different dimensions). Best-effort: the model is - // already active for new saves even if this pass fails. - if let Err(e) = state + // Switch the embedder and re-embed existing notes in one locked, atomic step: either the + // new model becomes active with a freshly-rebuilt index, or — on any failure — the + // previous model and its index are left untouched and the error surfaces to the UI. This + // avoids both a half-applied switch (new embedder vs old-dimension vectors) and a + // silently-wiped index. + let embedder = resolve_embedder(&state, model_id)?; + state .library .lock() .map_err(|_| "library lock poisoned".to_owned())? - .reindex_all() - { - eprintln!("reindex after embedding-model change failed: {e}"); - } + .set_embedder_and_reindex(Some(embedder)) + .map_err(|e| format!("indexing notes with {model_id} failed: {e}"))?; + + let _ = fs::write(&state.embed_model_path, model_id); } None => { state diff --git a/crates/wisp-embed/src/lib.rs b/crates/wisp-embed/src/lib.rs index 8664415..869fd45 100644 --- a/crates/wisp-embed/src/lib.rs +++ b/crates/wisp-embed/src/lib.rs @@ -267,9 +267,24 @@ fn fetch_atomic(url: &str, dest: &Path) -> Result<()> { let resp = ureq::get(url) .call() .map_err(|e| LibraryError::Embed(format!("download {url}: {e}")))?; + + // Hold the transfer to the declared length. Without this, a connection dropped mid-body on a + // length-less (close-delimited) response reads as a clean EOF, and the truncated file gets + // renamed into place as "complete" — then it fails to load forever, since the present (short) + // file makes the download skip it on every retry. + let expected: Option = resp.header("Content-Length").and_then(|h| h.parse().ok()); + let mut reader = resp.into_reader(); let mut out = std::fs::File::create(&part).map_err(io_err)?; - std::io::copy(&mut reader, &mut out).map_err(io_err)?; + let written = std::io::copy(&mut reader, &mut out).map_err(io_err)?; + + if let Some(total) = expected { + if written != total { + return Err(LibraryError::Embed(format!( + "download {url}: truncated ({written} of {total} bytes)" + ))); + } + } Ok(()) })(); diff --git a/crates/wisp-embed/src/ort_embed.rs b/crates/wisp-embed/src/ort_embed.rs index 7721ac0..67fe853 100644 --- a/crates/wisp-embed/src/ort_embed.rs +++ b/crates/wisp-embed/src/ort_embed.rs @@ -130,6 +130,14 @@ impl OrtEmbedder { let seq = encodings[0].len(); let batch = encodings.len(); + // A zero-length encoding (a pathological input/tokenizer) would otherwise yield a 0-width + // tensor and an all-zero "embedding"; fail loudly instead. + if seq == 0 { + return Err(LibraryError::Embed( + "tokenizer produced an empty encoding".to_owned(), + )); + } + let mut ids = Vec::with_capacity(batch * seq); let mut mask = Vec::with_capacity(batch * seq); let mut types = Vec::with_capacity(batch * seq); @@ -279,11 +287,18 @@ impl Embedder for OrtEmbedder { } fn embed_passages(&self, texts: &[&str]) -> Result>> { - let prefixed: Vec = texts + // Embed one text at a time (batch size 1). Batching pads the shorter rows, and a quantized + // model's per-batch scales make a padded row drift from the same text embedded alone — so a + // stored passage would not match a (single-text) query embedding, and two notes would index + // the same sentence differently depending on their neighbours. One at a time keeps every + // vector identical regardless of the batch it was embedded in. + texts .iter() - .map(|t| format!("{}{t}", self.passage_prefix)) - .collect(); - self.run(&prefixed) + .map(|t| { + let prefixed = format!("{}{t}", self.passage_prefix); + Ok(self.run(&[prefixed])?.pop().unwrap_or_default()) + }) + .collect() } fn embed_query(&self, text: &str) -> Result> { @@ -747,6 +762,52 @@ mod tests { ); } + // A passage must embed identically whether alone or alongside others. `embed_passages` embeds one + // text at a time precisely so a quantized model's per-batch scales can't make a stored vector + // depend on its neighbours (which would also stop it matching a single-text query). Uses the + // cached Qwen3 model. `cargo test -- --ignored`. + #[test] + #[ignore] + fn qwen3_passage_embedding_is_independent_of_neighbours() { + let dir = std::env::temp_dir().join("wisp-ort-qwen3-0_6b"); + let repo = "onnx-community/Qwen3-Embedding-0.6B-ONNX"; + for f in [ + "onnx/model_quantized.onnx", + "tokenizer.json", + "tokenizer_config.json", + "config.json", + "special_tokens_map.json", + ] { + download_file(repo, f, &dir.join(f)); + } + + let recipe = Recipe { + onnx_file: "onnx/model_quantized.onnx", + pooling: Pooling::LastToken, + passage_prefix: "", + query_prefix: "", + normalize: true, + dim: 1024, + max_length: 512, + }; + let emb = OrtEmbedder::load(&dir, &recipe).unwrap(); + + let short = "budget"; + let long = "the annual operating budget was approved by the board after a long and detailed \ + discussion covering every department and each contingency line item in the plan"; + + let alone = emb.embed_passages(&[short]).unwrap().pop().unwrap(); + // Batching the short row with a much longer one forces heavy right-padding on the short row. + let batched = emb.embed_passages(&[short, long]).unwrap(); + + let cos = |a: &[f32], b: &[f32]| a.iter().zip(b).map(|(x, y)| x * y).sum::(); + let sim = cos(&alone, &batched[0]); + assert!( + sim > 0.999, + "a passage embedded differently next to a neighbour (cos {sim}); embed_passages must embed one at a time" + ); + } + // BGE-M3 (XLM-RoBERTa encoder, CLS pooling, 1024-dim, 100+ languages) via the int8 ONNX (~600 MB). #[test] #[ignore] diff --git a/crates/wisp-library/src/store.rs b/crates/wisp-library/src/store.rs index 50b08ed..ba94c56 100644 --- a/crates/wisp-library/src/store.rs +++ b/crates/wisp-library/src/store.rs @@ -113,6 +113,26 @@ impl Library { self.embedder = embedder; } + /// Atomically switches the active embedder and re-embeds the whole corpus with it. On any failure + /// the library is left exactly as it was — previous embedder, previous chunks — so a model switch + /// never wipes the index or leaves a new embedder pointed at stale-dimension vectors. Use this + /// (not [`set_embedder`] + [`reindex_all`]) whenever the model changes. + pub fn set_embedder_and_reindex( + &mut self, + embedder: Option>, + ) -> Result { + let previous = self.embedder.take(); + self.embedder = embedder; + + match self.reindex_all() { + Ok(n) => Ok(n), + Err(e) => { + self.embedder = previous; + Err(e) + } + } + } + fn init(conn: Connection) -> Result { conn.execute_batch("PRAGMA foreign_keys = ON;")?; let mut lib = Self { @@ -199,14 +219,13 @@ impl Library { /// Returns how many notes produced chunks; clears all chunks and returns 0 with no embedder set. /// A personal library is small, so it re-embeds the whole corpus in a single pass. pub fn reindex_all(&mut self) -> Result { - self.conn.execute("DELETE FROM chunk", [])?; - if self.embedder.is_none() { + self.conn.execute("DELETE FROM chunk", [])?; return Ok(0); } // Embed every note before opening the write transaction — inference touches no DB and may be - // slow. + // slow. A note that fails to embed returns here with the existing index still intact. let mut indexed: Vec<(String, Vec)> = Vec::new(); for id in self.note_ids()? { let texts = self.segment_texts(&id)?; @@ -219,7 +238,11 @@ impl Library { } } + // Clear the old chunks and insert the new ones in one transaction, so a failure (or a crash) + // never leaves the corpus de-indexed. The DELETE used to run + commit before embedding, so a + // mid-embed error wiped the whole index for good. let tx = self.conn.transaction()?; + tx.execute("DELETE FROM chunk", [])?; for (id, chunks) in &indexed { insert_chunks(&tx, id, chunks)?; } @@ -707,6 +730,28 @@ mod tests { } } + /// A test embedder whose passage embedding always fails — used to drive a reindex error and prove + /// the corpus + previous embedder survive it. + struct FailingEmbedder; + + impl Embedder for FailingEmbedder { + fn dim(&self) -> usize { + 64 + } + + fn embed_passages(&self, _texts: &[&str]) -> Result>> { + Err(crate::LibraryError::Embed( + "simulated inference failure".to_owned(), + )) + } + + fn embed_query(&self, _text: &str) -> Result> { + Err(crate::LibraryError::Embed( + "simulated inference failure".to_owned(), + )) + } + } + /// An in-memory library with the deterministic test embedder configured. fn lib_with_embedder() -> Library { let mut lib = Library::open_in_memory().unwrap(); @@ -1179,6 +1224,45 @@ mod tests { ); } + #[test] + fn a_failed_model_switch_keeps_the_old_index_and_embedder() { + let mut lib = lib_with_embedder(); + lib.save_note( + "m1", + &meta("M"), + 0, + &[seg(1, 0, 100, "budget review", AudioSourceKind::Microphone)], + ) + .unwrap(); + assert_eq!(lib.reindex_all().unwrap(), 1); + assert_eq!( + lib.search_semantic("budget", 10).unwrap()[0].meeting_id, + "m1" + ); + + // Switching to a model that fails to embed must NOT wipe the index or leave the failing model + // active: the call errors, and the previous (working) embedder + chunks are restored intact. + assert!(lib + .set_embedder_and_reindex(Some(Box::new(FailingEmbedder))) + .is_err()); + assert_eq!( + lib.search_semantic("budget", 10).unwrap()[0].meeting_id, + "m1", + "a failed model switch must leave the old index + embedder intact" + ); + + // A switch to a working model still succeeds and rebuilds the index. + assert_eq!( + lib.set_embedder_and_reindex(Some(Box::new(HashingEmbedder { dim: 128 }))) + .unwrap(), + 1 + ); + assert_eq!( + lib.search_semantic("budget", 10).unwrap()[0].meeting_id, + "m1" + ); + } + #[test] fn reindex_without_an_embedder_clears_all_chunks() { let mut lib = lib_with_embedder();