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:
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user