diff --git a/crates/clawhdf5-bench/src/bin/longmemeval_bench.rs b/crates/clawhdf5-bench/src/bin/longmemeval_bench.rs index 4369804..b508bab 100644 --- a/crates/clawhdf5-bench/src/bin/longmemeval_bench.rs +++ b/crates/clawhdf5-bench/src/bin/longmemeval_bench.rs @@ -57,7 +57,8 @@ mod embedder; use clawhdf5_agent::bm25::TokenFilter; use clawhdf5_agent::hybrid::Fusion; -use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry}; +use clawhdf5_agent::reranker::{ReRankConfig, RerankInput, rerank}; +use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchResult}; use serde::Deserialize; use tempfile::TempDir; @@ -69,9 +70,17 @@ fn describe(mode: Mode) -> String { Fusion::Weighted { vector, keyword } => format!("vector_{vector:.1}_keyword_{keyword:.1}"), Fusion::Rrf { k } => format!("rrf_k{k:.0}"), }; - match mode.tokens { + let tokens = match mode.tokens { TokenFilter::Plain => fusion, TokenFilter::Stemmed => format!("{fusion}_stemmed"), + }; + match mode.rerank { + None => tokens, + Some(cfg) if cfg.relevance_weight == 0.0 => format!("{tokens}_rerank_metadata"), + Some(cfg) => format!( + "{tokens}_rerank_blended_hl{:.0}d", + cfg.temporal_half_life_secs / 86_400.0 + ), } } @@ -83,6 +92,9 @@ struct Mode { fusion: Fusion, /// How keyword tokens are normalised before indexing and querying. tokens: TokenFilter, + /// Re-rank the retrieved candidates with recency and friends, relative to + /// the question's own date. + rerank: Option, } impl Mode { @@ -91,9 +103,16 @@ impl Mode { label, fusion: Fusion::Weighted { vector, keyword }, tokens: TokenFilter::Plain, + rerank: None, } } + fn reranked(mut self, label: &'static str, rerank: ReRankConfig) -> Self { + self.label = label; + self.rerank = Some(rerank); + self + } + const fn stemmed(mut self, label: &'static str) -> Self { self.label = label; self.tokens = TokenFilter::Stemmed; @@ -121,11 +140,61 @@ const RRF: Mode = Mode { label: "Hybrid (reciprocal rank fusion, k=60)", fusion: Fusion::Rrf { k: 60.0 }, tokens: TokenFilter::Plain, + rerank: None, }; /// 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"); + +/// Re-ranking as it behaved before `relevance` was an input: the combined +/// score was recency + authority + activation only, so the retriever's own +/// ordering was discarded. +#[cfg(feature = "embeddings")] +fn hybrid_rerank_metadata_only() -> Mode { + HYBRID.reranked( + "Hybrid + rerank (metadata only, pre-fix)", + ReRankConfig { + relevance_weight: 0.0, + ..ReRankConfig::default() + }, + ) +} + +/// Re-ranking as it behaves now: relevance leads, recency nudges. +#[cfg(feature = "embeddings")] +fn hybrid_rerank_blended() -> Mode { + HYBRID.reranked( + "Hybrid + rerank (relevance + recency)", + ReRankConfig::default(), + ) +} + +/// The same blend at several half-lives. Decay is `2^(-age / half_life)`, so a +/// half-life far shorter than the gaps between memories sends every score to +/// zero and the signal vanishes; far longer and everything scores ~1 and it +/// vanishes the other way. The right value tracks how far apart the memories +/// actually are. +#[cfg(feature = "embeddings")] +fn hybrid_rerank_half_lives() -> Vec { + [ + ("1 day", 86_400.0), + ("7 days", 7.0 * 86_400.0), + ("30 days", 30.0 * 86_400.0), + ("90 days", 90.0 * 86_400.0), + ] + .into_iter() + .map(|(label, half_life)| { + HYBRID.reranked( + Box::leak(format!("Hybrid + rerank, half-life {label}").into_boxed_str()), + ReRankConfig { + temporal_half_life_secs: half_life, + ..ReRankConfig::default() + }, + ) + }) + .collect() +} #[cfg(feature = "embeddings")] const HYBRID_STEMMED: Mode = HYBRID.stemmed("Hybrid 0.4/0.6, stemmed tokens"); @@ -217,6 +286,37 @@ struct Question { haystack_session_ids: Vec, haystack_sessions: Vec>, answer_session_ids: Vec, + /// One timestamp per haystack session, e.g. "2023/05/25 (Thu) 20:21". + #[serde(default)] + haystack_dates: Vec, +} + +/// Seconds since the epoch for a LongMemEval session date, which looks like +/// `2023/05/25 (Thu) 20:21`. Sessions are stored in chronological order, so a +/// date that cannot be parsed falls back to its position — order is preserved +/// even if the interval is not. +fn session_time(date: &str, position: usize) -> f64 { + let stamp = |y: i64, mo: i64, d: i64, h: i64, mi: i64| -> f64 { + // Days since 1970-01-01 via the civil-from-days algorithm. + let (y, mo) = if mo <= 2 { (y - 1, mo + 12) } else { (y, mo) }; + let era = y.div_euclid(400); + let yoe = y - era * 400; + let doy = (153 * (mo - 3) + 2) / 5 + d - 1; + let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy; + let days = era * 146_097 + doe - 719_468; + (days * 86_400 + h * 3_600 + mi * 60) as f64 + }; + let parse = || -> Option { + let (ymd, rest) = date.split_once(' ')?; + let mut ymd = ymd.split('/'); + let y = ymd.next()?.parse().ok()?; + let mo = ymd.next()?.parse().ok()?; + let d = ymd.next()?.parse().ok()?; + let hm = rest.rsplit(' ').next()?; + let (h, mi) = hm.split_once(':')?; + Some(stamp(y, mo, d, h.parse().ok()?, mi.parse().ok()?)) + }; + parse().unwrap_or(1_000_000.0 + position as f64 * 86_400.0) } // --------------------------------------------------------------------------- @@ -235,11 +335,21 @@ struct Metrics { rr_turn: f64, abstention_correct: u32, abstention_total: u32, + /// Questions where the newest gold session outranked the older ones, out + /// of those with more than one gold session and at least one retrieved. + newest_gold_first: u32, + newest_gold_total: u32, latency_ns: Vec, count: u32, } impl Metrics { + /// `None` when no question in this bucket had multiple gold sessions. + fn newest_gold_first_pct(&self) -> Option { + (self.newest_gold_total > 0) + .then(|| self.newest_gold_first as f64 / self.newest_gold_total as f64 * 100.0) + } + fn hit1_session_pct(&self) -> f64 { self.hit1_session as f64 / self.count.max(1) as f64 * 100.0 } @@ -297,6 +407,16 @@ struct EvalResult { hit5_turn: bool, hit10_turn: bool, rr_turn: Option, + /// For a question whose evidence spans several dated sessions (a + /// `knowledge-update`, where an earlier fact is superseded by a later + /// one): did the *newest* gold session outrank every older gold session + /// that was returned? `None` when the question has one gold session, or + /// when none were retrieved, so there is nothing to discriminate. + /// + /// Plain recall cannot see this. LongMemEval labels *both* the stale and + /// the updated session as gold, so returning either counts as a hit — yet + /// only one of them answers the question correctly. + newest_gold_first: Option, latency: Duration, } @@ -317,15 +437,21 @@ fn evaluate_question( // Build MemoryEntry list from all haystack sessions let mut entries: Vec = Vec::new(); let mut turn_has_answer: Vec = Vec::new(); - let mut ts = 1_000_000.0f64; - for (sess_idx, session) in q.haystack_sessions.iter().enumerate() { let sess_id = q .haystack_session_ids .get(sess_idx) .map(String::as_str) .unwrap_or("unknown"); - for turn in session { + // Real session dates, not a synthetic counter: anything that decays + // with age needs true intervals, not just the right order. + let session_start = q + .haystack_dates + .get(sess_idx) + .map_or(sess_idx as f64 * 86_400.0, |d| session_time(d, sess_idx)); + for (turn_idx, turn) in session.iter().enumerate() { + // Spread a session's turns over the minutes following its start. + let ts = session_start + turn_idx as f64 * 60.0; entries.push(MemoryEntry { chunk: turn.content.clone(), embedding: embedding_for(embeddings, &turn.content), @@ -339,7 +465,6 @@ fn evaluate_question( }, }); turn_has_answer.push(turn.has_answer); - ts += 1.0; } } @@ -356,11 +481,87 @@ fn evaluate_question( // Set of session IDs that contain the answer let answer_sess_set: HashSet<&str> = q.answer_session_ids.iter().map(String::as_str).collect(); + // When each gold session was recorded, so "newest" is by date rather than + // by position (the two agree in this dataset, but the metric should not + // depend on that). + let gold_times: HashMap<&str, f64> = q + .haystack_session_ids + .iter() + .enumerate() + .filter(|(_, sid)| answer_sess_set.contains(sid.as_str())) + .map(|(i, sid)| { + let t = q + .haystack_dates + .get(i) + .map_or(i as f64 * 86_400.0, |d| session_time(d, i)); + (sid.as_str(), t) + }) + .collect(); + let query_emb = embedding_for(embeddings, &q.question); let t0 = Instant::now(); - let results = memory.hybrid_search_with(&query_emb, &q.question, mode.fusion, top_k); + // Re-ranking only reorders; it needs a candidate pool larger than `top_k` + // to have anything to promote. + let pool = if mode.rerank.is_some() { + top_k * 4 + } else { + top_k + }; + let mut results = memory.hybrid_search_with(&query_emb, &q.question, mode.fusion, pool); + if let Some(config) = mode.rerank { + // "Now" is the moment the question was asked, so decay measures how + // stale each memory was at that point. + let now = session_time(&q.question_date, q.haystack_sessions.len()); + let inputs: Vec = results + .iter() + .map(|r| RerankInput { + index: r.index, + timestamp: r.timestamp, + source_channel: r.source_channel.clone(), + raw_activation: r.activation, + relevance: r.score, + }) + .collect(); + let order: Vec = rerank(&inputs, &config, now) + .into_iter() + .map(|r| r.index) + .collect(); + let by_index: HashMap = + results.into_iter().map(|r| (r.index, r)).collect(); + results = order + .into_iter() + .filter_map(|i| by_index.get(&i).cloned()) + .collect(); + } + results.truncate(top_k); let latency = t0.elapsed(); + // Rank of the best-placed result from each gold session. + let mut first_rank: HashMap<&str, usize> = HashMap::new(); + for (rank, result) in results.iter().enumerate() { + let sid = memory.cache.session_ids[result.index].as_str(); + if let Some((gold_sid, _)) = gold_times.get_key_value(sid) { + first_rank.entry(gold_sid).or_insert(rank); + } + } + let newest_gold_first = if gold_times.len() < 2 || first_rank.is_empty() { + None + } else { + // The newest gold session must be retrieved, and no older gold session + // may outrank it. + let newest = gold_times + .iter() + .max_by(|a, b| a.1.total_cmp(b.1)) + .map(|(sid, _)| *sid) + .expect("at least two gold sessions"); + Some(match first_rank.get(newest) { + Some(&newest_rank) => first_rank + .iter() + .all(|(sid, &rank)| *sid == newest || rank > newest_rank), + None => false, + }) + }; + // Session-level recall let mut hit1_session = false; let mut hit5_session = false; @@ -415,6 +616,7 @@ fn evaluate_question( hit5_turn, hit10_turn, rr_turn, + newest_gold_first, latency, } } @@ -566,6 +768,24 @@ fn print_report( ); println!(); + if let Some(pct) = overall.newest_gold_first_pct() { + println!( + "## Recency Discrimination (n={})", + overall.newest_gold_total + ); + println!( + " Newest gold session ranked first: {}/{} ({pct:.1}%)", + overall.newest_gold_first, overall.newest_gold_total + ); + println!( + " Questions whose evidence spans several dated sessions — a fact and\n \ + its later correction. Both sessions are labelled gold, so recall\n \ + scores either as a hit; this asks whether the *current* one came\n \ + first. A retriever with no sense of time scores near chance." + ); + println!(); + } + if overall.abstention_total > 0 { println!("## Abstention Accuracy"); println!( @@ -679,6 +899,14 @@ fn print_report( } else { println!(" \"abstention_accuracy\": null,"); } + match overall.newest_gold_first_pct() { + Some(pct) => println!( + " \"newest_gold_first\": {:.4}, \"newest_gold_n\": {},", + pct / 100.0, + overall.newest_gold_total + ), + None => println!(" \"newest_gold_first\": null,"), + } println!(" \"latency_us\": {{"); println!( " \"avg\": {:.1}, \"p50\": {:.1}, \"p95\": {:.1}, \"p99\": {:.1}", @@ -701,6 +929,7 @@ fn main() { let mut limit: Option = None; let mut weights_dir: Option = None; let mut sweep = false; + let mut rerank_sweep = false; let mut args = std::env::args().skip(1); while let Some(arg) = args.next() { match arg.as_str() { @@ -709,6 +938,7 @@ fn main() { limit = Some(v.parse().expect("--limit must be a positive integer")); } "--sweep" => sweep = true, + "--rerank-sweep" => rerank_sweep = true, "--embeddings" => { weights_dir = Some(args.next().expect("--embeddings needs a directory")); } @@ -727,6 +957,9 @@ fn main() { BM25-only, vector-only, and hybrid separately. Requires\n\ --features embeddings; without it the vector stage is\n\ inert and only the BM25 row is produced.\n\ + --rerank-sweep\n\ + compare re-ranking off, metadata-only (the old\n\ + behaviour) and blended at several half-lives.\n\ --sweep instead of the three named modes, sweep vector_weight\n\ from 0.0 to 1.0 in 0.1 steps. The 0.7/0.3 default was\n\ never searched; this is what searches it." @@ -792,6 +1025,10 @@ fn main() { { if sweep { sweep_modes() + } else if rerank_sweep { + let mut modes = vec![HYBRID, hybrid_rerank_metadata_only()]; + modes.extend(hybrid_rerank_half_lives()); + modes } else { vec![ BM25_ONLY, @@ -800,6 +1037,8 @@ fn main() { RRF, BM25_STEMMED, HYBRID_STEMMED, + hybrid_rerank_metadata_only(), + hybrid_rerank_blended(), ] } } @@ -916,6 +1155,14 @@ fn run_mode( entry.rr_turn += rr; overall.rr_turn += rr; } + if let Some(newest_first) = result.newest_gold_first { + entry.newest_gold_total += 1; + overall.newest_gold_total += 1; + if newest_first { + entry.newest_gold_first += 1; + overall.newest_gold_first += 1; + } + } let ns = result.latency.as_nanos() as u64; entry.latency_ns.push(ns); @@ -927,3 +1174,30 @@ fn run_mode( eprintln!(); print_report(&overall, &by_type, profile, mode); } + +#[cfg(test)] +mod tests { + use super::session_time; + + #[test] + fn session_dates_parse_to_the_right_instant() { + // Reference values from Python's datetime, UTC. + for (date, expected) in [ + ("2023/05/25 (Thu) 20:21", 1_685_046_060.0), + ("1970/01/01 (Thu) 00:00", 0.0), + ("2000/02/29 (Tue) 12:00", 951_825_600.0), + ("2023/12/31 (Sun) 23:59", 1_704_067_140.0), + ("2024/03/01 (Fri) 00:00", 1_709_251_200.0), + ] { + assert_eq!(session_time(date, 0), expected, "{date}"); + } + } + + #[test] + fn unparseable_dates_fall_back_to_position_order() { + let a = session_time("not a date", 0); + let b = session_time("", 1); + let c = session_time("2023/13/99 (???) 99:99", 2); + assert!(a < b && b < c, "fallback must preserve session order"); + } +}