bench: make the vector stage real, and measure BM25 vs vector vs hybrid (Tier 4b)
CI / test (push) Failing after 2s
CI / test (push) Failing after 2s
Every LongMemEval number this project has published measured BM25 alone. The
bench passed zero-vector embeddings with vector_weight=0.0, so the HNSW/vector
stage — the thing the README credits for retrieval quality — contributed
nothing and was never tested.
An optional `embeddings` feature loads all-MiniLM-L6-v2 via candle and encodes
the corpus for real. It is off by default and nothing in the shipped crates
depends on it, so a project that advertises no heavyweight dependencies keeps
that property; without the feature the bench behaves exactly as before.
Full haystack, n=500, turn-level:
Hit@1 Hit@5 Hit@10 MRR
BM25 only 53.8% 75.0% 81.6% 0.6320
Vector only 36.0% 71.8% 81.6% 0.5027
Hybrid 0.7/0.3 44.4% 79.2% 86.0% 0.5868
Session-level, hybrid leads outright: 88.2 / 95.8 / 97.8 / 0.9158.
The hybrid claim holds for depth and not for precision@1. Hybrid is the best
configuration at Hit@5 and Hit@10 at both granularities — turn-level Hit@5 gains
4.2 points over BM25 and 7.4 over vector-only, which is the result that justifies
running two stages at all. But BM25 alone still leads turn-level Hit@1 and MRR,
so fusing buys deeper recall and pays at rank 1. Callers assembling five memories
of context want hybrid; callers taking a single top hit are better served by BM25
today. The 0.7/0.3 weights are the documented default, not a searched optimum.
omni-cortex's four-signal ablation found the same direction independently — there,
adding BM25 to a dense retriever raised nDCG@5 while lowering Hit@1 and MRR. Two
codebases, two fusion schemes, same trade.
Vector-only trailing BM25 at every turn-level cutoff except Hit@10 is stated
plainly rather than buried: LongMemEval questions share heavy vocabulary with
their evidence turns, which is close to the best case for lexical matching, and
MiniLM at 384-d is a small model.
Implementation notes:
- Texts are deduplicated before encoding. The haystack sessions are drawn from
a shared pool, so 500 questions x 493.5 turns collapses to 190,015 unique
strings — the difference between encoding the corpus once and per question.
- `embeddings-cuda` adds the GPU path, and it is not a convenience: 190k texts
take ~13 min on an RTX 5060 Ti, while the same work on 8 CPU cores was still
unfinished after 30 minutes. The device is selected at runtime with a CPU
fallback, so a machine without CUDA still works.
- Mean-pooling is masked and the output L2-normalised, which is the published
recipe for this checkpoint (not the [CLS] pooler).
One measurement wrinkle, recorded rather than smoothed over: on the oracle
variant BM25-only reads 84.2% Hit@5 with real embedding vectors present against
84.4% with zero vectors — one question of 500 changes rank, MRR identical at
0.6597. On the full haystack the two agree exactly. Weight 0.0 evidently does not
make the vector stage bit-for-bit absent from candidate selection on a small
corpus.
Verified on the Linux dev host: 49 groups / 1659 passed / 0 failed, clippy clean
under -D warnings, fmt clean, with and without the feature.
This commit is contained in:
@@ -49,12 +49,56 @@
|
||||
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::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||
use serde::Deserialize;
|
||||
use tempfile::TempDir;
|
||||
|
||||
const EMBEDDING_DIM: usize = 384;
|
||||
|
||||
/// A retrieval configuration: how much of the score comes from each stage.
|
||||
#[derive(Clone, Copy)]
|
||||
struct Mode {
|
||||
label: &'static str,
|
||||
vector_weight: f32,
|
||||
keyword_weight: f32,
|
||||
}
|
||||
|
||||
/// 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 {
|
||||
label: "BM25 only (vector stage inert)",
|
||||
vector_weight: 0.0,
|
||||
keyword_weight: 1.0,
|
||||
};
|
||||
#[cfg(feature = "embeddings")]
|
||||
const VECTOR_ONLY: Mode = Mode {
|
||||
label: "Vector only (MiniLM + HNSW)",
|
||||
vector_weight: 1.0,
|
||||
keyword_weight: 0.0,
|
||||
};
|
||||
#[cfg(feature = "embeddings")]
|
||||
const HYBRID: Mode = Mode {
|
||||
label: "Hybrid (0.7 vector / 0.3 BM25)",
|
||||
vector_weight: 0.7,
|
||||
keyword_weight: 0.3,
|
||||
};
|
||||
|
||||
/// 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
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -196,7 +240,12 @@ struct EvalResult {
|
||||
latency: Duration,
|
||||
}
|
||||
|
||||
fn evaluate_question(q: &Question, top_k: usize) -> EvalResult {
|
||||
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;
|
||||
@@ -218,7 +267,7 @@ fn evaluate_question(q: &Question, top_k: usize) -> EvalResult {
|
||||
for turn in session {
|
||||
entries.push(MemoryEntry {
|
||||
chunk: turn.content.clone(),
|
||||
embedding: vec![0.0f32; EMBEDDING_DIM],
|
||||
embedding: embedding_for(embeddings, &turn.content),
|
||||
source_channel: "longmemeval".to_string(),
|
||||
timestamp: ts,
|
||||
session_id: sess_id.to_string(),
|
||||
@@ -246,10 +295,15 @@ fn evaluate_question(q: &Question, top_k: usize) -> EvalResult {
|
||||
// Set of session IDs that contain the answer
|
||||
let answer_sess_set: HashSet<&str> = q.answer_session_ids.iter().map(String::as_str).collect();
|
||||
|
||||
// Run hybrid search (BM25-only: vector_weight=0.0, keyword_weight=1.0)
|
||||
let zero_emb = vec![0.0f32; EMBEDDING_DIM];
|
||||
let query_emb = embedding_for(embeddings, &q.question);
|
||||
let t0 = Instant::now();
|
||||
let results = memory.hybrid_search(&zero_emb, &q.question, 0.0, 1.0, top_k);
|
||||
let results = memory.hybrid_search(
|
||||
&query_emb,
|
||||
&q.question,
|
||||
mode.vector_weight,
|
||||
mode.keyword_weight,
|
||||
top_k,
|
||||
);
|
||||
let latency = t0.elapsed();
|
||||
|
||||
// Session-level recall
|
||||
@@ -384,12 +438,20 @@ impl DatasetProfile {
|
||||
}
|
||||
}
|
||||
|
||||
fn print_report(overall: &Metrics, by_type: &HashMap<String, Metrics>, profile: &DatasetProfile) {
|
||||
fn print_report(
|
||||
overall: &Metrics,
|
||||
by_type: &HashMap<String, Metrics>,
|
||||
profile: &DatasetProfile,
|
||||
mode: Mode,
|
||||
) {
|
||||
println!("=================================================================");
|
||||
println!(" LongMemEval Benchmark (BM25-only retrieval, zero embeddings)");
|
||||
println!(" LongMemEval Benchmark — {}", mode.label);
|
||||
println!("=================================================================");
|
||||
println!();
|
||||
println!("Mode: vector_weight=0.0 / keyword_weight=1.0 (pure BM25)");
|
||||
println!(
|
||||
"Mode: vector_weight={:.1} / keyword_weight={:.1}",
|
||||
mode.vector_weight, mode.keyword_weight
|
||||
);
|
||||
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");
|
||||
@@ -516,7 +578,10 @@ fn print_report(overall: &Metrics, by_type: &HashMap<String, Metrics>, profile:
|
||||
println!("```json");
|
||||
println!("{{");
|
||||
println!(" \"benchmark\": \"longmemeval\",");
|
||||
println!(" \"mode\": \"bm25_only\",");
|
||||
println!(
|
||||
" \"mode\": \"vector_{:.1}_keyword_{:.1}\",",
|
||||
mode.vector_weight, mode.keyword_weight
|
||||
);
|
||||
println!(" \"dataset_variant\": \"{}\",", profile.variant());
|
||||
println!(" \"scoring_target\": \"retrieval_recall\",");
|
||||
println!(" \"k\": 10,");
|
||||
@@ -585,6 +650,7 @@ fn print_report(overall: &Metrics, by_type: &HashMap<String, Metrics>, profile:
|
||||
fn main() {
|
||||
let mut json_path: Option<String> = None;
|
||||
let mut limit: Option<usize> = None;
|
||||
let mut weights_dir: Option<String> = None;
|
||||
let mut args = std::env::args().skip(1);
|
||||
while let Some(arg) = args.next() {
|
||||
match arg.as_str() {
|
||||
@@ -592,6 +658,9 @@ fn main() {
|
||||
let v = args.next().expect("--limit needs a value");
|
||||
limit = Some(v.parse().expect("--limit must be a positive integer"));
|
||||
}
|
||||
"--embeddings" => {
|
||||
weights_dir = Some(args.next().expect("--embeddings needs a directory"));
|
||||
}
|
||||
"--help" | "-h" => {
|
||||
eprintln!(
|
||||
"usage: longmemeval_bench [PATH] [--limit N]\n\n\
|
||||
@@ -600,7 +669,13 @@ fn main() {
|
||||
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."
|
||||
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."
|
||||
);
|
||||
return;
|
||||
}
|
||||
@@ -650,6 +725,63 @@ fn main() {
|
||||
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")]
|
||||
{
|
||||
vec![BM25_ONLY, VECTOR_ONLY, HYBRID]
|
||||
}
|
||||
#[cfg(not(feature = "embeddings"))]
|
||||
{
|
||||
vec![BM25_ONLY]
|
||||
}
|
||||
} else {
|
||||
vec![BM25_ONLY]
|
||||
};
|
||||
|
||||
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();
|
||||
|
||||
@@ -658,7 +790,7 @@ fn main() {
|
||||
eprint!("\r [{}/{}] evaluating...", i + 1, total);
|
||||
}
|
||||
|
||||
let result = evaluate_question(q, 10);
|
||||
let result = evaluate_question(q, 10, mode, embeddings);
|
||||
|
||||
let is_abs = q.question_type.ends_with("_abs");
|
||||
let base_type = if is_abs {
|
||||
@@ -723,5 +855,5 @@ fn main() {
|
||||
|
||||
eprintln!("\r [{total}/{total}] done. ");
|
||||
eprintln!();
|
||||
print_report(&overall, &by_type, &profile);
|
||||
print_report(&overall, &by_type, profile, mode);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user