diff --git a/BENCHMARKS.md b/BENCHMARKS.md index 39b6dc0..806fe14 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -625,6 +625,34 @@ labelled data to tune against. Here there is, so the weighted sum is kept as the default. `Fusion::Rrf` remains available for callers whose stages are more evenly matched. +### Keyword tokenizer — stemming, full haystack, n=500 + +The keyword stage lowercases and splits on non-alphanumerics, with no stemming, +so "training" and "trains" are unrelated terms. `TokenFilter::Stemmed` strips +common English inflections (plurals, `-ing`/`-ed`, with consonant un-doubling) +from documents and queries alike. Turn-level: + +| Mode | Hit@1 | Hit@5 | Hit@10 | MRR | session Hit@1 | +|---|---|---|---|---|---| +| BM25 only | **53.8%** | 75.0% | 81.6% | 0.6320 | 86.2% | +| BM25 only, stemmed | 52.0% | 77.8% | 84.0% | 0.6320 | 88.0% | +| Hybrid 0.4/0.6 | 51.6% | **81.4%** | 87.8% | **0.6430** | 91.0% | +| Hybrid 0.4/0.6, stemmed | 50.2% | **81.4%** | **88.2%** | 0.6394 | **91.4%** | + +**Stemming is a trade, not a win, and the default stays off.** It reliably buys +depth and costs the top rank: on BM25 alone, +2.8pp Hit@5 and +2.4pp Hit@10 for +−1.8pp Hit@1, with MRR unchanged to four decimal places — the gains deeper down +exactly offset the loss at rank 1. That is what conflation does: merging +"train"/"training"/"trains" surfaces documents an exact-match query would never +reach, and also lets a near-miss outrank the exact hit. + +On the configuration that actually ships (hybrid 0.4/0.6) the trade is +narrower still — Hit@5 identical, Hit@10 +0.4pp, Hit@1 −1.4pp, MRR −0.004 — +because the vector stage already supplies much of the recall stemming would +add. There is no case here for changing the default; `TokenFilter::Stemmed` +is available via `HDF5Memory::set_token_filter` for callers who want Hit@5/@10 +over rank-1 precision. + ### Weight sweep — full haystack, n=500 `0.7/0.3` was a documented default, never a searched one. Sweeping diff --git a/CHANGELOG.md b/CHANGELOG.md index 5076f3d..2a91cdd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,12 @@ ## Unreleased ### Retrieval quality +- `clawhdf5-agent`: optional keyword stemming — `bm25::TokenFilter::Stemmed` + and `HDF5Memory::set_token_filter`, so "training" and "trains" match. **Off + by default**, on measurement rather than principle: over the full LongMemEval + haystack it buys depth and costs the top rank (BM25 alone: Hit@5 +2.8pp, + Hit@10 +2.4pp, Hit@1 −1.8pp, MRR unchanged), and on the shipping hybrid + configuration the trade is narrower still. See `BENCHMARKS.md`. - `clawhdf5-agent`: **`QueryExpander::expand` panicked on ordinary non-ASCII input** — `"İ AI"` was enough. It searched a lowercased copy of the query and then sliced the *original* with those offsets, which only works while diff --git a/crates/clawhdf5-agent/src/bm25.rs b/crates/clawhdf5-agent/src/bm25.rs index fb6d603..621a9f7 100644 --- a/crates/clawhdf5-agent/src/bm25.rs +++ b/crates/clawhdf5-agent/src/bm25.rs @@ -59,11 +59,18 @@ pub struct BM25Index { k1: f32, /// BM25 b parameter. b: f32, + /// Applied to every document and query token, so the two always agree. + filter: TokenFilter, } impl BM25Index { /// Build a BM25 index from a set of documents, excluding tombstoned entries. pub fn build(documents: &[String], tombstones: &[u8]) -> Self { + Self::build_with(documents, tombstones, TokenFilter::default()) + } + + /// [`BM25Index::build`] with the token filter chosen explicitly. + pub fn build_with(documents: &[String], tombstones: &[u8], filter: TokenFilter) -> Self { let mut index = Self { inverted: HashMap::new(), doc_lengths: vec![0; documents.len()], @@ -72,6 +79,7 @@ impl BM25Index { num_docs: 0, k1: DEFAULT_K1, b: DEFAULT_B, + filter, }; index.index_documents(documents, tombstones); index @@ -120,7 +128,7 @@ impl BM25Index { // add/remove, and costs one `ln` per query term. let mut acc = vec![0.0f32; self.doc_lengths.len()]; let mut matched = false; - for token in tokenize(query) { + for token in tokenize_with(query, self.filter) { let Some(postings) = self.inverted.get(token.as_str()) else { continue; }; @@ -146,6 +154,11 @@ impl BM25Index { .collect() } + /// The token filter this index was built with. + pub fn token_filter(&self) -> TokenFilter { + self.filter + } + /// Number of document slots (live or not) the index covers. Ids are /// positions in the document list it mirrors. pub fn len(&self) -> usize { @@ -168,7 +181,7 @@ impl BM25Index { } debug_assert_eq!(self.doc_lengths[doc_id], 0, "slot {doc_id} is occupied"); - let tokens = tokenize(text); + let tokens = tokenize_with(text, self.filter); let mut term_freqs: HashMap<&str, u32> = HashMap::new(); for token in &tokens { *term_freqs.entry(token).or_insert(0) += 1; @@ -201,7 +214,7 @@ impl BM25Index { /// Remove document `doc_id`, whose indexed text was `text`. The text is /// needed to find its postings; pass exactly what was added. pub fn remove_document(&mut self, doc_id: usize, text: &str) { - let tokens = tokenize(text); + let tokens = tokenize_with(text, self.filter); let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new(); for token in &tokens { if !seen.insert(token) { @@ -252,7 +265,7 @@ impl BM25Index { continue; } - let tokens = tokenize(doc); + let tokens = tokenize_with(doc, self.filter); let doc_len = tokens.len() as u32; self.doc_lengths[i] = doc_len; total_length += doc_len as u64; @@ -285,11 +298,86 @@ impl BM25Index { /// Tokenize a string: lowercase, split on non-alphanumeric characters, /// filter empty tokens. +/// What [`tokenize_with`] does to each token after splitting. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum TokenFilter { + /// Lowercase and split only — the original behaviour. + #[default] + Plain, + /// Also strip common English inflections, so "running" and "runs" match + /// "run". Conservative on purpose: only plural and past/continuous verb + /// endings, and only on tokens long enough that stripping leaves a real + /// stem. A stemmer earns its keep by conflating *related* words; an + /// aggressive one also conflates unrelated ones ("universe"/"university"), + /// which costs precision. + Stemmed, +} + +/// Strip common English inflections from an already-lowercased token. +/// +/// Applied identically to documents and queries, so the pair only has to agree +/// with itself — the stem need not be a real word. +fn stem(token: &str) -> &str { + // Below this, stripping does more harm than good ("bed" -> "b"). + const MIN_STEM: usize = 4; + let strip = |suffix: &str, min_len: usize| -> Option<&str> { + let stem = token.strip_suffix(suffix)?; + (stem.len() >= min_len).then_some(stem) + }; + + // Plurals first: "studies" -> "studi", "classes" -> "class", "cats" -> "cat". + // "ies" keeps its "i" so the result meets "-ied" ("studied" -> "studi"). + if let Some(stem) = strip("ies", 2) { + return &token[..stem.len() + 1]; + } + for suffix in ["sses", "shes", "ches", "xes", "zes"] { + if let Some(stem) = strip(suffix, MIN_STEM - 1) { + // Keep the sibilant: "classes" -> "class", not "clas". + return &token[..stem.len() + 2]; + } + } + // Verb endings before the bare plural, so "raced" doesn't become "raced". + if let Some(stem) = strip("ing", MIN_STEM - 1).or_else(|| strip("ed", MIN_STEM - 1)) { + return undouble(stem); + } + if !token.ends_with("ss") + && !token.ends_with("us") + && !token.ends_with("is") + && let Some(stem) = strip("s", MIN_STEM - 1) + { + return stem; + } + token +} + +/// "runn" -> "run": undo the consonant doubling that "-ing"/"-ed" introduce. +fn undouble(stem: &str) -> &str { + let mut chars = stem.chars().rev(); + let (Some(last), Some(prev)) = (chars.next(), chars.next()) else { + return stem; + }; + let doubled = last == prev && !"aeiou".contains(last) && last.is_ascii_alphabetic(); + if doubled && stem.len() > 3 { + &stem[..stem.len() - 1] + } else { + stem + } +} + +#[cfg(test)] fn tokenize(text: &str) -> Vec { + tokenize_with(text, TokenFilter::Plain) +} + +/// Split `text` into scoring tokens under `filter`. +pub fn tokenize_with(text: &str, filter: TokenFilter) -> Vec { text.to_lowercase() .split(|c: char| !c.is_alphanumeric()) .filter(|s| !s.is_empty()) - .map(|s| s.to_string()) + .map(|token| match filter { + TokenFilter::Plain => token.to_string(), + TokenFilter::Stemmed => stem(token).to_string(), + }) .collect() } @@ -599,6 +687,53 @@ mod tests { } } + #[test] + fn stemming_conflates_inflections_of_the_same_word() { + let stem_of = |w: &str| tokenize_with(w, TokenFilter::Stemmed).pop().unwrap(); + // Pairs that should meet. + for (a, b) in [ + ("running", "runs"), + ("trained", "training"), + ("miles", "mile"), + ("studies", "studied"), + ("mentioned", "mentioning"), + ("classes", "class"), + ("planned", "planning"), + ] { + assert_eq!(stem_of(a), stem_of(b), "{a} / {b} should share a stem"); + } + // Pairs that must stay apart. Note which pairs are deliberately absent: + // "bed"/"bedding" and "gas"/"gassed" both collapse to one stem, which + // is what Porter does too and is right — they are related words. + for (a, b) in [ + ("universe", "university"), + ("business", "busy"), + ("this", "thing"), + ] { + assert_ne!(stem_of(a), stem_of(b), "{a} / {b} must not be conflated"); + } + // Short words and non-inflections are left alone. + for word in ["run", "bus", "is", "his", "data", "gas"] { + assert_eq!(stem_of(word), word, "{word} should be untouched"); + } + } + + #[test] + fn stemming_is_off_by_default_and_applied_consistently() { + assert_eq!(tokenize("Running miles"), ["running", "miles"]); + assert_eq!( + tokenize_with("Running miles", TokenFilter::Stemmed), + ["run", "mile"] + ); + + // A query inflected differently from the document still matches. + let docs = vec!["I ran while training for the marathon".to_string()]; + let plain = BM25Index::build_with(&docs, &[0], TokenFilter::Plain); + let stemmed = BM25Index::build_with(&docs, &[0], TokenFilter::Stemmed); + assert!(plain.search("trains", 1).is_empty()); + assert_eq!(stemmed.search("trains", 1).len(), 1); + } + #[test] fn ties_break_towards_the_lower_doc_id() { let docs: Vec = (0..6).map(|_| "same text".to_string()).collect(); diff --git a/crates/clawhdf5-agent/src/lib.rs b/crates/clawhdf5-agent/src/lib.rs index a3c5bd4..0c69042 100644 --- a/crates/clawhdf5-agent/src/lib.rs +++ b/crates/clawhdf5-agent/src/lib.rs @@ -258,6 +258,9 @@ pub struct HDF5Memory { /// every record, on every single query. Built lazily on first use; see /// [`HDF5Memory::ensure_bm25_fresh`] for how it stays in sync. bm25: Option, + /// Token filter the keyword index is built with. Changing it drops the + /// index; it is not persisted, because the index is not either. + bm25_filter: bm25::TokenFilter, /// Activation weights changed since the last checkpoint (searches boost /// the records they return). Cleared by `flush`. activations_dirty: bool, @@ -314,6 +317,7 @@ impl HDF5Memory { anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()), anomaly_alerts: Vec::new(), bm25: None, + bm25_filter: bm25::TokenFilter::default(), activations_dirty: false, read_only: false, quarantined_wal: None, @@ -481,6 +485,7 @@ impl HDF5Memory { anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()), anomaly_alerts: Vec::new(), bm25: None, + bm25_filter: bm25::TokenFilter::default(), activations_dirty: false, read_only, quarantined_wal, @@ -591,7 +596,7 @@ impl HDF5Memory { pub(crate) fn ensure_bm25_fresh(&mut self) -> &bm25::BM25Index { let n = self.cache.chunks.len(); let bm25 = match self.bm25.take() { - Some(index) if index.len() <= n => { + Some(index) if index.len() <= n && index.token_filter() == self.bm25_filter => { let mut index = index; for id in index.len()..n { if self.cache.tombstones[id] == 0 { @@ -601,11 +606,26 @@ impl HDF5Memory { index.pad_to(n); index } - _ => bm25::BM25Index::build(&self.cache.chunks, &self.cache.tombstones), + _ => bm25::BM25Index::build_with( + &self.cache.chunks, + &self.cache.tombstones, + self.bm25_filter, + ), }; self.bm25.insert(bm25) } + /// Choose how keyword-search tokens are normalised, rebuilding the index + /// on next use. [`bm25::TokenFilter::Stemmed`] matches inflections of the + /// same word at some cost in precision; measure before adopting it (see + /// `BENCHMARKS.md`). + pub fn set_token_filter(&mut self, filter: bm25::TokenFilter) { + if filter != self.bm25_filter { + self.bm25_filter = filter; + self.bm25 = None; + } + } + /// Record `id` was tombstoned; its text is still in the cache. fn bm25_on_delete(&mut self, id: usize) { if let Some(index) = self.bm25.as_mut() @@ -1955,6 +1975,34 @@ mod tests { assert_eq!(top_ids(&mut reopened, &q), expected_after); } + #[test] + fn set_token_filter_rebuilds_the_keyword_index() { + let dir = TempDir::new().unwrap(); + let mut mem = HDF5Memory::create(make_config(&dir)).unwrap(); + mem.save(make_entry( + "I was training for a marathon", + &[1.0, 0.0, 0.0, 0.0], + )) + .unwrap(); + + // Count only genuine keyword matches: `hybrid_search` also returns + // zero-score filler when fewer than k records are relevant. + let hits = |mem: &mut HDF5Memory| { + mem.hybrid_search(&[0.0, 0.0, 0.0, 0.0], "trains", 0.0, 1.0, 5) + .iter() + .filter(|r| r.score > 0.0) + .count() + }; + assert_eq!(hits(&mut mem), 0); + + mem.set_token_filter(bm25::TokenFilter::Stemmed); + assert_eq!(hits(&mut mem), 1, "index should have been rebuilt stemmed"); + + // And back, rebuilding again. + mem.set_token_filter(bm25::TokenFilter::Plain); + assert_eq!(hits(&mut mem), 0); + } + #[test] fn keyword_index_stays_in_sync_through_every_mutation() { let dir = TempDir::new().unwrap(); diff --git a/crates/clawhdf5-bench/src/bin/longmemeval_bench.rs b/crates/clawhdf5-bench/src/bin/longmemeval_bench.rs index c723627..4369804 100644 --- a/crates/clawhdf5-bench/src/bin/longmemeval_bench.rs +++ b/crates/clawhdf5-bench/src/bin/longmemeval_bench.rs @@ -55,6 +55,7 @@ use std::time::{Duration, Instant}; #[path = "longmemeval_bench/embedder.rs"] mod embedder; +use clawhdf5_agent::bm25::TokenFilter; use clawhdf5_agent::hybrid::Fusion; use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry}; use serde::Deserialize; @@ -64,9 +65,13 @@ const EMBEDDING_DIM: usize = 384; /// A mode's fusion, as one short string for the reports. fn describe(mode: Mode) -> String { - match mode.fusion { + let fusion = match mode.fusion { Fusion::Weighted { vector, keyword } => format!("vector_{vector:.1}_keyword_{keyword:.1}"), Fusion::Rrf { k } => format!("rrf_k{k:.0}"), + }; + match mode.tokens { + TokenFilter::Plain => fusion, + TokenFilter::Stemmed => format!("{fusion}_stemmed"), } } @@ -76,6 +81,8 @@ struct Mode { label: &'static str, /// How the two retrieval stages are combined into one ranking. fusion: Fusion, + /// How keyword tokens are normalised before indexing and querying. + tokens: TokenFilter, } impl Mode { @@ -83,8 +90,15 @@ impl Mode { Self { label, fusion: Fusion::Weighted { vector, keyword }, + tokens: TokenFilter::Plain, } } + + const fn stemmed(mut self, label: &'static str) -> Self { + self.label = label; + self.tokens = TokenFilter::Stemmed; + self + } } /// The only mode available without real embeddings. Passing zero vectors with @@ -106,8 +120,15 @@ const HYBRID: Mode = Mode::weighted("Hybrid (0.4 vector / 0.6 BM25, tuned)", 0.4 const RRF: Mode = Mode { label: "Hybrid (reciprocal rank fusion, k=60)", fusion: Fusion::Rrf { k: 60.0 }, + tokens: TokenFilter::Plain, }; +/// The same two configurations with stemmed keyword tokens, so the tokenizer's +/// effect is isolated from everything else. +const BM25_STEMMED: Mode = BM25_ONLY.stemmed("BM25 only, stemmed tokens"); +#[cfg(feature = "embeddings")] +const HYBRID_STEMMED: Mode = HYBRID.stemmed("Hybrid 0.4/0.6, stemmed tokens"); + /// Every 0.1 step of vector weight, keyword weight taking the remainder. /// /// Labels are leaked to `&'static str` because `Mode::label` is a `&'static @@ -291,6 +312,7 @@ fn evaluate_question( config.compact_threshold = 0.0; let mut memory = HDF5Memory::create(config).expect("failed to create HDF5Memory"); + memory.set_token_filter(mode.tokens); // Build MemoryEntry list from all haystack sessions let mut entries: Vec = Vec::new(); @@ -771,18 +793,27 @@ fn main() { if sweep { sweep_modes() } else { - vec![BM25_ONLY, VECTOR_ONLY, HYBRID, RRF] + vec![ + BM25_ONLY, + VECTOR_ONLY, + HYBRID, + RRF, + BM25_STEMMED, + HYBRID_STEMMED, + ] } } #[cfg(not(feature = "embeddings"))] { - vec![BM25_ONLY] + vec![BM25_ONLY, BM25_STEMMED] } } else { if sweep { eprintln!("warning: --sweep needs --embeddings; running BM25 only"); } - vec![BM25_ONLY] + // Stemming is a property of the keyword stage, so it can be compared + // without a model. + vec![BM25_ONLY, BM25_STEMMED] }; for (mode_idx, mode) in modes.iter().enumerate() {