Files
clawhdf5/crates/clawhdf5-bench/src/bin/longmemeval_bench.rs
T
2026-09-19 19:42:05 -07:00

1215 lines
44 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! LongMemEval Benchmark Harness (Track 8.1)
//!
//! Evaluates BM25-based retrieval recall against the LongMemEval dataset (500 questions).
//! Since no embedding model is available at bench time, all embeddings are zero vectors
//! and `hybrid_search` operates in BM25-only mode (vector_weight=0.0, keyword_weight=1.0).
//!
//! # Scoring target (read before citing any number from this harness)
//!
//! - **Metric: retrieval recall.** A "hit" means the gold-labelled memory appeared in
//! the top-k. No answer is generated and none is scored — the dataset's `answer`
//! field is deserialized and deliberately never read. This is **not** the official
//! LongMemEval metric, which is end-to-end QA accuracy (retrieve → generate → LLM
//! judge). Reporting retrieval recall as QA accuracy overstates by 2030 points.
//! - **Dataset: whichever variant you point it at.** Both `longmemeval_oracle`
//! (evidence sessions only — a substantially easier corpus) and the full
//! `longmemeval_s` haystack are supported. The harness does not trust the
//! filename: [`DatasetProfile`] measures evidence-session density from the
//! data and labels the run from that, so a mislabelled input cannot produce a
//! mislabelled result.
//! - **Session-level metrics are degenerate when evidence density is high**, and
//! the report says so per run rather than assuming it. On the oracle variant
//! the haystack is essentially all-evidence, so any returned document is a
//! session-level hit at rank 0 by construction; only turn-level
//! (`has_answer == true` on the source turn) measures the retriever there. On
//! the full haystack, session-level recall is meaningful.
//! - **Not comparable to MemX's Hit@5=51.6% / MRR=0.380**, which is *fact-level*
//! granularity over 220,349 records from 19,195 sessions.
//!
//! See `BENCHMARKS.md` § "Retracted: session-level recall and the MemX comparison".
//!
//! # Usage
//! ```
//! cargo run --release --bin longmemeval_bench [PATH] [--limit N]
//!
//! # Usage: full haystack
//! ```
//! cargo run --release --bin longmemeval_bench -- \
//! benchmarks/longmemeval/longmemeval_s_cleaned.json --limit 50
//! ```
//! ```
//!
//! # WASM Note
//! `#[cfg(target_arch = "wasm32")]` is not supported here. Changes required for wasm32:
//! - `std::fs::read_to_string` → fetch-based async loader (e.g. `wasm_bindgen_futures`)
//! - `TempDir` → virtual in-memory HDF5 backend (separate effort; tracked in ROADMAP)
//! - `std::time::Instant` → `web_sys::Performance::now()`
//! - HDF5 I/O layer would need a wasm32 storage backend (out of scope for this bench)
use std::collections::{HashMap, HashSet};
use std::time::{Duration, Instant};
// `#[path]` keeps the module beside its binary without Cargo autodiscovering it
// as a second bin target (which a bare `src/bin/embedder.rs` would be).
#[cfg(feature = "embeddings")]
#[path = "longmemeval_bench/embedder.rs"]
mod embedder;
use clawhdf5_agent::bm25::TokenFilter;
use clawhdf5_agent::hybrid::Fusion;
use clawhdf5_agent::reranker::{ReRankConfig, RerankInput, rerank};
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchResult};
use serde::Deserialize;
use tempfile::TempDir;
const EMBEDDING_DIM: usize = 384;
/// A mode's fusion, as one short string for the reports.
fn describe(mode: Mode) -> String {
let fusion = match mode.fusion {
Fusion::Weighted { vector, keyword } => format!("vector_{vector:.1}_keyword_{keyword:.1}"),
Fusion::Rrf { k } => format!("rrf_k{k:.0}"),
};
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
),
}
}
/// A retrieval configuration: how much of the score comes from each stage.
#[derive(Clone, Copy)]
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,
/// Re-rank the retrieved candidates with recency and friends, relative to
/// the question's own date.
rerank: Option<ReRankConfig>,
}
impl Mode {
const fn weighted(label: &'static str, vector: f32, keyword: f32) -> Self {
Self {
label,
fusion: Fusion::Weighted { vector, keyword },
tokens: TokenFilter::Plain,
rerank: None,
}
}
#[cfg_attr(not(feature = "embeddings"), allow(dead_code))]
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;
self
}
}
/// The only mode available without real embeddings. Passing zero vectors with
/// `vector_weight = 0.0` is what made the vector stage inert.
const BM25_ONLY: Mode = Mode::weighted("BM25 only (vector stage inert)", 0.0, 1.0);
#[cfg(feature = "embeddings")]
const VECTOR_ONLY: Mode = Mode::weighted("Vector only (MiniLM + HNSW)", 1.0, 0.0);
/// Tuned by `--sweep` over the full haystack. The former 0.7/0.3 was a
/// documented default that had never been searched, and the sweep found it
/// strictly dominated: 0.4/0.6 is better on Hit@1, Hit@5, Hit@10 and MRR at
/// both granularities.
#[cfg(feature = "embeddings")]
const HYBRID: Mode = Mode::weighted("Hybrid (0.4 vector / 0.6 BM25, tuned)", 0.4, 0.6);
/// Reciprocal rank fusion, the documented alternative to the weighted sum.
/// It ignores score magnitudes, so there is nothing to tune — which is the
/// claim being tested.
#[cfg(feature = "embeddings")]
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<Mode> {
[
("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");
/// Every 0.1 step of vector weight, keyword weight taking the remainder.
///
/// Labels are leaked to `&'static str` because `Mode::label` is a `&'static
/// str` for the eleven named modes and a sweep is a short-lived process; the
/// alternative is threading a lifetime through the whole report path for a
/// diagnostic mode.
#[cfg(feature = "embeddings")]
fn sweep_modes() -> Vec<Mode> {
(0..=10)
.map(|i| {
let v = i as f32 / 10.0;
Mode::weighted(
Box::leak(format!("sweep v={v:.1} / k={:.1}", 1.0 - v).into_boxed_str()),
v,
1.0 - v,
)
})
.collect()
}
/// Text -> embedding, built once for the whole corpus.
type EmbeddingMap = HashMap<String, Vec<f32>>;
/// Look up a real embedding, falling back to zeros when running BM25-only.
fn embedding_for(map: Option<&EmbeddingMap>, text: &str) -> Vec<f32> {
map.and_then(|m| m.get(text))
.cloned()
.unwrap_or_else(|| vec![0.0f32; EMBEDDING_DIM])
}
// ---------------------------------------------------------------------------
// JSON data types
// ---------------------------------------------------------------------------
/// The `answer` field in LongMemEval can be a string or a number.
fn deserialize_answer<'de, D>(deserializer: D) -> Result<String, D::Error>
where
D: serde::Deserializer<'de>,
{
use serde::de;
struct AnswerVisitor;
impl<'de> de::Visitor<'de> for AnswerVisitor {
type Value = String;
fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("a string or number")
}
fn visit_str<E: de::Error>(self, v: &str) -> Result<String, E> {
Ok(v.to_string())
}
fn visit_string<E: de::Error>(self, v: String) -> Result<String, E> {
Ok(v)
}
fn visit_i64<E: de::Error>(self, v: i64) -> Result<String, E> {
Ok(v.to_string())
}
fn visit_u64<E: de::Error>(self, v: u64) -> Result<String, E> {
Ok(v.to_string())
}
fn visit_f64<E: de::Error>(self, v: f64) -> Result<String, E> {
Ok(v.to_string())
}
}
deserializer.deserialize_any(AnswerVisitor)
}
#[derive(Deserialize)]
struct Turn {
#[allow(dead_code)]
role: String,
content: String,
#[serde(default)]
has_answer: bool,
}
#[derive(Deserialize)]
struct Question {
#[allow(dead_code)]
question_id: String,
question_type: String,
question: String,
#[allow(dead_code)]
#[serde(deserialize_with = "deserialize_answer")]
answer: String,
#[allow(dead_code)]
question_date: String,
haystack_session_ids: Vec<String>,
haystack_sessions: Vec<Vec<Turn>>,
answer_session_ids: Vec<String>,
/// One timestamp per haystack session, e.g. "2023/05/25 (Thu) 20:21".
#[serde(default)]
haystack_dates: Vec<String>,
}
/// 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<f64> {
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)
}
// ---------------------------------------------------------------------------
// Per-type metrics accumulator
// ---------------------------------------------------------------------------
#[derive(Default)]
struct Metrics {
hit1_session: u32,
hit5_session: u32,
hit10_session: u32,
rr_session: f64,
hit1_turn: u32,
hit5_turn: u32,
hit10_turn: u32,
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<u64>,
count: u32,
}
impl Metrics {
/// `None` when no question in this bucket had multiple gold sessions.
fn newest_gold_first_pct(&self) -> Option<f64> {
(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
}
fn hit5_session_pct(&self) -> f64 {
self.hit5_session as f64 / self.count.max(1) as f64 * 100.0
}
fn hit10_session_pct(&self) -> f64 {
self.hit10_session as f64 / self.count.max(1) as f64 * 100.0
}
fn mrr_session(&self) -> f64 {
self.rr_session / self.count.max(1) as f64
}
fn hit1_turn_pct(&self) -> f64 {
self.hit1_turn as f64 / self.count.max(1) as f64 * 100.0
}
fn hit5_turn_pct(&self) -> f64 {
self.hit5_turn as f64 / self.count.max(1) as f64 * 100.0
}
fn hit10_turn_pct(&self) -> f64 {
self.hit10_turn as f64 / self.count.max(1) as f64 * 100.0
}
fn mrr_turn(&self) -> f64 {
self.rr_turn / self.count.max(1) as f64
}
fn abstention_pct(&self) -> f64 {
self.abstention_correct as f64 / self.abstention_total.max(1) as f64 * 100.0
}
fn latency_avg_us(&self) -> f64 {
if self.latency_ns.is_empty() {
return 0.0;
}
self.latency_ns.iter().sum::<u64>() as f64 / self.latency_ns.len() as f64 / 1000.0
}
fn latency_pct_us(&self, p: usize) -> f64 {
if self.latency_ns.is_empty() {
return 0.0;
}
let mut v = self.latency_ns.clone();
v.sort_unstable();
let idx = (p * v.len() / 100).min(v.len() - 1);
v[idx] as f64 / 1000.0
}
}
// ---------------------------------------------------------------------------
// Per-question evaluation result
// ---------------------------------------------------------------------------
struct EvalResult {
hit1_session: bool,
hit5_session: bool,
hit10_session: bool,
rr_session: Option<f64>,
hit1_turn: bool,
hit5_turn: bool,
hit10_turn: bool,
rr_turn: Option<f64>,
/// 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<bool>,
latency: Duration,
}
fn evaluate_question(
q: &Question,
top_k: usize,
mode: Mode,
embeddings: Option<&EmbeddingMap>,
) -> EvalResult {
let dir = TempDir::new().expect("failed to create temp dir");
let mut config = MemoryConfig::new(dir.path().join("lme.h5"), "lme-bench", EMBEDDING_DIM);
config.wal_enabled = false;
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();
let mut turn_has_answer: Vec<bool> = Vec::new();
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");
// 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),
source_channel: "longmemeval".to_string(),
timestamp: ts,
session_id: sess_id.to_string(),
tags: if turn.has_answer {
"has_answer".to_string()
} else {
String::new()
},
});
turn_has_answer.push(turn.has_answer);
}
}
let indices = memory.save_batch(entries).expect("failed to save entries");
// Map memory index → has_answer
let has_answer_indices: HashSet<usize> = indices
.iter()
.zip(turn_has_answer.iter())
.filter(|(_, ha)| **ha)
.map(|(idx, _)| *idx)
.collect();
// 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();
// 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<RerankInput> = 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<usize> = rerank(&inputs, &config, now)
.into_iter()
.map(|r| r.index)
.collect();
let by_index: HashMap<usize, SearchResult> =
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;
let mut hit10_session = false;
let mut rr_session: Option<f64> = None;
for (rank, result) in results.iter().enumerate() {
let sess_id = memory.cache.session_ids[result.index].as_str();
if answer_sess_set.contains(sess_id) {
hit10_session = true;
if rank < 5 {
hit5_session = true;
}
if rank == 0 {
hit1_session = true;
}
if rr_session.is_none() {
rr_session = Some(1.0 / (rank + 1) as f64);
}
break;
}
}
// Turn-level recall
let mut hit1_turn = false;
let mut hit5_turn = false;
let mut hit10_turn = false;
let mut rr_turn: Option<f64> = None;
for (rank, result) in results.iter().enumerate() {
if has_answer_indices.contains(&result.index) {
hit10_turn = true;
if rank < 5 {
hit5_turn = true;
}
if rank == 0 {
hit1_turn = true;
}
if rr_turn.is_none() {
rr_turn = Some(1.0 / (rank + 1) as f64);
}
break;
}
}
EvalResult {
hit1_session,
hit5_session,
hit10_session,
rr_session,
hit1_turn,
hit5_turn,
hit10_turn,
rr_turn,
newest_gold_first,
latency,
}
}
// ---------------------------------------------------------------------------
// Report printing
// ---------------------------------------------------------------------------
// ---------------------------------------------------------------------------
// Dataset profile — measured, not assumed
// ---------------------------------------------------------------------------
/// Shape of the loaded corpus, computed from the data itself.
///
/// The variant used to be a hardcoded `"oracle"` string in the report and the
/// JSON summary, so pointing the harness at `longmemeval_s` would have produced
/// full-haystack numbers labelled oracle. Everything here is derived from the
/// questions instead, which means the label cannot drift from the corpus and a
/// mislabelled input file cannot produce a mislabelled result.
struct DatasetProfile {
n_questions: usize,
mean_sessions: f64,
mean_turns: f64,
/// Mean over questions of `|answer_sessions| / |haystack_sessions|`.
///
/// This is what actually decides whether session-level recall means
/// anything. At ~1.0 every haystack session is an evidence session, so any
/// returned document is a session-level hit by construction.
evidence_density: f64,
}
impl DatasetProfile {
fn measure(questions: &[Question]) -> Self {
let n = questions.len().max(1) as f64;
let mut sessions = 0.0;
let mut turns = 0.0;
let mut density = 0.0;
for q in questions {
let n_sess = q.haystack_sessions.len();
sessions += n_sess as f64;
turns += q.haystack_sessions.iter().map(Vec::len).sum::<usize>() as f64;
if n_sess > 0 {
let evidence: HashSet<&str> =
q.answer_session_ids.iter().map(String::as_str).collect();
let hit = q
.haystack_session_ids
.iter()
.filter(|id| evidence.contains(id.as_str()))
.count();
density += hit as f64 / n_sess as f64;
}
}
Self {
n_questions: questions.len(),
mean_sessions: sessions / n,
mean_turns: turns / n,
evidence_density: density / n,
}
}
/// Above this share of evidence sessions, session-level recall is measuring
/// the corpus shape rather than the retriever.
const DEGENERACY_THRESHOLD: f64 = 0.9;
const fn session_level_degenerate(&self) -> bool {
self.evidence_density > Self::DEGENERACY_THRESHOLD
}
/// Variant name inferred from evidence density, not from the filename.
const fn variant(&self) -> &'static str {
if self.session_level_degenerate() {
"oracle"
} else {
"full_haystack"
}
}
}
fn print_report(
overall: &Metrics,
by_type: &HashMap<String, Metrics>,
profile: &DatasetProfile,
mode: Mode,
) {
println!("=================================================================");
println!(" LongMemEval Benchmark — {}", mode.label);
println!("=================================================================");
println!();
println!("Mode: {}", describe(mode));
println!();
println!("Scoring target: RETRIEVAL RECALL (did the gold memory land in top-k).");
println!(" No answer is generated or scored. This is NOT the official");
println!(" LongMemEval metric (QA accuracy via retrieve+generate+judge).");
println!(
"Dataset: {}{} questions, {:.1} sessions and {:.0} turns per question,",
profile.variant(),
profile.n_questions,
profile.mean_sessions,
profile.mean_turns,
);
println!(
" {:.1}% of haystack sessions are evidence sessions.",
profile.evidence_density * 100.0
);
if profile.session_level_degenerate() {
println!(" This is the evidence-only corpus, NOT the full longmemeval_s");
println!(" haystack — a substantially easier retrieval problem.");
} else {
println!(" This is a full-haystack corpus: evidence sessions are a small");
println!(" minority, so retrieval has to actually discriminate.");
}
println!();
println!("Do NOT compare these to MemX's Hit@5=51.6% / MRR=0.380: that is");
println!(" fact-level granularity over 220,349 records from 19,195 sessions.");
println!(" Different granularity and a corpus larger by orders of magnitude.");
println!();
println!("## Session-Level Recall (n={})", overall.count);
if profile.session_level_degenerate() {
println!(
" [DEGENERATE — {:.1}% of haystack sessions are evidence sessions, so a",
profile.evidence_density * 100.0
);
println!(" returned document is a session-level hit almost by construction.");
println!(" This measures the corpus shape, not the retriever. Use turn-level.]");
} else {
println!(
" [Meaningful on this corpus — only {:.1}% of haystack sessions are",
profile.evidence_density * 100.0
);
println!(" evidence sessions, so a hit reflects the retriever's discrimination.]");
}
println!(
" Hit@1: {:5.1}% Hit@5: {:5.1}% Hit@10: {:5.1}% MRR: {:.4}",
overall.hit1_session_pct(),
overall.hit5_session_pct(),
overall.hit10_session_pct(),
overall.mrr_session()
);
println!();
println!("## Turn-Level Recall (n={})", overall.count);
println!(
" Hit@1: {:5.1}% Hit@5: {:5.1}% Hit@10: {:5.1}% MRR: {:.4}",
overall.hit1_turn_pct(),
overall.hit5_turn_pct(),
overall.hit10_turn_pct(),
overall.mrr_turn()
);
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!(
" Correct: {}/{} ({:.1}%)",
overall.abstention_correct,
overall.abstention_total,
overall.abstention_pct()
);
println!(" (abstention = system correctly returns no matching session)");
println!();
}
println!(
"## Search Latency (n={} queries, BM25 over variable haystack sizes)",
overall.latency_ns.len()
);
println!(
" avg={:.1} µs p50={:.1} µs p95={:.1} µs p99={:.1} µs",
overall.latency_avg_us(),
overall.latency_pct_us(50),
overall.latency_pct_us(95),
overall.latency_pct_us(99)
);
println!();
println!("## Per-Type Breakdown (session-level)");
println!();
println!(
"{:<32} {:>5} {:>7} {:>7} {:>7} {:>7}",
"Question Type", "N", "Hit@1", "Hit@5", "Hit@10", "MRR"
);
println!("{}", "-".repeat(72));
let mut types: Vec<(&String, &Metrics)> = by_type.iter().collect();
types.sort_by_key(|(t, _)| t.as_str());
for (qtype, m) in &types {
if m.count > 0 {
println!(
"{:<32} {:>5} {:>6.1}% {:>6.1}% {:>6.1}% {:>7.4}",
qtype,
m.count,
m.hit1_session_pct(),
m.hit5_session_pct(),
m.hit10_session_pct(),
m.mrr_session()
);
}
if m.abstention_total > 0 {
println!(
"{:<32} {:>5} abstention accuracy: {:>5.1}%",
format!("{qtype}_abs"),
m.abstention_total,
m.abstention_pct()
);
}
}
// Machine-parseable JSON summary
println!();
println!("## JSON Summary");
println!("```json");
println!("{{");
println!(" \"benchmark\": \"longmemeval\",");
println!(" \"mode\": \"{}\",", describe(mode));
println!(" \"dataset_variant\": \"{}\",", profile.variant());
println!(" \"scoring_target\": \"retrieval_recall\",");
println!(" \"k\": 10,");
println!(
" \"session_level_degenerate\": {},",
profile.session_level_degenerate()
);
println!(
" \"evidence_session_density\": {:.4},",
profile.evidence_density
);
println!(
" \"mean_sessions_per_question\": {:.2},",
profile.mean_sessions
);
println!(" \"mean_turns_per_question\": {:.1},", profile.mean_turns);
println!(
" \"total_questions\": {},",
overall.count + overall.abstention_total
);
println!(" \"session_level\": {{");
println!(
" \"hit_at_1\": {:.4}, \"hit_at_5\": {:.4}, \"hit_at_10\": {:.4}, \"mrr\": {:.4}",
overall.hit1_session_pct() / 100.0,
overall.hit5_session_pct() / 100.0,
overall.hit10_session_pct() / 100.0,
overall.mrr_session()
);
println!(" }},");
println!(" \"turn_level\": {{");
println!(
" \"hit_at_1\": {:.4}, \"hit_at_5\": {:.4}, \"hit_at_10\": {:.4}, \"mrr\": {:.4}",
overall.hit1_turn_pct() / 100.0,
overall.hit5_turn_pct() / 100.0,
overall.hit10_turn_pct() / 100.0,
overall.mrr_turn()
);
println!(" }},");
// `null`, not 0.0 — a corpus with no abstention questions has no abstention
// accuracy, and emitting 0.0 reads as total failure at a task never posed.
if overall.abstention_total > 0 {
println!(
" \"abstention_accuracy\": {:.4},",
overall.abstention_pct() / 100.0
);
} 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}",
overall.latency_avg_us(),
overall.latency_pct_us(50),
overall.latency_pct_us(95),
overall.latency_pct_us(99)
);
println!(" }}");
println!("}}");
println!("```");
}
// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------
fn main() {
let mut json_path: Option<String> = None;
let mut limit: Option<usize> = None;
let mut weights_dir: Option<String> = None;
let mut sweep = false;
#[cfg_attr(not(feature = "embeddings"), allow(unused_mut, unused_variables))]
let mut rerank_sweep = false;
let mut args = std::env::args().skip(1);
while let Some(arg) = args.next() {
match arg.as_str() {
"--limit" => {
let v = args.next().expect("--limit needs a value");
limit = Some(v.parse().expect("--limit must be a positive integer"));
}
"--sweep" => sweep = true,
"--rerank-sweep" => {
// Re-ranking needs the vector stage to have candidates worth
// reordering, so this is an embeddings-only comparison.
#[cfg(feature = "embeddings")]
{
rerank_sweep = true;
}
#[cfg(not(feature = "embeddings"))]
eprintln!("warning: --rerank-sweep needs --features embeddings; ignoring");
}
"--embeddings" => {
weights_dir = Some(args.next().expect("--embeddings needs a directory"));
}
"--help" | "-h" => {
eprintln!(
"usage: longmemeval_bench [PATH] [--limit N]\n\n\
PATH dataset JSON; defaults to the oracle variant.\n\
longmemeval_s works too — the harness measures which\n\
variant it was given rather than trusting the filename.\n\
--limit evaluate N questions, sampled evenly across the file\n\
rather than as a prefix — the dataset is ordered by\n\
question type, so a prefix samples one type only.\n\
--embeddings DIR\n\
directory holding all-MiniLM-L6-v2's model.safetensors\n\
and tokenizer.json. Enables the vector stage and reports\n\
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."
);
return;
}
other => json_path = Some(other.to_string()),
}
}
let json_path =
json_path.unwrap_or_else(|| "benchmarks/longmemeval/longmemeval_oracle.json".to_string());
eprintln!("Loading: {json_path}");
let data = std::fs::read_to_string(&json_path)
.unwrap_or_else(|e| panic!("Failed to read {json_path}: {e}"));
let mut questions: Vec<Question> = serde_json::from_str(&data).expect("Failed to parse JSON");
if let Some(n) = limit
&& n < questions.len()
{
// Stride rather than truncate. The dataset is ordered by question type,
// so taking a prefix samples one type: `--limit 20` on longmemeval_s
// returns 20 `single-session-user` questions and nothing else, which
// reads as a whole-dataset result but is not one.
let total = questions.len();
let step = total as f64 / n as f64;
let keep: HashSet<usize> = (0..n)
.map(|i| ((i as f64 * step) as usize).min(total - 1))
.collect();
questions = questions
.into_iter()
.enumerate()
.filter(|(i, _)| keep.contains(i))
.map(|(_, q)| q)
.collect();
eprintln!(
"Sampling {} of {total} questions, evenly strided (--limit)",
questions.len()
);
}
let total = questions.len();
eprintln!("Loaded {total} questions");
let profile = DatasetProfile::measure(&questions);
eprintln!(
"Corpus: {} variant — {:.1} sessions / {:.0} turns per question, \
{:.1}% evidence-session density",
profile.variant(),
profile.mean_sessions,
profile.mean_turns,
profile.evidence_density * 100.0,
);
// Build the embedding table once for the whole corpus, if asked for.
let embeddings: Option<EmbeddingMap> = weights_dir
.as_deref()
.map(|dir| load_embeddings(dir, &questions));
if embeddings.is_none() && weights_dir.is_some() {
eprintln!("warning: --embeddings ignored (build with --features embeddings)");
}
let modes: Vec<Mode> = if embeddings.is_some() {
#[cfg(feature = "embeddings")]
{
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,
VECTOR_ONLY,
HYBRID,
RRF,
BM25_STEMMED,
HYBRID_STEMMED,
hybrid_rerank_metadata_only(),
hybrid_rerank_blended(),
]
}
}
#[cfg(not(feature = "embeddings"))]
{
vec![BM25_ONLY, BM25_STEMMED]
}
} else {
if sweep {
eprintln!("warning: --sweep needs --embeddings; running 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() {
eprintln!("[{}/{}] {}", mode_idx + 1, modes.len(), mode.label);
run_mode(&questions, *mode, embeddings.as_ref(), &profile);
}
}
/// Load and encode the corpus. Returns `None` unless the `embeddings` feature
/// is compiled in, so the flag degrades to a warning rather than a hard error.
#[cfg(feature = "embeddings")]
fn load_embeddings(dir: &str, questions: &[Question]) -> EmbeddingMap {
let enc = embedder::Embedder::load(std::path::Path::new(dir))
.unwrap_or_else(|e| panic!("failed to load embedder from {dir}: {e}"));
let texts = questions.iter().flat_map(|q| {
q.haystack_sessions
.iter()
.flatten()
.map(|t| t.content.clone())
.chain(std::iter::once(q.question.clone()))
});
enc.encode_unique(texts)
.unwrap_or_else(|e| panic!("embedding failed: {e}"))
}
#[cfg(not(feature = "embeddings"))]
fn load_embeddings(_dir: &str, _questions: &[Question]) -> EmbeddingMap {
EmbeddingMap::new()
}
/// Evaluate every question under one retrieval mode and print its report.
fn run_mode(
questions: &[Question],
mode: Mode,
embeddings: Option<&EmbeddingMap>,
profile: &DatasetProfile,
) {
let total = questions.len();
let mut overall = Metrics::default();
let mut by_type: HashMap<String, Metrics> = HashMap::new();
for (i, q) in questions.iter().enumerate() {
if (i + 1) % 50 == 0 || i + 1 == total {
eprint!("\r [{}/{}] evaluating...", i + 1, total);
}
let result = evaluate_question(q, 10, mode, embeddings);
let is_abs = q.question_type.ends_with("_abs");
let base_type = if is_abs {
q.question_type.trim_end_matches("_abs").to_string()
} else {
q.question_type.clone()
};
let entry = by_type.entry(base_type).or_default();
if is_abs {
entry.abstention_total += 1;
overall.abstention_total += 1;
// Correct abstention: no session-level hit in top 10
if !result.hit10_session {
entry.abstention_correct += 1;
overall.abstention_correct += 1;
}
} else {
entry.count += 1;
overall.count += 1;
if result.hit1_session {
entry.hit1_session += 1;
overall.hit1_session += 1;
}
if result.hit5_session {
entry.hit5_session += 1;
overall.hit5_session += 1;
}
if result.hit10_session {
entry.hit10_session += 1;
overall.hit10_session += 1;
}
if let Some(rr) = result.rr_session {
entry.rr_session += rr;
overall.rr_session += rr;
}
if result.hit1_turn {
entry.hit1_turn += 1;
overall.hit1_turn += 1;
}
if result.hit5_turn {
entry.hit5_turn += 1;
overall.hit5_turn += 1;
}
if result.hit10_turn {
entry.hit10_turn += 1;
overall.hit10_turn += 1;
}
if let Some(rr) = result.rr_turn {
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);
overall.latency_ns.push(ns);
}
}
eprintln!("\r [{total}/{total}] done. ");
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");
}
}