Merge feat/bm25-tokenizer: optional keyword stemming, measured and left off by default

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>

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

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

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<String> {
+    tokenize_with(text, TokenFilter::Plain)
+}
+
+/// Split `text` into scoring tokens under `filter`.
+pub fn tokenize_with(text: &str, filter: TokenFilter) -> Vec<String> {
     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<String> = (0..6).map(|_| "same text".to_string()).collect();

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<bm25::BM25Index>,
+    /// 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();

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<MemoryEntry> = 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() {