osobhandClaude Fable 5.1
3ed0489faa
test_hebbian_activation_boost failed intermittently. Root causes, all in the query path: - normalize_scores mapped a set of identical scores — including the single-candidate case — to 0.0, so a lone perfect match contributed nothing to the fused score. Identical positive scores now normalise to 1.0 (all equally the best match); identical non-positive scores stay 0.0. - merge_vector_keyword sorted a HashMap's entries by score alone and then truncated, so which ties survived varied from run to run; hybrid_search had the same problem in its final sort. Both now break ties by index. - hybrid_search applied the Hebbian boost to every returned record, including the zero-score filler that pads the list when fewer than k records match. With random tie-breaking a filler record could collect as many boosts as the real hit. Only records with a positive fused score are reinforced now. Co-Authored-By: Claude Fable 5.1 <[email protected]>
498 lines
15 KiB
Rust
498 lines
15 KiB
Rust
//! Hybrid search combining vector similarity and BM25 keyword scores.
|
|
//!
|
|
//! Normalizes both score sets to [0, 1] and computes a weighted merge.
|
|
|
|
use std::collections::HashMap;
|
|
|
|
use crate::bm25::BM25Index;
|
|
use crate::vector_search;
|
|
|
|
/// Perform hybrid search combining cosine vector similarity and BM25 keyword search.
|
|
///
|
|
/// Both score distributions are independently normalized to [0, 1] before
|
|
/// being combined with the specified weights. Uses pre-computed norms when
|
|
/// available for faster vector search.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `query_embedding` - The query vector for cosine similarity.
|
|
/// * `query_text` - The query text for BM25 keyword search.
|
|
/// * `vectors` - All stored embedding vectors.
|
|
/// * `_chunks` - All stored text chunks (parallel to `vectors`).
|
|
/// * `tombstones` - Tombstone flags (non-zero = deleted).
|
|
/// * `bm25_index` - Pre-built BM25 index.
|
|
/// * `vector_weight` - Weight for vector similarity scores (default 0.7).
|
|
/// * `keyword_weight` - Weight for keyword search scores (default 0.3).
|
|
/// * `k` - Number of top results to return.
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn hybrid_search(
|
|
query_embedding: &[f32],
|
|
query_text: &str,
|
|
vectors: &[Vec<f32>],
|
|
_chunks: &[String],
|
|
tombstones: &[u8],
|
|
bm25_index: &BM25Index,
|
|
vector_weight: f32,
|
|
keyword_weight: f32,
|
|
k: usize,
|
|
) -> Vec<(usize, f32)> {
|
|
// Get raw scores from both systems. Request all results so normalization
|
|
// covers the full distribution.
|
|
// Use parallel search when rayon feature is enabled and vector count > 10K.
|
|
let vec_scores = {
|
|
#[cfg(feature = "parallel")]
|
|
{
|
|
if vectors.len() > 10_000 {
|
|
vector_search::parallel_cosine_batch(
|
|
query_embedding,
|
|
vectors,
|
|
tombstones,
|
|
vectors.len(),
|
|
)
|
|
} else {
|
|
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
|
}
|
|
}
|
|
#[cfg(not(feature = "parallel"))]
|
|
{
|
|
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
|
}
|
|
};
|
|
let kw_scores = bm25_index.search(query_text, vectors.len());
|
|
|
|
merge_vector_keyword(vec_scores, kw_scores, vector_weight, keyword_weight, k)
|
|
}
|
|
|
|
/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
|
|
///
|
|
/// Both score sets are independently min-max normalized to [0, 1] and combined
|
|
/// with the given weights. This is the shared core of [`hybrid_search`]; it is
|
|
/// also used by the optional HNSW path, which supplies vector scores from an
|
|
/// approximate-nearest-neighbour index instead of a full linear scan.
|
|
pub fn merge_vector_keyword(
|
|
vec_scores: Vec<(usize, f32)>,
|
|
kw_scores: Vec<(usize, f32)>,
|
|
vector_weight: f32,
|
|
keyword_weight: f32,
|
|
k: usize,
|
|
) -> Vec<(usize, f32)> {
|
|
// Normalize each set to [0, 1].
|
|
let vec_normalized = normalize_scores(&vec_scores);
|
|
let kw_normalized = normalize_scores(&kw_scores);
|
|
|
|
// Merge scores with weights.
|
|
let mut merged: HashMap<usize, f32> = HashMap::new();
|
|
|
|
for (idx, score) in &vec_normalized {
|
|
*merged.entry(*idx).or_insert(0.0) += vector_weight * score;
|
|
}
|
|
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();
|
|
// Index tie-break: `merged` is a HashMap, so without it the ties that
|
|
// survive `truncate` differ from run to run.
|
|
results.sort_by(|a, b| {
|
|
b.1.partial_cmp(&a.1)
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
.then(a.0.cmp(&b.0))
|
|
});
|
|
results.truncate(k);
|
|
results
|
|
}
|
|
|
|
/// Normalize a set of scores to the [0, 1] range using min-max normalization.
|
|
///
|
|
/// If all scores are identical there is no spread to normalise: each entry
|
|
/// gets 1.0 when that score is positive (all equally the best match) and 0.0
|
|
/// otherwise (nothing matched).
|
|
fn normalize_scores(scores: &[(usize, f32)]) -> Vec<(usize, f32)> {
|
|
if scores.is_empty() {
|
|
return Vec::new();
|
|
}
|
|
|
|
let min = scores.iter().map(|(_, s)| *s).fold(f32::INFINITY, f32::min);
|
|
let max = scores
|
|
.iter()
|
|
.map(|(_, s)| *s)
|
|
.fold(f32::NEG_INFINITY, f32::max);
|
|
|
|
let range = max - min;
|
|
if range == 0.0 {
|
|
// All candidates scored the same (including the single-candidate
|
|
// case), so min-max has no spread to work with. They are all equally
|
|
// the best match if that score is positive, and all non-matches
|
|
// otherwise. This used to return 0.0 unconditionally, which erased a
|
|
// lone perfect match from the fused score.
|
|
let level = if max > 0.0 { 1.0 } else { 0.0 };
|
|
return scores.iter().map(|(idx, _)| (*idx, level)).collect();
|
|
}
|
|
|
|
scores
|
|
.iter()
|
|
.map(|(idx, s)| (*idx, (s - min) / range))
|
|
.collect()
|
|
}
|
|
|
|
/// Perform hybrid search using Reciprocal Rank Fusion (RRF).
|
|
///
|
|
/// RRF combines rankings from multiple retrieval systems without requiring
|
|
/// score normalization. Each result is scored as:
|
|
///
|
|
/// `score = Σ 1 / (k + rank_i)`
|
|
///
|
|
/// where `k = 60` (standard constant that dampens the impact of high ranks)
|
|
/// and `rank_i` is the 1-based rank of the document in retrieval system `i`.
|
|
///
|
|
/// Documents only present in one system still receive a partial score.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `query_embedding` - The query vector for cosine similarity.
|
|
/// * `query_text` - The query text for BM25 keyword search.
|
|
/// * `vectors` - All stored embedding vectors.
|
|
/// * `_chunks` - All stored text chunks (parallel to `vectors`).
|
|
/// * `tombstones` - Tombstone flags (non-zero = deleted).
|
|
/// * `bm25_index` - Pre-built BM25 index.
|
|
/// * `k` - Number of top results to return.
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn rrf_hybrid_search(
|
|
query_embedding: &[f32],
|
|
query_text: &str,
|
|
vectors: &[Vec<f32>],
|
|
_chunks: &[String],
|
|
tombstones: &[u8],
|
|
bm25_index: &BM25Index,
|
|
k: usize,
|
|
) -> Vec<(usize, f32)> {
|
|
const RRF_K: f32 = 60.0;
|
|
|
|
// Retrieve all results from both systems sorted descending by score.
|
|
let mut vec_scores = {
|
|
#[cfg(feature = "parallel")]
|
|
{
|
|
if vectors.len() > 10_000 {
|
|
vector_search::parallel_cosine_batch(
|
|
query_embedding,
|
|
vectors,
|
|
tombstones,
|
|
vectors.len(),
|
|
)
|
|
} else {
|
|
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
|
}
|
|
}
|
|
#[cfg(not(feature = "parallel"))]
|
|
{
|
|
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
|
}
|
|
};
|
|
let mut kw_scores = bm25_index.search(query_text, vectors.len());
|
|
|
|
// Sort both lists descending so rank 1 = best.
|
|
vec_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
|
kw_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
|
|
|
// Accumulate RRF scores.
|
|
let mut rrf_scores: HashMap<usize, f32> = HashMap::new();
|
|
|
|
for (rank, (idx, _score)) in vec_scores.iter().enumerate() {
|
|
let rrf = 1.0 / (RRF_K + (rank + 1) as f32);
|
|
*rrf_scores.entry(*idx).or_insert(0.0) += rrf;
|
|
}
|
|
for (rank, (idx, _score)) in kw_scores.iter().enumerate() {
|
|
let rrf = 1.0 / (RRF_K + (rank + 1) as f32);
|
|
*rrf_scores.entry(*idx).or_insert(0.0) += rrf;
|
|
}
|
|
|
|
let mut results: Vec<(usize, f32)> = rrf_scores.into_iter().collect();
|
|
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
|
results.truncate(k);
|
|
results
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn make_test_data() -> (Vec<Vec<f32>>, Vec<String>, Vec<u8>, BM25Index) {
|
|
// 4 documents with 3-dim embeddings
|
|
let vectors = vec![
|
|
vec![1.0, 0.0, 0.0], // doc 0: points in x direction
|
|
vec![0.0, 1.0, 0.0], // doc 1: points in y direction
|
|
vec![0.7, 0.7, 0.0], // doc 2: between x and y
|
|
vec![0.0, 0.0, 1.0], // doc 3: points in z direction
|
|
];
|
|
let chunks = vec![
|
|
"rust programming language".to_string(),
|
|
"python scripting language".to_string(),
|
|
"rust and python comparison".to_string(),
|
|
"javascript web development".to_string(),
|
|
];
|
|
let tombstones = vec![0u8; 4];
|
|
let bm25 = BM25Index::build(&chunks, &tombstones);
|
|
(vectors, chunks, tombstones, bm25)
|
|
}
|
|
|
|
#[test]
|
|
fn vector_only_search() {
|
|
let (vectors, chunks, tombstones, bm25) = make_test_data();
|
|
let query_emb = vec![1.0, 0.0, 0.0]; // points in x, should match doc 0
|
|
|
|
let results = hybrid_search(
|
|
&query_emb,
|
|
"nonexistent_xyz",
|
|
&vectors,
|
|
&chunks,
|
|
&tombstones,
|
|
&bm25,
|
|
1.0, // vector only
|
|
0.0, // no keyword
|
|
4,
|
|
);
|
|
|
|
assert!(!results.is_empty());
|
|
assert_eq!(
|
|
results[0].0, 0,
|
|
"doc 0 should be top match for x-direction query"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn keyword_only_search() {
|
|
let (vectors, chunks, tombstones, bm25) = make_test_data();
|
|
// Use a zero vector so vector similarity contributes nothing meaningful
|
|
let query_emb = vec![0.0, 0.0, 0.0];
|
|
|
|
let results = hybrid_search(
|
|
&query_emb,
|
|
"rust programming",
|
|
&vectors,
|
|
&chunks,
|
|
&tombstones,
|
|
&bm25,
|
|
0.0, // no vector
|
|
1.0, // keyword only
|
|
4,
|
|
);
|
|
|
|
assert!(!results.is_empty());
|
|
// Doc 0 ("rust programming language") should rank highest for "rust programming"
|
|
assert_eq!(results[0].0, 0);
|
|
}
|
|
|
|
#[test]
|
|
fn balanced_merge_ranking() {
|
|
let (vectors, chunks, tombstones, bm25) = make_test_data();
|
|
// Query embedding close to doc 0, text query for "rust"
|
|
let query_emb = vec![0.9, 0.1, 0.0];
|
|
|
|
let results = hybrid_search(
|
|
&query_emb,
|
|
"rust",
|
|
&vectors,
|
|
&chunks,
|
|
&tombstones,
|
|
&bm25,
|
|
0.7,
|
|
0.3,
|
|
4,
|
|
);
|
|
|
|
assert!(!results.is_empty());
|
|
// Doc 0 should rank high (good vector match + contains "rust")
|
|
// Doc 2 should also appear (contains "rust" + decent vector match)
|
|
let top_ids: Vec<usize> = results.iter().map(|(idx, _)| *idx).collect();
|
|
assert!(top_ids.contains(&0), "doc 0 should appear in results");
|
|
assert!(top_ids.contains(&2), "doc 2 should appear in results");
|
|
}
|
|
|
|
#[test]
|
|
fn empty_results_when_no_data() {
|
|
let vectors: Vec<Vec<f32>> = Vec::new();
|
|
let chunks: Vec<String> = Vec::new();
|
|
let tombstones: Vec<u8> = Vec::new();
|
|
let bm25 = BM25Index::build(&chunks, &tombstones);
|
|
|
|
let results = hybrid_search(
|
|
&[],
|
|
"anything",
|
|
&vectors,
|
|
&chunks,
|
|
&tombstones,
|
|
&bm25,
|
|
0.7,
|
|
0.3,
|
|
10,
|
|
);
|
|
|
|
assert!(results.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn normalize_scores_empty() {
|
|
let result = normalize_scores(&[]);
|
|
assert!(result.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn normalize_scores_single() {
|
|
// A lone positive score is the best match there is, not a non-match.
|
|
let result = normalize_scores(&[(0, 5.0)]);
|
|
assert_eq!(result.len(), 1);
|
|
assert_eq!(result[0].1, 1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn normalize_scores_all_equal() {
|
|
let matched = normalize_scores(&[(0, 0.4), (1, 0.4)]);
|
|
assert!(matched.iter().all(|(_, s)| *s == 1.0));
|
|
let unmatched = normalize_scores(&[(0, 0.0), (1, 0.0)]);
|
|
assert!(unmatched.iter().all(|(_, s)| *s == 0.0));
|
|
}
|
|
|
|
#[test]
|
|
fn normalize_scores_range() {
|
|
let scores = vec![(0, 2.0), (1, 4.0), (2, 6.0)];
|
|
let result = normalize_scores(&scores);
|
|
|
|
assert_eq!(result.len(), 3);
|
|
assert!((result[0].1 - 0.0).abs() < 1e-6); // min -> 0
|
|
assert!((result[1].1 - 0.5).abs() < 1e-6); // mid -> 0.5
|
|
assert!((result[2].1 - 1.0).abs() < 1e-6); // max -> 1
|
|
}
|
|
|
|
#[test]
|
|
fn hybrid_respects_k_limit() {
|
|
let (vectors, chunks, tombstones, bm25) = make_test_data();
|
|
let query_emb = vec![0.5, 0.5, 0.0];
|
|
|
|
let results = hybrid_search(
|
|
&query_emb,
|
|
"language",
|
|
&vectors,
|
|
&chunks,
|
|
&tombstones,
|
|
&bm25,
|
|
0.5,
|
|
0.5,
|
|
2,
|
|
);
|
|
|
|
assert!(results.len() <= 2);
|
|
}
|
|
|
|
// --- RRF tests ---
|
|
|
|
#[test]
|
|
fn rrf_vector_dominant_query() {
|
|
let (vectors, chunks, tombstones, bm25) = make_test_data();
|
|
let query_emb = vec![1.0, 0.0, 0.0]; // strong match on doc 0
|
|
|
|
let results = rrf_hybrid_search(
|
|
&query_emb,
|
|
"nonexistent_xyz",
|
|
&vectors,
|
|
&chunks,
|
|
&tombstones,
|
|
&bm25,
|
|
4,
|
|
);
|
|
|
|
assert!(!results.is_empty());
|
|
assert_eq!(
|
|
results[0].0, 0,
|
|
"doc 0 should top RRF for x-direction query"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn rrf_keyword_dominant_query() {
|
|
let (vectors, chunks, tombstones, bm25) = make_test_data();
|
|
let query_emb = vec![0.0, 0.0, 0.0];
|
|
|
|
let results = rrf_hybrid_search(
|
|
&query_emb,
|
|
"rust programming",
|
|
&vectors,
|
|
&chunks,
|
|
&tombstones,
|
|
&bm25,
|
|
4,
|
|
);
|
|
|
|
assert!(!results.is_empty());
|
|
assert_eq!(
|
|
results[0].0, 0,
|
|
"doc 0 should top RRF for rust programming query"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn rrf_respects_k_limit() {
|
|
let (vectors, chunks, tombstones, bm25) = make_test_data();
|
|
let query_emb = vec![0.5, 0.5, 0.0];
|
|
|
|
let results = rrf_hybrid_search(
|
|
&query_emb,
|
|
"language",
|
|
&vectors,
|
|
&chunks,
|
|
&tombstones,
|
|
&bm25,
|
|
2,
|
|
);
|
|
|
|
assert!(results.len() <= 2);
|
|
}
|
|
|
|
#[test]
|
|
fn rrf_scores_are_positive() {
|
|
let (vectors, chunks, tombstones, bm25) = make_test_data();
|
|
let query_emb = vec![0.5, 0.5, 0.0];
|
|
|
|
let results =
|
|
rrf_hybrid_search(&query_emb, "rust", &vectors, &chunks, &tombstones, &bm25, 4);
|
|
|
|
for (_, score) in &results {
|
|
assert!(*score > 0.0, "RRF scores must be positive");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn rrf_empty_data() {
|
|
let vectors: Vec<Vec<f32>> = Vec::new();
|
|
let chunks: Vec<String> = Vec::new();
|
|
let tombstones: Vec<u8> = Vec::new();
|
|
let bm25 = BM25Index::build(&chunks, &tombstones);
|
|
|
|
let results = rrf_hybrid_search(&[], "anything", &vectors, &chunks, &tombstones, &bm25, 10);
|
|
|
|
assert!(results.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn rrf_scores_sorted_descending() {
|
|
let (vectors, chunks, tombstones, bm25) = make_test_data();
|
|
let query_emb = vec![0.7, 0.3, 0.0];
|
|
|
|
let results = rrf_hybrid_search(
|
|
&query_emb,
|
|
"rust language",
|
|
&vectors,
|
|
&chunks,
|
|
&tombstones,
|
|
&bm25,
|
|
4,
|
|
);
|
|
|
|
for window in results.windows(2) {
|
|
assert!(
|
|
window[0].1 >= window[1].1,
|
|
"RRF results must be sorted descending"
|
|
);
|
|
}
|
|
}
|
|
}
|