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
+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();