Merge feat/retrieval-quality: tuned fusion defaults, RRF measured, query-expansion fixes
CI / test (push) Failing after 2s
CI / test (push) Failing after 2s
Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
+30
-1
@@ -597,6 +597,34 @@ Session-level:
|
|||||||
| Vector only | 85.4% | 94.2% | 96.6% | 0.8901 |
|
| Vector only | 85.4% | 94.2% | 96.6% | 0.8901 |
|
||||||
| Hybrid | **88.2%** | **95.8%** | **97.8%** | **0.9158** |
|
| Hybrid | **88.2%** | **95.8%** | **97.8%** | **0.9158** |
|
||||||
|
|
||||||
|
### Fusion method — weighted vs. RRF, full haystack, n=500
|
||||||
|
|
||||||
|
Reciprocal rank fusion has been in the codebase since early on but was only
|
||||||
|
reachable as a free function over a linear scan, so it had never been compared
|
||||||
|
with the weighted sum on equal terms. `HDF5Memory::hybrid_search_with` now
|
||||||
|
takes a `Fusion`, and both run over the same HNSW + BM25 candidates:
|
||||||
|
|
||||||
|
| Mode | turn Hit@1 | Hit@5 | Hit@10 | MRR | session Hit@1 | session MRR |
|
||||||
|
|---|---|---|---|---|---|---|
|
||||||
|
| BM25 only | **53.8%** | 75.0% | 81.6% | 0.6320 | 86.2% | 0.8948 |
|
||||||
|
| Vector only | 36.0% | 71.8% | 81.6% | 0.5031 | 85.4% | 0.8901 |
|
||||||
|
| **Weighted 0.4 / 0.6** | 51.6% | **81.4%** | **87.8%** | **0.6430** | **91.0%** | **0.9347** |
|
||||||
|
| RRF (k=60) | 45.0% | 78.8% | 87.6% | 0.5967 | 89.6% | 0.9253 |
|
||||||
|
|
||||||
|
**RRF loses to the tuned weighted sum** — 6.6pp of turn Hit@1 and 0.046 of MRR
|
||||||
|
— and lands almost exactly where the old `0.7/0.3` weighting did (44.2% /
|
||||||
|
0.5856). That is not a coincidence: RRF combines the two stages by rank with
|
||||||
|
*equal* influence, and on this corpus the stages are not equally good. BM25
|
||||||
|
alone beats the vector stage by 17.8pp at Hit@1, so any scheme that treats them
|
||||||
|
as peers gives up rank-1 accuracy, and RRF discards the score magnitudes that
|
||||||
|
would say which stage to believe.
|
||||||
|
|
||||||
|
This is a property of the corpus, not a defect in RRF: its selling point is
|
||||||
|
robustness when the two stages' scores are not comparable and there is no
|
||||||
|
labelled data to tune against. Here there is, so the weighted sum is kept as
|
||||||
|
the default. `Fusion::Rrf` remains available for callers whose stages are more
|
||||||
|
evenly matched.
|
||||||
|
|
||||||
### Weight sweep — full haystack, n=500
|
### Weight sweep — full haystack, n=500
|
||||||
|
|
||||||
`0.7/0.3` was a documented default, never a searched one. Sweeping
|
`0.7/0.3` was a documented default, never a searched one. Sweeping
|
||||||
@@ -639,7 +667,8 @@ BM25 at Hit@1. Both dominate `0.7/0.3`.
|
|||||||
|
|
||||||
The rows below are kept at the three original settings because they are what the
|
The rows below are kept at the three original settings because they are what the
|
||||||
mode ablation measured — read them as "the shape of each stage in isolation",
|
mode ablation measured — read them as "the shape of each stage in isolation",
|
||||||
and take the operating point from the sweep.
|
and take the operating point from the sweep. `0.4/0.6` is now the shipped
|
||||||
|
default (`hybrid::DEFAULT_FUSION`).
|
||||||
|
|
||||||
The same pattern shows up independently in omni-cortex's four-signal RRF ablation,
|
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.
|
where adding BM25 to a dense retriever raised nDCG@5 while lowering Hit@1 and MRR.
|
||||||
|
|||||||
@@ -2,6 +2,39 @@
|
|||||||
|
|
||||||
## Unreleased
|
## Unreleased
|
||||||
|
|
||||||
|
### Retrieval quality
|
||||||
|
- `clawhdf5-agent`: **`QueryExpander::expand` panicked on ordinary non-ASCII
|
||||||
|
input** — `"İ AI"` was enough. It searched a lowercased copy of the query and
|
||||||
|
then sliced the *original* with those offsets, which only works while
|
||||||
|
lowercasing preserves byte length (Turkish `İ` is 2 bytes and lowercases to
|
||||||
|
3). Depending on where the offsets drifted it either corrupted the output
|
||||||
|
("İstanbul AI trip" lost a character) or panicked. Matching now walks the
|
||||||
|
original string.
|
||||||
|
- `clawhdf5-agent`: query expansion no longer rewrites text inside words.
|
||||||
|
`replace_word_case_insensitive` did a plain substring replace despite its
|
||||||
|
name, so "training" became "trArtificial Intelligencening" and "programming"
|
||||||
|
became "Pull Requestogramming" — every acronym expansion of ordinary prose
|
||||||
|
was corrupt. Matches now require word boundaries; genuine acronyms
|
||||||
|
(`API`, `database`) still expand.
|
||||||
|
- `clawhdf5-agent`: **the default fusion weights are now the measured ones.**
|
||||||
|
A sweep of every 0.1 step over the full LongMemEval haystack (500 questions,
|
||||||
|
real MiniLM embeddings) shows the long-standing `0.7/0.3` default is
|
||||||
|
*strictly dominated* by `0.4/0.6` — turn-level Hit@1 51.6% vs 44.2%, Hit@5
|
||||||
|
81.4% vs 79.2%, Hit@10 87.8% vs 85.8%, MRR 0.643 vs 0.586, and better at
|
||||||
|
session level too. The finding was recorded in `BENCHMARKS.md` but had never
|
||||||
|
been applied: `unified_search` and the OpenClaw backend both hardcoded
|
||||||
|
`0.7/0.3`. They now use `hybrid::DEFAULT_FUSION`. **Callers passing weights
|
||||||
|
to `hybrid_search` explicitly are unaffected** — pass `0.4`/`0.6` (or use
|
||||||
|
`hybrid_search_with`) to get the tuned behaviour.
|
||||||
|
- `clawhdf5-agent`: fusion is now selectable. New `hybrid::Fusion`
|
||||||
|
(`Weighted { vector, keyword }` or `Rrf { k }`), `hybrid::fuse`,
|
||||||
|
`hybrid::hybrid_search_fused` and `HDF5Memory::hybrid_search_with`.
|
||||||
|
Reciprocal rank fusion existed but was unreachable from the store, so it had
|
||||||
|
never been measured against the weighted sum; the LongMemEval bench now has
|
||||||
|
an `RRF` mode.
|
||||||
|
|
||||||
|
## Unreleased
|
||||||
|
|
||||||
### HDF5 Read Path
|
### HDF5 Read Path
|
||||||
- **Selection reads cost what the selection costs.** `read_*_selection` decoded
|
- **Selection reads cost what the selection costs.** `read_*_selection` decoded
|
||||||
the *entire* dataset and then picked elements out, so a 64 x 64 window of a
|
the *entire* dataset and then picked elements out, so a 64 x 64 window of a
|
||||||
|
|||||||
@@ -37,7 +37,7 @@
|
|||||||
//! let mem = AsyncHDF5Memory::open_with(path, config).await?;
|
//! let mem = AsyncHDF5Memory::open_with(path, config).await?;
|
||||||
//! mem.save(entry).await?; // buffered → background writer
|
//! mem.save(entry).await?; // buffered → background writer
|
||||||
//! mem.save_batch(entries).await?; // also buffered
|
//! mem.save_batch(entries).await?; // also buffered
|
||||||
//! let results = mem.hybrid_search(emb, "query".into(), 0.7, 0.3, 5).await;
|
//! let results = mem.hybrid_search(emb, "query".into(), 0.4, 0.6, 5).await;
|
||||||
//! mem.shutdown().await?; // final flush + stop
|
//! mem.shutdown().await?; // final flush + stop
|
||||||
//! ```
|
//! ```
|
||||||
|
|
||||||
|
|||||||
@@ -29,12 +29,39 @@ pub fn hybrid_search(
|
|||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
vectors: &[Vec<f32>],
|
vectors: &[Vec<f32>],
|
||||||
_chunks: &[String],
|
chunks: &[String],
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
bm25_index: &BM25Index,
|
bm25_index: &BM25Index,
|
||||||
vector_weight: f32,
|
vector_weight: f32,
|
||||||
keyword_weight: f32,
|
keyword_weight: f32,
|
||||||
k: usize,
|
k: usize,
|
||||||
|
) -> Vec<(usize, f32)> {
|
||||||
|
hybrid_search_fused(
|
||||||
|
query_embedding,
|
||||||
|
query_text,
|
||||||
|
vectors,
|
||||||
|
chunks,
|
||||||
|
tombstones,
|
||||||
|
bm25_index,
|
||||||
|
Fusion::Weighted {
|
||||||
|
vector: vector_weight,
|
||||||
|
keyword: keyword_weight,
|
||||||
|
},
|
||||||
|
k,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// [`hybrid_search`] with the fusion method chosen explicitly.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
pub fn hybrid_search_fused(
|
||||||
|
query_embedding: &[f32],
|
||||||
|
query_text: &str,
|
||||||
|
vectors: &[Vec<f32>],
|
||||||
|
_chunks: &[String],
|
||||||
|
tombstones: &[u8],
|
||||||
|
bm25_index: &BM25Index,
|
||||||
|
fusion: Fusion,
|
||||||
|
k: usize,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
// Get raw scores from both systems. Request all results so normalization
|
// Get raw scores from both systems. Request all results so normalization
|
||||||
// covers the full distribution.
|
// covers the full distribution.
|
||||||
@@ -60,7 +87,7 @@ pub fn hybrid_search(
|
|||||||
};
|
};
|
||||||
let kw_scores = bm25_index.scores(query_text);
|
let kw_scores = bm25_index.scores(query_text);
|
||||||
|
|
||||||
merge_vector_keyword(vec_scores, kw_scores, vector_weight, keyword_weight, k)
|
fuse(vec_scores, kw_scores, fusion, k)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
|
/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
|
||||||
@@ -76,18 +103,92 @@ pub fn merge_vector_keyword(
|
|||||||
keyword_weight: f32,
|
keyword_weight: f32,
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
// Normalize each set to [0, 1].
|
fuse(
|
||||||
let vec_normalized = normalize_scores(&vec_scores);
|
vec_scores,
|
||||||
let kw_normalized = normalize_scores(&kw_scores);
|
kw_scores,
|
||||||
|
Fusion::Weighted {
|
||||||
|
vector: vector_weight,
|
||||||
|
keyword: keyword_weight,
|
||||||
|
},
|
||||||
|
k,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
// Merge scores with weights.
|
/// How the vector and keyword stages are combined into one ranking.
|
||||||
let mut merged: HashMap<usize, f32> = HashMap::new();
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||||
|
pub enum Fusion {
|
||||||
|
/// Min-max normalise each stage over its own candidates, then take a
|
||||||
|
/// weighted sum. Uses the *scores*, so a stage that separates its
|
||||||
|
/// candidates sharply keeps that separation — and a stage whose candidates
|
||||||
|
/// are all near-identical contributes little.
|
||||||
|
Weighted {
|
||||||
|
/// Weight on the vector stage.
|
||||||
|
vector: f32,
|
||||||
|
/// Weight on the keyword stage.
|
||||||
|
keyword: f32,
|
||||||
|
},
|
||||||
|
/// Reciprocal rank fusion: each stage contributes `1 / (k + rank)`,
|
||||||
|
/// ignoring score magnitudes entirely. Robust when the two stages'
|
||||||
|
/// scores aren't comparable, at the cost of discarding confidence.
|
||||||
|
Rrf {
|
||||||
|
/// The rank-damping constant; 60 is the value from the original paper.
|
||||||
|
k: f32,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
for (idx, score) in &vec_normalized {
|
impl Default for Fusion {
|
||||||
*merged.entry(*idx).or_insert(0.0) += vector_weight * score;
|
fn default() -> Self {
|
||||||
|
DEFAULT_FUSION
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The fusion `hybrid_search` uses unless told otherwise.
|
||||||
|
///
|
||||||
|
/// The weights are not a guess: a sweep of every 0.1 step over the full
|
||||||
|
/// LongMemEval haystack (500 questions, real MiniLM embeddings) found the
|
||||||
|
/// long-standing 0.7/0.3 default *strictly dominated* — 0.4/0.6 is better at
|
||||||
|
/// Hit@1, Hit@5, Hit@10 and MRR, at both turn and session granularity. See
|
||||||
|
/// `BENCHMARKS.md`, "Weight sweep".
|
||||||
|
pub const DEFAULT_FUSION: Fusion = Fusion::Weighted {
|
||||||
|
vector: 0.4,
|
||||||
|
keyword: 0.6,
|
||||||
|
};
|
||||||
|
|
||||||
|
/// Combine one ranked candidate list from each stage into a single top-`k`.
|
||||||
|
///
|
||||||
|
/// Neither list need be sorted; both are consumed.
|
||||||
|
pub fn fuse(
|
||||||
|
vec_scores: Vec<(usize, f32)>,
|
||||||
|
kw_scores: Vec<(usize, f32)>,
|
||||||
|
fusion: Fusion,
|
||||||
|
k: usize,
|
||||||
|
) -> Vec<(usize, f32)> {
|
||||||
|
let mut merged: HashMap<usize, f32> = HashMap::new();
|
||||||
|
match fusion {
|
||||||
|
Fusion::Weighted { vector, keyword } => {
|
||||||
|
// Normalize each set to [0, 1].
|
||||||
|
for (idx, score) in &normalize_scores(&vec_scores) {
|
||||||
|
*merged.entry(*idx).or_insert(0.0) += vector * score;
|
||||||
|
}
|
||||||
|
for (idx, score) in &normalize_scores(&kw_scores) {
|
||||||
|
*merged.entry(*idx).or_insert(0.0) += keyword * score;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Fusion::Rrf { k: damping } => {
|
||||||
|
for mut stage in [vec_scores, kw_scores] {
|
||||||
|
// Rank 1 is the best score. Ties break by index so a stage's
|
||||||
|
// contribution doesn't depend on the candidate order it
|
||||||
|
// happened to be produced in.
|
||||||
|
stage.sort_by(|a, b| {
|
||||||
|
b.1.partial_cmp(&a.1)
|
||||||
|
.unwrap_or(std::cmp::Ordering::Equal)
|
||||||
|
.then(a.0.cmp(&b.0))
|
||||||
|
});
|
||||||
|
for (rank, (idx, _)) in stage.iter().enumerate() {
|
||||||
|
*merged.entry(*idx).or_insert(0.0) += 1.0 / (damping + (rank + 1) as f32);
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
for (idx, score) in &kw_normalized {
|
|
||||||
*merged.entry(*idx).or_insert(0.0) += keyword_weight * score;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut results: Vec<(usize, f32)> = merged.into_iter().collect();
|
let mut results: Vec<(usize, f32)> = merged.into_iter().collect();
|
||||||
@@ -353,6 +454,49 @@ mod tests {
|
|||||||
assert_eq!(result[0].1, 1.0);
|
assert_eq!(result[0].1, 1.0);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn default_fusion_is_the_tuned_operating_point() {
|
||||||
|
// A sweep over the full LongMemEval haystack found 0.7/0.3 strictly
|
||||||
|
// dominated by 0.4/0.6 (BENCHMARKS.md). This guards the finding
|
||||||
|
// against being quietly undone.
|
||||||
|
assert_eq!(
|
||||||
|
DEFAULT_FUSION,
|
||||||
|
Fusion::Weighted {
|
||||||
|
vector: 0.4,
|
||||||
|
keyword: 0.6
|
||||||
|
}
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn rrf_rewards_agreement_between_the_stages_and_ignores_magnitudes() {
|
||||||
|
// Doc 1 is second-best in both stages; doc 0 is best in one and absent
|
||||||
|
// from the other. RRF prefers the doc both stages liked.
|
||||||
|
let vec_scores = vec![(0, 100.0), (1, 0.9)];
|
||||||
|
let kw_scores = vec![(2, 5.0), (1, 4.9)];
|
||||||
|
let ranked = fuse(vec_scores, kw_scores, Fusion::Rrf { k: 60.0 }, 3);
|
||||||
|
assert_eq!(ranked[0].0, 1, "{ranked:?}");
|
||||||
|
|
||||||
|
// Scaling one stage's scores cannot change an RRF ranking, only the
|
||||||
|
// order within that stage can.
|
||||||
|
let a = fuse(
|
||||||
|
vec![(0, 1.0), (1, 0.5)],
|
||||||
|
vec![(1, 2.0), (0, 1.0)],
|
||||||
|
Fusion::Rrf { k: 60.0 },
|
||||||
|
2,
|
||||||
|
);
|
||||||
|
let b = fuse(
|
||||||
|
vec![(0, 1e6), (1, -3.0)],
|
||||||
|
vec![(1, 0.002), (0, 0.001)],
|
||||||
|
Fusion::Rrf { k: 60.0 },
|
||||||
|
2,
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
a.iter().map(|r| r.0).collect::<Vec<_>>(),
|
||||||
|
b.iter().map(|r| r.0).collect::<Vec<_>>()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn merge_top_k_matches_a_full_sort() {
|
fn merge_top_k_matches_a_full_sort() {
|
||||||
// Many ties (scores repeat) so the index tie-break is exercised.
|
// Many ties (scores repeat) so the index tie-break is exercised.
|
||||||
|
|||||||
@@ -1350,7 +1350,8 @@ impl HDF5Memory {
|
|||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<SearchResult> {
|
) -> Vec<SearchResult> {
|
||||||
// Persistent tier.
|
// Persistent tier.
|
||||||
let persistent = self.hybrid_search(query_embedding, query_text, 0.7, 0.3, k);
|
let persistent =
|
||||||
|
self.hybrid_search_with(query_embedding, query_text, hybrid::DEFAULT_FUSION, k);
|
||||||
const EPHEMERAL_BOOST: f32 = 1.2;
|
const EPHEMERAL_BOOST: f32 = 1.2;
|
||||||
let mut results = persistent;
|
let mut results = persistent;
|
||||||
|
|
||||||
|
|||||||
@@ -531,11 +531,14 @@ impl MemoryBackend for ClawhdfBackend {
|
|||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<MemorySearchResult> {
|
) -> Vec<MemorySearchResult> {
|
||||||
// 1. Hybrid retrieval (RRF-blended vector + BM25).
|
// 1. Hybrid retrieval (vector + BM25, fused by score).
|
||||||
let candidates = k.saturating_mul(3).max(10);
|
let candidates = k.saturating_mul(3).max(10);
|
||||||
let raw = self
|
let raw = self.memory.hybrid_search_with(
|
||||||
.memory
|
query_embedding,
|
||||||
.hybrid_search(query_embedding, query_text, 0.7, 0.3, candidates);
|
query_text,
|
||||||
|
crate::hybrid::DEFAULT_FUSION,
|
||||||
|
candidates,
|
||||||
|
);
|
||||||
|
|
||||||
if raw.is_empty() {
|
if raw.is_empty() {
|
||||||
return Vec::new();
|
return Vec::new();
|
||||||
|
|||||||
@@ -6,6 +6,11 @@
|
|||||||
//! - Temporal expansion (time-related rewrites)
|
//! - Temporal expansion (time-related rewrites)
|
||||||
//! - Morphological variants (stemming-like transforms)
|
//! - Morphological variants (stemming-like transforms)
|
||||||
//! - Knowledge graph expansion (entity aliases and neighbors)
|
//! - Knowledge graph expansion (entity aliases and neighbors)
|
||||||
|
//!
|
||||||
|
//! The morphological rules are crude suffix swaps, so some variants are not
|
||||||
|
//! words ("during" -> "dured"). That is tolerable for a BM25 stage, which
|
||||||
|
//! simply finds no postings for a nonsense term, but it means expansion is not
|
||||||
|
//! free: measure before enabling it on a retrieval path.
|
||||||
|
|
||||||
use crate::knowledge::KnowledgeCache;
|
use crate::knowledge::KnowledgeCache;
|
||||||
|
|
||||||
@@ -340,18 +345,85 @@ fn contains_phrase(text: &str, phrase: &str) -> bool {
|
|||||||
|
|
||||||
/// Replace a phrase in `text` case-insensitively, preserving surrounding case.
|
/// Replace a phrase in `text` case-insensitively, preserving surrounding case.
|
||||||
fn replace_word_case_insensitive(text: &str, from: &str, to: &str) -> String {
|
fn replace_word_case_insensitive(text: &str, from: &str, to: &str) -> String {
|
||||||
case_insensitive_replace(text, from, to)
|
replace_first(text, from, to, MatchKind::WholeWord)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn case_insensitive_replace(text: &str, from: &str, to: &str) -> String {
|
fn case_insensitive_replace(text: &str, from: &str, to: &str) -> String {
|
||||||
let lower = text.to_lowercase();
|
replace_first(text, from, to, MatchKind::Substring)
|
||||||
let lower_from = from.to_lowercase();
|
}
|
||||||
if let Some(pos) = lower.find(&lower_from) {
|
|
||||||
let end = pos + from.len();
|
/// Whether a match may fall inside a larger word.
|
||||||
format!("{}{}{}", &text[..pos], to, &text[end..])
|
#[derive(Clone, Copy, PartialEq)]
|
||||||
} else {
|
enum MatchKind {
|
||||||
text.to_string()
|
/// Match anywhere, including inside another word.
|
||||||
|
Substring,
|
||||||
|
/// Match only when both ends sit on a word boundary.
|
||||||
|
WholeWord,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Replace the first case-insensitive match of `from` in `text` with `to`.
|
||||||
|
///
|
||||||
|
/// Matching walks the *original* string rather than a lowercased copy. The
|
||||||
|
/// previous implementation searched `text.to_lowercase()` and then sliced
|
||||||
|
/// `text` with the offsets it found, which only holds while lowercasing
|
||||||
|
/// preserves byte length. It does not: Turkish `İ` (2 bytes) lowercases to
|
||||||
|
/// `i` + U+0307 (3 bytes), so every later offset was wrong — silently
|
||||||
|
/// corrupting the output, or panicking when an offset landed inside a
|
||||||
|
/// character or past the end. `"İ AI"` was enough to panic.
|
||||||
|
fn replace_first(text: &str, from: &str, to: &str, kind: MatchKind) -> String {
|
||||||
|
match find_case_insensitive(text, from, kind) {
|
||||||
|
Some((start, end)) => {
|
||||||
|
let mut out = String::with_capacity(text.len() - (end - start) + to.len());
|
||||||
|
out.push_str(&text[..start]);
|
||||||
|
out.push_str(to);
|
||||||
|
out.push_str(&text[end..]);
|
||||||
|
out
|
||||||
}
|
}
|
||||||
|
None => text.to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Byte range of the first case-insensitive match of `needle` in `haystack`.
|
||||||
|
fn find_case_insensitive(haystack: &str, needle: &str, kind: MatchKind) -> Option<(usize, usize)> {
|
||||||
|
if needle.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let lowered: Vec<char> = needle.chars().flat_map(char::to_lowercase).collect();
|
||||||
|
let is_word = |c: char| c.is_alphanumeric() || c == '_';
|
||||||
|
|
||||||
|
for (start, _) in haystack.char_indices() {
|
||||||
|
if kind == MatchKind::WholeWord
|
||||||
|
&& haystack[..start].chars().next_back().is_some_and(is_word)
|
||||||
|
{
|
||||||
|
continue; // mid-word: "ai" inside "training"
|
||||||
|
}
|
||||||
|
let mut matched = 0usize;
|
||||||
|
let mut end = start;
|
||||||
|
for (offset, ch) in haystack[start..].char_indices() {
|
||||||
|
if matched == lowered.len() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
let mut consumed_all = true;
|
||||||
|
for lc in ch.to_lowercase() {
|
||||||
|
if lowered.get(matched) != Some(&lc) {
|
||||||
|
consumed_all = false;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
matched += 1;
|
||||||
|
}
|
||||||
|
if !consumed_all {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
end = start + offset + ch.len_utf8();
|
||||||
|
}
|
||||||
|
if matched == lowered.len()
|
||||||
|
&& !(kind == MatchKind::WholeWord
|
||||||
|
&& haystack[end..].chars().next().is_some_and(is_word))
|
||||||
|
{
|
||||||
|
return Some((start, end));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Simple whitespace/punctuation tokenizer.
|
/// Simple whitespace/punctuation tokenizer.
|
||||||
@@ -637,4 +709,86 @@ mod tests {
|
|||||||
expanded.iter().map(|x| &x.text).collect::<Vec<_>>()
|
expanded.iter().map(|x| &x.text).collect::<Vec<_>>()
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
#[test]
|
||||||
|
fn acronyms_only_match_whole_words() {
|
||||||
|
let ex = QueryExpander::new(QueryExpansionConfig::default());
|
||||||
|
// "training" contains "ai", "programming" contains "pr". These used to
|
||||||
|
// be rewritten to "trArtificial Intelligencening" and
|
||||||
|
// "Pull Requestogramming".
|
||||||
|
for query in [
|
||||||
|
"How many miles during my marathon training?",
|
||||||
|
"Which programming language did I pick?",
|
||||||
|
"I updated the maintainer list",
|
||||||
|
] {
|
||||||
|
for expansion in ex.expand(query) {
|
||||||
|
assert!(
|
||||||
|
expansion.expansion_type != "acronym",
|
||||||
|
"{query:?} produced {expansion:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// A real acronym still expands, in both directions.
|
||||||
|
let texts: Vec<String> = ex
|
||||||
|
.expand("What about the API and the database?")
|
||||||
|
.into_iter()
|
||||||
|
.filter(|e| e.expansion_type == "acronym")
|
||||||
|
.map(|e| e.text)
|
||||||
|
.collect();
|
||||||
|
assert!(
|
||||||
|
texts
|
||||||
|
.iter()
|
||||||
|
.any(|t| t.contains("Application Programming Interface")),
|
||||||
|
"{texts:?}"
|
||||||
|
);
|
||||||
|
assert!(texts.iter().any(|t| t.contains("DB")), "{texts:?}");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn non_ascii_queries_do_not_panic_or_corrupt() {
|
||||||
|
let ex = QueryExpander::new(QueryExpansionConfig::default());
|
||||||
|
// Turkish 'İ' is 2 bytes but lowercases to 3, so offsets taken from a
|
||||||
|
// lowercased copy no longer line up with the original. `"İ AI"` used
|
||||||
|
// to panic; `"İstanbul AI trip"` used to silently eat a character.
|
||||||
|
for query in ["İ AI", "İé AI", "İİ ML", "İstanbul AI trip", "ǰ ML notes"] {
|
||||||
|
for expansion in ex.expand(query) {
|
||||||
|
assert!(
|
||||||
|
expansion.text.contains('İ') || expansion.text.contains('ǰ'),
|
||||||
|
"{query:?} lost its leading character: {expansion:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let expanded = ex.expand("İstanbul AI trip");
|
||||||
|
assert!(
|
||||||
|
expanded
|
||||||
|
.iter()
|
||||||
|
.any(|e| e.text == "İstanbul Artificial Intelligence trip"),
|
||||||
|
"{expanded:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn whole_word_matching_handles_string_edges_and_case() {
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("ai tools", "AI", "Artificial Intelligence"),
|
||||||
|
"Artificial Intelligence tools"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("tools for ai", "AI", "Artificial Intelligence"),
|
||||||
|
"tools for Artificial Intelligence"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("the aim", "AI", "Artificial Intelligence"),
|
||||||
|
"the aim",
|
||||||
|
"must not match inside a word"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("no match here", "xyz", "abc"),
|
||||||
|
"no match here"
|
||||||
|
);
|
||||||
|
// Only the first occurrence is replaced, as before.
|
||||||
|
assert_eq!(
|
||||||
|
replace_word_case_insensitive("ai and ai", "ai", "ML"),
|
||||||
|
"ML and ai"
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -20,8 +20,7 @@ impl HDF5Memory {
|
|||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
bm25: &bm25::BM25Index,
|
bm25: &bm25::BM25Index,
|
||||||
vector_weight: f32,
|
fusion: hybrid::Fusion,
|
||||||
keyword_weight: f32,
|
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
self.ensure_hnsw_fresh();
|
self.ensure_hnsw_fresh();
|
||||||
@@ -38,23 +37,16 @@ impl HDF5Memory {
|
|||||||
// Fusion normalises over every keyword match, so it needs all
|
// Fusion normalises over every keyword match, so it needs all
|
||||||
// the scores — but not ranked.
|
// the scores — but not ranked.
|
||||||
let kw_scores = bm25.scores(query_text);
|
let kw_scores = bm25.scores(query_text);
|
||||||
hybrid::merge_vector_keyword(
|
hybrid::fuse(vec_scores, kw_scores, fusion, k)
|
||||||
vec_scores,
|
|
||||||
kw_scores,
|
|
||||||
vector_weight,
|
|
||||||
keyword_weight,
|
|
||||||
k,
|
|
||||||
)
|
|
||||||
}
|
}
|
||||||
_ => hybrid::hybrid_search(
|
_ => hybrid::hybrid_search_fused(
|
||||||
query_embedding,
|
query_embedding,
|
||||||
query_text,
|
query_text,
|
||||||
&self.cache.embeddings,
|
&self.cache.embeddings,
|
||||||
&self.cache.chunks,
|
&self.cache.chunks,
|
||||||
&self.cache.tombstones,
|
&self.cache.tombstones,
|
||||||
bm25,
|
bm25,
|
||||||
vector_weight,
|
fusion,
|
||||||
keyword_weight,
|
|
||||||
k,
|
k,
|
||||||
),
|
),
|
||||||
}
|
}
|
||||||
@@ -66,19 +58,17 @@ impl HDF5Memory {
|
|||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
bm25: &bm25::BM25Index,
|
bm25: &bm25::BM25Index,
|
||||||
vector_weight: f32,
|
fusion: hybrid::Fusion,
|
||||||
keyword_weight: f32,
|
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
hybrid::hybrid_search(
|
hybrid::hybrid_search_fused(
|
||||||
query_embedding,
|
query_embedding,
|
||||||
query_text,
|
query_text,
|
||||||
&self.cache.embeddings,
|
&self.cache.embeddings,
|
||||||
&self.cache.chunks,
|
&self.cache.chunks,
|
||||||
&self.cache.tombstones,
|
&self.cache.tombstones,
|
||||||
bm25,
|
bm25,
|
||||||
vector_weight,
|
fusion,
|
||||||
keyword_weight,
|
|
||||||
k,
|
k,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -91,20 +81,36 @@ impl HDF5Memory {
|
|||||||
vector_weight: f32,
|
vector_weight: f32,
|
||||||
keyword_weight: f32,
|
keyword_weight: f32,
|
||||||
k: usize,
|
k: usize,
|
||||||
|
) -> Vec<SearchResult> {
|
||||||
|
self.hybrid_search_with(
|
||||||
|
query_embedding,
|
||||||
|
query_text,
|
||||||
|
hybrid::Fusion::Weighted {
|
||||||
|
vector: vector_weight,
|
||||||
|
keyword: keyword_weight,
|
||||||
|
},
|
||||||
|
k,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// [`HDF5Memory::hybrid_search`] with the fusion method chosen explicitly.
|
||||||
|
///
|
||||||
|
/// [`hybrid::DEFAULT_FUSION`] is what the weighted form defaults to;
|
||||||
|
/// [`hybrid::Fusion::Rrf`] combines the two stages by rank instead of by
|
||||||
|
/// score.
|
||||||
|
pub fn hybrid_search_with(
|
||||||
|
&mut self,
|
||||||
|
query_embedding: &[f32],
|
||||||
|
query_text: &str,
|
||||||
|
fusion: hybrid::Fusion,
|
||||||
|
k: usize,
|
||||||
) -> Vec<SearchResult> {
|
) -> Vec<SearchResult> {
|
||||||
// The keyword index lives for the life of the store and is updated
|
// The keyword index lives for the life of the store and is updated
|
||||||
// incrementally. Take it out for the duration of the call so the
|
// incrementally. Take it out for the duration of the call so the
|
||||||
// vector stage can borrow `self` mutably, then put it back.
|
// vector stage can borrow `self` mutably, then put it back.
|
||||||
self.ensure_bm25_fresh();
|
self.ensure_bm25_fresh();
|
||||||
let bm25 = self.bm25.take().expect("ensure_bm25_fresh leaves an index");
|
let bm25 = self.bm25.take().expect("ensure_bm25_fresh leaves an index");
|
||||||
let scored = self.vector_keyword_search(
|
let scored = self.vector_keyword_search(query_embedding, query_text, &bm25, fusion, k);
|
||||||
query_embedding,
|
|
||||||
query_text,
|
|
||||||
&bm25,
|
|
||||||
vector_weight,
|
|
||||||
keyword_weight,
|
|
||||||
k,
|
|
||||||
);
|
|
||||||
let mut results: Vec<SearchResult> = scored
|
let mut results: Vec<SearchResult> = scored
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|(idx, score)| {
|
.map(|(idx, score)| {
|
||||||
|
|||||||
@@ -55,42 +55,57 @@ use std::time::{Duration, Instant};
|
|||||||
#[path = "longmemeval_bench/embedder.rs"]
|
#[path = "longmemeval_bench/embedder.rs"]
|
||||||
mod embedder;
|
mod embedder;
|
||||||
|
|
||||||
|
use clawhdf5_agent::hybrid::Fusion;
|
||||||
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 mode's fusion, as one short string for the reports.
|
||||||
|
fn describe(mode: Mode) -> String {
|
||||||
|
match mode.fusion {
|
||||||
|
Fusion::Weighted { vector, keyword } => format!("vector_{vector:.1}_keyword_{keyword:.1}"),
|
||||||
|
Fusion::Rrf { k } => format!("rrf_k{k:.0}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// A retrieval configuration: how much of the score comes from each stage.
|
/// A retrieval configuration: how much of the score comes from each stage.
|
||||||
#[derive(Clone, Copy)]
|
#[derive(Clone, Copy)]
|
||||||
struct Mode {
|
struct Mode {
|
||||||
label: &'static str,
|
label: &'static str,
|
||||||
vector_weight: f32,
|
/// How the two retrieval stages are combined into one ranking.
|
||||||
keyword_weight: f32,
|
fusion: Fusion,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Mode {
|
||||||
|
const fn weighted(label: &'static str, vector: f32, keyword: f32) -> Self {
|
||||||
|
Self {
|
||||||
|
label,
|
||||||
|
fusion: Fusion::Weighted { vector, keyword },
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// The only mode available without real embeddings. Passing zero vectors with
|
/// The only mode available without real embeddings. Passing zero vectors with
|
||||||
/// `vector_weight = 0.0` is what made the vector stage inert.
|
/// `vector_weight = 0.0` is what made the vector stage inert.
|
||||||
const BM25_ONLY: Mode = Mode {
|
const BM25_ONLY: Mode = Mode::weighted("BM25 only (vector stage inert)", 0.0, 1.0);
|
||||||
label: "BM25 only (vector stage inert)",
|
|
||||||
vector_weight: 0.0,
|
|
||||||
keyword_weight: 1.0,
|
|
||||||
};
|
|
||||||
#[cfg(feature = "embeddings")]
|
#[cfg(feature = "embeddings")]
|
||||||
const VECTOR_ONLY: Mode = Mode {
|
const VECTOR_ONLY: Mode = Mode::weighted("Vector only (MiniLM + HNSW)", 1.0, 0.0);
|
||||||
label: "Vector only (MiniLM + HNSW)",
|
|
||||||
vector_weight: 1.0,
|
|
||||||
keyword_weight: 0.0,
|
|
||||||
};
|
|
||||||
/// Tuned by `--sweep` over the full haystack. The former 0.7/0.3 was a
|
/// 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
|
/// 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
|
/// strictly dominated: 0.4/0.6 is better on Hit@1, Hit@5, Hit@10 and MRR at
|
||||||
/// both granularities.
|
/// both granularities.
|
||||||
#[cfg(feature = "embeddings")]
|
#[cfg(feature = "embeddings")]
|
||||||
const HYBRID: Mode = Mode {
|
const HYBRID: Mode = Mode::weighted("Hybrid (0.4 vector / 0.6 BM25, tuned)", 0.4, 0.6);
|
||||||
label: "Hybrid (0.4 vector / 0.6 BM25, tuned)",
|
|
||||||
vector_weight: 0.4,
|
/// Reciprocal rank fusion, the documented alternative to the weighted sum.
|
||||||
keyword_weight: 0.6,
|
/// 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 },
|
||||||
};
|
};
|
||||||
|
|
||||||
/// Every 0.1 step of vector weight, keyword weight taking the remainder.
|
/// Every 0.1 step of vector weight, keyword weight taking the remainder.
|
||||||
@@ -104,11 +119,11 @@ fn sweep_modes() -> Vec<Mode> {
|
|||||||
(0..=10)
|
(0..=10)
|
||||||
.map(|i| {
|
.map(|i| {
|
||||||
let v = i as f32 / 10.0;
|
let v = i as f32 / 10.0;
|
||||||
Mode {
|
Mode::weighted(
|
||||||
label: Box::leak(format!("sweep v={v:.1} / k={:.1}", 1.0 - v).into_boxed_str()),
|
Box::leak(format!("sweep v={v:.1} / k={:.1}", 1.0 - v).into_boxed_str()),
|
||||||
vector_weight: v,
|
v,
|
||||||
keyword_weight: 1.0 - v,
|
1.0 - v,
|
||||||
}
|
)
|
||||||
})
|
})
|
||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
@@ -321,13 +336,7 @@ fn evaluate_question(
|
|||||||
|
|
||||||
let query_emb = embedding_for(embeddings, &q.question);
|
let query_emb = embedding_for(embeddings, &q.question);
|
||||||
let t0 = Instant::now();
|
let t0 = Instant::now();
|
||||||
let results = memory.hybrid_search(
|
let results = memory.hybrid_search_with(&query_emb, &q.question, mode.fusion, top_k);
|
||||||
&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
|
||||||
@@ -472,10 +481,7 @@ fn print_report(
|
|||||||
println!(" LongMemEval Benchmark — {}", mode.label);
|
println!(" LongMemEval Benchmark — {}", mode.label);
|
||||||
println!("=================================================================");
|
println!("=================================================================");
|
||||||
println!();
|
println!();
|
||||||
println!(
|
println!("Mode: {}", describe(mode));
|
||||||
"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");
|
||||||
@@ -602,10 +608,7 @@ fn print_report(
|
|||||||
println!("```json");
|
println!("```json");
|
||||||
println!("{{");
|
println!("{{");
|
||||||
println!(" \"benchmark\": \"longmemeval\",");
|
println!(" \"benchmark\": \"longmemeval\",");
|
||||||
println!(
|
println!(" \"mode\": \"{}\",", describe(mode));
|
||||||
" \"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,");
|
||||||
@@ -768,7 +771,7 @@ fn main() {
|
|||||||
if sweep {
|
if sweep {
|
||||||
sweep_modes()
|
sweep_modes()
|
||||||
} else {
|
} else {
|
||||||
vec![BM25_ONLY, VECTOR_ONLY, HYBRID]
|
vec![BM25_ONLY, VECTOR_ONLY, HYBRID, RRF]
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
#[cfg(not(feature = "embeddings"))]
|
#[cfg(not(feature = "embeddings"))]
|
||||||
|
|||||||
Reference in New Issue
Block a user