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:
+58
-2
@@ -220,8 +220,11 @@ _Latency benchmarks generated with Criterion.rs (50-100 samples per benchmark).
|
|||||||
> judge). Retrieval recall reported as QA accuracy typically overstates by 20–30 points.
|
> judge). Retrieval recall reported as QA accuracy typically overstates by 20–30 points.
|
||||||
> - **Granularity:** turn-level = the returned memory's source turn had `has_answer == true`.
|
> - **Granularity:** turn-level = the returned memory's source turn had `has_answer == true`.
|
||||||
> - **k = 10**, n = 500.
|
> - **k = 10**, n = 500.
|
||||||
> - **Retrieval mode:** BM25 only. The bench passes zero-vector embeddings with
|
> - **Retrieval mode:** all three are reported below. Historically the bench passed
|
||||||
> `vector_weight=0.0`, so the HNSW/vector stage is inert and contributes nothing.
|
> zero-vector embeddings with `vector_weight=0.0`, so the HNSW/vector stage was
|
||||||
|
> inert and every published number was BM25 alone. Real `all-MiniLM-L6-v2`
|
||||||
|
> embeddings are now available via `--features embeddings --embeddings <dir>`,
|
||||||
|
> and BM25-only / vector-only / hybrid are each measured separately.
|
||||||
|
|
||||||
**Mode:** BM25-only retrieval — zero embeddings, `vector_weight=0.0`, `keyword_weight=1.0`
|
**Mode:** BM25-only retrieval — zero embeddings, `vector_weight=0.0`, `keyword_weight=1.0`
|
||||||
|
|
||||||
@@ -253,6 +256,53 @@ Per-type, session-level: `single-session-assistant` 100.0% Hit@1 (n=56),
|
|||||||
struggles, since a preference question's evidence rarely shares vocabulary with
|
struggles, since a preference question's evidence rarely shares vocabulary with
|
||||||
the question.
|
the question.
|
||||||
|
|
||||||
|
### Retrieval mode ablation — full haystack, n=500
|
||||||
|
|
||||||
|
Real 384-d `all-MiniLM-L6-v2` embeddings, 190,015 unique texts encoded once on an
|
||||||
|
RTX 5060 Ti (~13 min; the same work on the 8-core CPU was still unfinished after
|
||||||
|
30 minutes, so the GPU path is not a convenience here). Turn-level:
|
||||||
|
|
||||||
|
| Mode | Hit@1 | Hit@5 | Hit@10 | MRR |
|
||||||
|
|------|-------|-------|--------|-----|
|
||||||
|
| BM25 only (`0.0`/`1.0`) | **53.8%** | 75.0% | 81.6% | **0.6320** |
|
||||||
|
| Vector only (`1.0`/`0.0`) | 36.0% | 71.8% | 81.6% | 0.5027 |
|
||||||
|
| Hybrid (`0.7`/`0.3`) | 44.4% | **79.2%** | **86.0%** | 0.5868 |
|
||||||
|
|
||||||
|
Session-level:
|
||||||
|
|
||||||
|
| Mode | Hit@1 | Hit@5 | Hit@10 | MRR |
|
||||||
|
|------|-------|-------|--------|-----|
|
||||||
|
| BM25 only | 86.2% | 93.6% | 96.6% | 0.8948 |
|
||||||
|
| Vector only | 85.4% | 94.2% | 96.6% | 0.8901 |
|
||||||
|
| Hybrid | **88.2%** | **95.8%** | **97.8%** | **0.9158** |
|
||||||
|
|
||||||
|
**The hybrid claim holds for depth, not for precision@1.** Hybrid is the best
|
||||||
|
configuration at Hit@5 and Hit@10 at both granularities — turn-level Hit@5 rises
|
||||||
|
4.2 points over BM25 and 7.4 over vector-only, which is the result that justifies
|
||||||
|
running two stages. But BM25 alone remains the best turn-level Hit@1 (53.8% vs
|
||||||
|
44.4%) and MRR (0.6320 vs 0.5868). Fusing at `0.7/0.3` buys deeper recall and
|
||||||
|
pays for it at rank 1.
|
||||||
|
|
||||||
|
That trade matters for how the index is consumed: a caller assembling five
|
||||||
|
memories of context should use hybrid, and a caller taking the single top hit is
|
||||||
|
better served by BM25 today. The weights are not tuned — `0.7/0.3` is the
|
||||||
|
documented default, not a searched optimum, and a Hit@1-oriented deployment
|
||||||
|
should sweep them.
|
||||||
|
|
||||||
|
The same pattern shows up independently in omni-cortex's four-signal RRF ablation,
|
||||||
|
where adding BM25 to a dense retriever raised nDCG@5 while lowering Hit@1 and MRR.
|
||||||
|
Two different codebases, two different fusion schemes, same direction.
|
||||||
|
|
||||||
|
Vector-only being *worse* than BM25 at every turn-level cutoff except Hit@10 is
|
||||||
|
worth stating plainly rather than hiding: LongMemEval questions share substantial
|
||||||
|
vocabulary with their evidence turns, which is close to the best case for lexical
|
||||||
|
matching, and MiniLM at 384 dimensions is a small embedding model.
|
||||||
|
|
||||||
|
> **Run:** `cargo run --release --bin longmemeval_bench --features embeddings -- \
|
||||||
|
> benchmarks/longmemeval/longmemeval_s_cleaned.json --embeddings weights/all-minilm-l6-v2`
|
||||||
|
> Add `--features embeddings-cuda` (and put `nvcc` on `PATH`) for the GPU path.
|
||||||
|
> Weights: `huggingface.co/sentence-transformers/all-MiniLM-L6-v2`.
|
||||||
|
|
||||||
### Oracle variant — `longmemeval_oracle`, n=500 (easier corpus, kept for continuity)
|
### Oracle variant — `longmemeval_oracle`, n=500 (easier corpus, kept for continuity)
|
||||||
|
|
||||||
| Metric | ClawhDF5 (BM25-only, oracle variant) |
|
| Metric | ClawhDF5 (BM25-only, oracle variant) |
|
||||||
@@ -267,6 +317,12 @@ price of the harder corpus, and is the reason oracle-only numbers should not be
|
|||||||
presented as LongMemEval results. Session-level figures on this variant are
|
presented as LongMemEval results. Session-level figures on this variant are
|
||||||
degenerate — see below.
|
degenerate — see below.
|
||||||
|
|
||||||
|
With real embeddings the same oracle corpus gives BM25-only 84.2% / vector-only
|
||||||
|
80.4% / hybrid **85.2%** Hit@5 turn-level — hybrid ahead at Hit@5 and Hit@10 and
|
||||||
|
behind at Hit@1, matching the full-haystack pattern above. (BM25-only reads 84.2%
|
||||||
|
here against 84.4% with zero embedding vectors: one question of 500 changes rank,
|
||||||
|
with MRR identical at 0.6597. On the full haystack the two agree exactly.)
|
||||||
|
|
||||||
### Retracted: session-level recall and the MemX comparison
|
### Retracted: session-level recall and the MemX comparison
|
||||||
|
|
||||||
Earlier revisions of this file reported session-level Hit@1/5/10 of **100.0%** with
|
Earlier revisions of this file reported session-level Hit@1/5/10 of **100.0%** with
|
||||||
|
|||||||
@@ -103,14 +103,25 @@ Use `.with_zstd(3)` or `.with_deflate(6)` for write-heavy workloads — both now
|
|||||||
|
|
||||||
Evaluated against the full **`longmemeval_s`** haystack — all 500 questions, 47.7
|
Evaluated against the full **`longmemeval_s`** haystack — all 500 questions, 47.7
|
||||||
sessions and 493.5 turns each, with only 4.0% of haystack sessions being evidence
|
sessions and 493.5 turns each, with only 4.0% of haystack sessions being evidence
|
||||||
sessions. BM25-only baseline (zero embeddings, so the vector stage is inert). See
|
sessions. See [BENCHMARKS.md § LongMemEval
|
||||||
[BENCHMARKS.md § LongMemEval Results](BENCHMARKS.md#longmemeval-results) for the
|
Results](BENCHMARKS.md#longmemeval-results) for the full scoring-target
|
||||||
full scoring-target declaration:
|
declaration:
|
||||||
|
|
||||||
| Metric | Turn-Level | Session-Level |
|
| Mode | Turn-Level Hit@5 | Session-Level Hit@5 |
|
||||||
|--------|------------|---------------|
|
|------|------------------|---------------------|
|
||||||
| Hit@5 | **75.0%** | **93.6%** |
|
| BM25 only | 75.0% | 93.6% |
|
||||||
| MRR | 0.6320 | 0.8948 |
|
| Vector only (MiniLM) | 71.8% | 94.2% |
|
||||||
|
| Hybrid (0.7/0.3) | **79.2%** | **95.8%** |
|
||||||
|
|
||||||
|
Hybrid is the strongest configuration at Hit@5 and Hit@10, which is what running
|
||||||
|
two retrieval stages is for. It is *not* strongest at rank 1 — BM25 alone leads
|
||||||
|
turn-level Hit@1 (53.8% vs 44.4%) and MRR (0.6320 vs 0.5868), so a caller taking
|
||||||
|
a single top hit is better served by BM25 today. The `0.7/0.3` weights are the
|
||||||
|
documented default, not a searched optimum.
|
||||||
|
|
||||||
|
Vector embeddings require `--features embeddings`; without it the vector stage is
|
||||||
|
inert and only the BM25 row is produced, which is what every previously published
|
||||||
|
number here measured.
|
||||||
|
|
||||||
On the easier `longmemeval_oracle` variant (evidence sessions only) the same
|
On the easier `longmemeval_oracle` variant (evidence sessions only) the same
|
||||||
harness scores 84.4% turn-level Hit@5 / MRR 0.6597, reproduced identically on a
|
harness scores 84.4% turn-level Hit@5 / MRR 0.6597, reproduced identically on a
|
||||||
|
|||||||
@@ -57,6 +57,13 @@ tempfile = { workspace = true }
|
|||||||
# Enable with: cargo bench -p clawhdf5-bench --features libhdf5-compare
|
# Enable with: cargo bench -p clawhdf5-bench --features libhdf5-compare
|
||||||
# Uses hdf5-metno (fork of hdf5 crate) which supports HDF5 1.14.x.
|
# Uses hdf5-metno (fork of hdf5 crate) which supports HDF5 1.14.x.
|
||||||
hdf5 = { version = "0.12", optional = true, package = "hdf5-metno" }
|
hdf5 = { version = "0.12", optional = true, package = "hdf5-metno" }
|
||||||
|
# Optional: real sentence embeddings for the LongMemEval bench's vector stage.
|
||||||
|
# Enable with: cargo run --release --bin longmemeval_bench --features embeddings
|
||||||
|
# Off by default — nothing in the shipped crates depends on these.
|
||||||
|
candle-core = { version = "0.9", optional = true }
|
||||||
|
candle-nn = { version = "0.9", optional = true }
|
||||||
|
candle-transformers = { version = "0.9", optional = true }
|
||||||
|
tokenizers = { version = "0.21", optional = true }
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
clawhdf5 = { path = "../clawhdf5", features = ["zstd", "pcodec"] }
|
clawhdf5 = { path = "../clawhdf5", features = ["zstd", "pcodec"] }
|
||||||
@@ -66,3 +73,8 @@ criterion = { workspace = true }
|
|||||||
# When enabled, benchmarks add matching libhdf5 variants for side-by-side comparison.
|
# When enabled, benchmarks add matching libhdf5 variants for side-by-side comparison.
|
||||||
libhdf5-compare = ["hdf5"]
|
libhdf5-compare = ["hdf5"]
|
||||||
mpi-io = ["clawhdf5-io/mpi-io", "mpi"]
|
mpi-io = ["clawhdf5-io/mpi-io", "mpi"]
|
||||||
|
# Real MiniLM embeddings for longmemeval_bench, so the vector stage is not inert.
|
||||||
|
embeddings = ["candle-core", "candle-nn", "candle-transformers", "tokenizers"]
|
||||||
|
# CUDA-accelerated embedding. MiniLM on a CPU takes hours over the full
|
||||||
|
# longmemeval_s haystack; on a GPU it is minutes.
|
||||||
|
embeddings-cuda = ["embeddings", "candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
|
||||||
|
|||||||
@@ -49,12 +49,56 @@
|
|||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
use std::time::{Duration, Instant};
|
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 clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
use tempfile::TempDir;
|
use tempfile::TempDir;
|
||||||
|
|
||||||
const EMBEDDING_DIM: usize = 384;
|
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
|
// JSON data types
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
@@ -196,7 +240,12 @@ struct EvalResult {
|
|||||||
latency: Duration,
|
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 dir = TempDir::new().expect("failed to create temp dir");
|
||||||
let mut config = MemoryConfig::new(dir.path().join("lme.h5"), "lme-bench", EMBEDDING_DIM);
|
let mut config = MemoryConfig::new(dir.path().join("lme.h5"), "lme-bench", EMBEDDING_DIM);
|
||||||
config.wal_enabled = false;
|
config.wal_enabled = false;
|
||||||
@@ -218,7 +267,7 @@ fn evaluate_question(q: &Question, top_k: usize) -> EvalResult {
|
|||||||
for turn in session {
|
for turn in session {
|
||||||
entries.push(MemoryEntry {
|
entries.push(MemoryEntry {
|
||||||
chunk: turn.content.clone(),
|
chunk: turn.content.clone(),
|
||||||
embedding: vec![0.0f32; EMBEDDING_DIM],
|
embedding: embedding_for(embeddings, &turn.content),
|
||||||
source_channel: "longmemeval".to_string(),
|
source_channel: "longmemeval".to_string(),
|
||||||
timestamp: ts,
|
timestamp: ts,
|
||||||
session_id: sess_id.to_string(),
|
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
|
// Set of session IDs that contain the answer
|
||||||
let answer_sess_set: HashSet<&str> = q.answer_session_ids.iter().map(String::as_str).collect();
|
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 query_emb = embedding_for(embeddings, &q.question);
|
||||||
let zero_emb = vec![0.0f32; EMBEDDING_DIM];
|
|
||||||
let t0 = Instant::now();
|
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();
|
let latency = t0.elapsed();
|
||||||
|
|
||||||
// Session-level recall
|
// 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!("=================================================================");
|
||||||
println!(" LongMemEval Benchmark (BM25-only retrieval, zero embeddings)");
|
println!(" LongMemEval Benchmark — {}", mode.label);
|
||||||
println!("=================================================================");
|
println!("=================================================================");
|
||||||
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!();
|
||||||
println!("Scoring target: RETRIEVAL RECALL (did the gold memory land in top-k).");
|
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!(" 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!("```json");
|
||||||
println!("{{");
|
println!("{{");
|
||||||
println!(" \"benchmark\": \"longmemeval\",");
|
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!(" \"dataset_variant\": \"{}\",", profile.variant());
|
||||||
println!(" \"scoring_target\": \"retrieval_recall\",");
|
println!(" \"scoring_target\": \"retrieval_recall\",");
|
||||||
println!(" \"k\": 10,");
|
println!(" \"k\": 10,");
|
||||||
@@ -585,6 +650,7 @@ fn print_report(overall: &Metrics, by_type: &HashMap<String, Metrics>, profile:
|
|||||||
fn main() {
|
fn main() {
|
||||||
let mut json_path: Option<String> = None;
|
let mut json_path: Option<String> = None;
|
||||||
let mut limit: Option<usize> = None;
|
let mut limit: Option<usize> = None;
|
||||||
|
let mut weights_dir: Option<String> = None;
|
||||||
let mut args = std::env::args().skip(1);
|
let mut args = std::env::args().skip(1);
|
||||||
while let Some(arg) = args.next() {
|
while let Some(arg) = args.next() {
|
||||||
match arg.as_str() {
|
match arg.as_str() {
|
||||||
@@ -592,6 +658,9 @@ fn main() {
|
|||||||
let v = args.next().expect("--limit needs a value");
|
let v = args.next().expect("--limit needs a value");
|
||||||
limit = Some(v.parse().expect("--limit must be a positive integer"));
|
limit = Some(v.parse().expect("--limit must be a positive integer"));
|
||||||
}
|
}
|
||||||
|
"--embeddings" => {
|
||||||
|
weights_dir = Some(args.next().expect("--embeddings needs a directory"));
|
||||||
|
}
|
||||||
"--help" | "-h" => {
|
"--help" | "-h" => {
|
||||||
eprintln!(
|
eprintln!(
|
||||||
"usage: longmemeval_bench [PATH] [--limit N]\n\n\
|
"usage: longmemeval_bench [PATH] [--limit N]\n\n\
|
||||||
@@ -600,7 +669,13 @@ fn main() {
|
|||||||
variant it was given rather than trusting the filename.\n\
|
variant it was given rather than trusting the filename.\n\
|
||||||
--limit evaluate N questions, sampled evenly across the file\n\
|
--limit evaluate N questions, sampled evenly across the file\n\
|
||||||
rather than as a prefix — the dataset is ordered by\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;
|
return;
|
||||||
}
|
}
|
||||||
@@ -650,6 +725,63 @@ fn main() {
|
|||||||
profile.evidence_density * 100.0,
|
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 overall = Metrics::default();
|
||||||
let mut by_type: HashMap<String, Metrics> = HashMap::new();
|
let mut by_type: HashMap<String, Metrics> = HashMap::new();
|
||||||
|
|
||||||
@@ -658,7 +790,7 @@ fn main() {
|
|||||||
eprint!("\r [{}/{}] evaluating...", i + 1, total);
|
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 is_abs = q.question_type.ends_with("_abs");
|
||||||
let base_type = if is_abs {
|
let base_type = if is_abs {
|
||||||
@@ -723,5 +855,5 @@ fn main() {
|
|||||||
|
|
||||||
eprintln!("\r [{total}/{total}] done. ");
|
eprintln!("\r [{total}/{total}] done. ");
|
||||||
eprintln!();
|
eprintln!();
|
||||||
print_report(&overall, &by_type, &profile);
|
print_report(&overall, &by_type, profile, mode);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,164 @@
|
|||||||
|
//! Optional MiniLM sentence embedder for the LongMemEval bench.
|
||||||
|
//!
|
||||||
|
//! Compiled only under the `embeddings` feature, so the default build of a
|
||||||
|
//! project that prides itself on having no heavyweight dependencies stays
|
||||||
|
//! exactly as it was. Without it the bench runs BM25-only, as it always has.
|
||||||
|
//!
|
||||||
|
//! Loads `sentence-transformers/all-MiniLM-L6-v2` — the same checkpoint
|
||||||
|
//! omni-cortex uses — and produces 384-d mean-pooled, L2-normalised sentence
|
||||||
|
//! embeddings, which is the published recipe for this model (mean over token
|
||||||
|
//! states weighted by the attention mask, *not* the `[CLS]` pooler output).
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::path::Path;
|
||||||
|
|
||||||
|
use candle_core::{DType, Device, Tensor};
|
||||||
|
use candle_nn::VarBuilder;
|
||||||
|
use candle_transformers::models::bert::{BertModel, Config, HiddenAct};
|
||||||
|
use tokenizers::Tokenizer;
|
||||||
|
|
||||||
|
/// Sequences encoded per forward pass. Larger batches amortise the transformer
|
||||||
|
/// call; 64 keeps peak memory modest while still saturating a CPU.
|
||||||
|
const BATCH: usize = 64;
|
||||||
|
|
||||||
|
/// A loaded MiniLM encoder.
|
||||||
|
pub struct Embedder {
|
||||||
|
model: BertModel,
|
||||||
|
tokenizer: Tokenizer,
|
||||||
|
device: Device,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Embedder {
|
||||||
|
/// Load from a directory holding `model.safetensors` and `tokenizer.json`.
|
||||||
|
///
|
||||||
|
/// `config.json` is read when present; otherwise the published MiniLM-L6-v2
|
||||||
|
/// architecture constants are used, which are pinned rather than guessed.
|
||||||
|
pub fn load(dir: &Path) -> Result<Self, Box<dyn std::error::Error>> {
|
||||||
|
// CUDA when the feature is on and a device is actually present; the CPU
|
||||||
|
// path is correct but roughly two orders of magnitude slower, which is
|
||||||
|
// the difference between minutes and most of a day on the full haystack.
|
||||||
|
let device = match Device::new_cuda(0) {
|
||||||
|
Ok(d) => {
|
||||||
|
eprintln!("Embedder: CUDA device 0");
|
||||||
|
d
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
eprintln!("Embedder: CPU ({e})");
|
||||||
|
Device::Cpu
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let weights = dir.join("model.safetensors");
|
||||||
|
let tok_path = dir.join("tokenizer.json");
|
||||||
|
|
||||||
|
let config: Config = match std::fs::read_to_string(dir.join("config.json")) {
|
||||||
|
Ok(raw) => serde_json::from_str(&raw)?,
|
||||||
|
Err(_) => Config {
|
||||||
|
vocab_size: 30_522,
|
||||||
|
hidden_size: 384,
|
||||||
|
num_hidden_layers: 6,
|
||||||
|
num_attention_heads: 12,
|
||||||
|
intermediate_size: 1_536,
|
||||||
|
hidden_act: HiddenAct::Gelu,
|
||||||
|
hidden_dropout_prob: 0.0,
|
||||||
|
max_position_embeddings: 512,
|
||||||
|
type_vocab_size: 2,
|
||||||
|
initializer_range: 0.02,
|
||||||
|
layer_norm_eps: 1e-12,
|
||||||
|
pad_token_id: 0,
|
||||||
|
position_embedding_type: Default::default(),
|
||||||
|
use_cache: false,
|
||||||
|
classifier_dropout: None,
|
||||||
|
model_type: None,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights], DType::F32, &device)? };
|
||||||
|
let model = BertModel::load(vb, &config)?;
|
||||||
|
let tokenizer = Tokenizer::from_file(&tok_path).map_err(|e| e.to_string())?;
|
||||||
|
|
||||||
|
Ok(Self {
|
||||||
|
model,
|
||||||
|
tokenizer,
|
||||||
|
device,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Encode `texts` into 384-d unit vectors, in order.
|
||||||
|
fn encode_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, Box<dyn std::error::Error>> {
|
||||||
|
let mut tk = self.tokenizer.clone();
|
||||||
|
let tk = tk
|
||||||
|
.with_padding(Some(tokenizers::PaddingParams::default()))
|
||||||
|
.with_truncation(Some(tokenizers::TruncationParams {
|
||||||
|
max_length: 512,
|
||||||
|
..Default::default()
|
||||||
|
}))
|
||||||
|
.map_err(|e| e.to_string())?;
|
||||||
|
let encodings = tk
|
||||||
|
.encode_batch(texts.to_vec(), true)
|
||||||
|
.map_err(|e| e.to_string())?;
|
||||||
|
|
||||||
|
let ids: Vec<u32> = encodings
|
||||||
|
.iter()
|
||||||
|
.flat_map(|e| e.get_ids().to_vec())
|
||||||
|
.collect();
|
||||||
|
let mask: Vec<u32> = encodings
|
||||||
|
.iter()
|
||||||
|
.flat_map(|e| e.get_attention_mask().to_vec())
|
||||||
|
.collect();
|
||||||
|
let (b, l) = (encodings.len(), encodings[0].get_ids().len());
|
||||||
|
|
||||||
|
let ids = Tensor::from_vec(ids, (b, l), &self.device)?;
|
||||||
|
let mask = Tensor::from_vec(mask, (b, l), &self.device)?;
|
||||||
|
let type_ids = ids.zeros_like()?;
|
||||||
|
|
||||||
|
let hidden = self.model.forward(&ids, &type_ids, Some(&mask))?;
|
||||||
|
|
||||||
|
// Mean-pool over real tokens only: sum(hidden * mask) / sum(mask).
|
||||||
|
let mask_f = mask.to_dtype(DType::F32)?.unsqueeze(2)?;
|
||||||
|
let summed = hidden.broadcast_mul(&mask_f)?.sum(1)?;
|
||||||
|
let counts = mask_f.sum(1)?.clamp(1e-9, f32::INFINITY)?;
|
||||||
|
let pooled = summed.broadcast_div(&counts)?;
|
||||||
|
|
||||||
|
// L2-normalise so cosine similarity is a plain dot product.
|
||||||
|
let norm = pooled
|
||||||
|
.sqr()?
|
||||||
|
.sum_keepdim(1)?
|
||||||
|
.sqrt()?
|
||||||
|
.clamp(1e-12, f32::INFINITY)?;
|
||||||
|
let normed = pooled.broadcast_div(&norm)?;
|
||||||
|
|
||||||
|
Ok(normed.to_vec2::<f32>()?)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Encode every distinct string in `texts` once, returning a lookup map.
|
||||||
|
///
|
||||||
|
/// LongMemEval's haystack sessions are drawn from a shared pool, so the same
|
||||||
|
/// turn text recurs across many questions. Deduplicating before encoding is
|
||||||
|
/// the difference between encoding the corpus once and encoding it per
|
||||||
|
/// question.
|
||||||
|
pub fn encode_unique(
|
||||||
|
&self,
|
||||||
|
texts: impl IntoIterator<Item = String>,
|
||||||
|
) -> Result<HashMap<String, Vec<f32>>, Box<dyn std::error::Error>> {
|
||||||
|
let mut unique: Vec<String> = texts.into_iter().collect();
|
||||||
|
unique.sort_unstable();
|
||||||
|
unique.dedup();
|
||||||
|
|
||||||
|
let total = unique.len();
|
||||||
|
eprintln!("Embedding {total} unique texts with MiniLM (batch {BATCH})...");
|
||||||
|
|
||||||
|
let mut out = HashMap::with_capacity(total);
|
||||||
|
for (n, chunk) in unique.chunks(BATCH).enumerate() {
|
||||||
|
let refs: Vec<&str> = chunk.iter().map(String::as_str).collect();
|
||||||
|
let vecs = self.encode_batch(&refs)?;
|
||||||
|
for (text, v) in chunk.iter().zip(vecs) {
|
||||||
|
out.insert(text.clone(), v);
|
||||||
|
}
|
||||||
|
if n % 50 == 0 {
|
||||||
|
eprint!("\r [{}/{}] embedded...", (n * BATCH).min(total), total);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
eprintln!("\r [{total}/{total}] embedded. ");
|
||||||
|
Ok(out)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user