feat(agent): optional keyword stemming, measured and left off by default

The keyword stage had no stemming, so "training" and "trains" were unrelated
terms. bm25::TokenFilter::Stemmed strips common English inflections (plurals,
-ing/-ed, with consonant un-doubling) from documents and queries alike;
BM25Index::build_with and HDF5Memory::set_token_filter select it, and the
index records which filter built it so a stale one is rebuilt rather than
mixed.

Measured over the full LongMemEval haystack (500 questions, real MiniLM
embeddings) rather than adopted on principle — and it is a trade, not a win:

  BM25 only         Hit@1 53.8%  Hit@5 75.0%  Hit@10 81.6%  MRR 0.6320
  BM25 stemmed      Hit@1 52.0%  Hit@5 77.8%  Hit@10 84.0%  MRR 0.6320
  Hybrid 0.4/0.6    Hit@1 51.6%  Hit@5 81.4%  Hit@10 87.8%  MRR 0.6430
  Hybrid stemmed    Hit@1 50.2%  Hit@5 81.4%  Hit@10 88.2%  MRR 0.6394

Conflation buys depth and costs the top rank: on BM25 alone MRR is unchanged
to four decimal places, the deeper gains exactly offsetting the rank-1 loss.
On the shipping hybrid configuration the vector stage already supplies most of
that recall, so the trade is narrower and slightly negative. Default stays
Plain; Stemmed is there for callers who want Hit@5/@10 over rank-1 precision.

The stemmer is deliberately conservative — it only strips inflections, and
only when the stem stays long enough to be meaningful, since an aggressive one
also conflates unrelated words. Tests pin both the pairs that must meet and
the pairs that must not.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
osobh
2026-09-19 18:20:19 -07:00
co-authored by Claude Opus 5
parent 55ed87d2e8
commit 84a39ef3c5
5 changed files with 259 additions and 11 deletions
+28
View File
@@ -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
+6
View File
@@ -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
+140 -5
View File
@@ -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();
+50 -2
View File
@@ -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();
@@ -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() {