mirror of
https://github.com/tinyhumansai/openhuman.git
synced 2026-07-27 21:08:00 +00:00
refactor(tinycortex): W5 memory_search vector+scoring shim + remaining-Ws readiness map (#4790)
This commit is contained in:
@@ -1,64 +1,30 @@
|
||||
//! Scoring weight profiles for hybrid retrieval.
|
||||
//! Scoring weight profiles for hybrid retrieval — thin host shim over
|
||||
//! `tinycortex::memory::WeightProfile` (W5).
|
||||
//!
|
||||
//! The weight profile (graph/vector/keyword/freshness weights + the
|
||||
//! `BALANCED`/`SEMANTIC`/`LEXICAL`/`GRAPH_FIRST` presets + `by_name`) is the
|
||||
//! crate's, a byte-identical port. The host keeps only [`compose_score`] — the
|
||||
//! trivial weighted combination the crate expresses via
|
||||
//! `retrieval::scoring::hybrid_score` at its own call sites; exposed here as a
|
||||
//! free function so `memory_search::tools::hybrid_search` keeps its call shape.
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct WeightProfile {
|
||||
pub graph: f64,
|
||||
pub vector: f64,
|
||||
pub keyword: f64,
|
||||
pub freshness: f64,
|
||||
}
|
||||
pub use tinycortex::memory::WeightProfile;
|
||||
|
||||
impl WeightProfile {
|
||||
pub const BALANCED: Self = Self {
|
||||
graph: 0.35,
|
||||
vector: 0.35,
|
||||
keyword: 0.15,
|
||||
freshness: 0.15,
|
||||
};
|
||||
|
||||
pub const SEMANTIC: Self = Self {
|
||||
graph: 0.15,
|
||||
vector: 0.65,
|
||||
keyword: 0.20,
|
||||
freshness: 0.0,
|
||||
};
|
||||
|
||||
pub const LEXICAL: Self = Self {
|
||||
graph: 0.25,
|
||||
vector: 0.15,
|
||||
keyword: 0.60,
|
||||
freshness: 0.0,
|
||||
};
|
||||
|
||||
pub const GRAPH_FIRST: Self = Self {
|
||||
graph: 0.55,
|
||||
vector: 0.30,
|
||||
keyword: 0.15,
|
||||
freshness: 0.0,
|
||||
};
|
||||
|
||||
pub fn from_name(name: &str) -> Option<Self> {
|
||||
match name {
|
||||
"balanced" => Some(Self::BALANCED),
|
||||
"semantic" => Some(Self::SEMANTIC),
|
||||
"lexical" => Some(Self::LEXICAL),
|
||||
"graph_first" => Some(Self::GRAPH_FIRST),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn compose_score(
|
||||
&self,
|
||||
graph_relevance: f64,
|
||||
vector_similarity: f64,
|
||||
keyword_relevance: f64,
|
||||
freshness: f64,
|
||||
) -> f64 {
|
||||
(self.graph * graph_relevance)
|
||||
+ (self.vector * vector_similarity)
|
||||
+ (self.keyword * keyword_relevance)
|
||||
+ (self.freshness * freshness)
|
||||
}
|
||||
/// Weighted composite of the four retrieval signals under `profile`.
|
||||
///
|
||||
/// `graph·graph_relevance + vector·vector_similarity + keyword·keyword_relevance
|
||||
/// + freshness·freshness`.
|
||||
pub fn compose_score(
|
||||
profile: &WeightProfile,
|
||||
graph_relevance: f64,
|
||||
vector_similarity: f64,
|
||||
keyword_relevance: f64,
|
||||
freshness: f64,
|
||||
) -> f64 {
|
||||
(profile.graph * graph_relevance)
|
||||
+ (profile.vector * vector_similarity)
|
||||
+ (profile.keyword * keyword_relevance)
|
||||
+ (profile.freshness * freshness)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -82,19 +48,30 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_name_resolves() {
|
||||
assert!(WeightProfile::from_name("balanced").is_some());
|
||||
assert!(WeightProfile::from_name("semantic").is_some());
|
||||
assert!(WeightProfile::from_name("lexical").is_some());
|
||||
assert!(WeightProfile::from_name("graph_first").is_some());
|
||||
assert!(WeightProfile::from_name("unknown").is_none());
|
||||
fn by_name_resolves_with_balanced_fallback() {
|
||||
assert_eq!(
|
||||
WeightProfile::by_name("semantic").vector,
|
||||
WeightProfile::SEMANTIC.vector
|
||||
);
|
||||
assert_eq!(
|
||||
WeightProfile::by_name("lexical").keyword,
|
||||
WeightProfile::LEXICAL.keyword
|
||||
);
|
||||
assert_eq!(
|
||||
WeightProfile::by_name("graph_first").graph,
|
||||
WeightProfile::GRAPH_FIRST.graph
|
||||
);
|
||||
// Unknown names fall back to balanced.
|
||||
assert_eq!(
|
||||
WeightProfile::by_name("unknown").graph,
|
||||
WeightProfile::BALANCED.graph
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compose_score_applies_weights() {
|
||||
let profile = WeightProfile::SEMANTIC;
|
||||
let score = profile.compose_score(0.5, 1.0, 0.5, 0.0);
|
||||
let expected = 0.15 * 0.5 + 0.65 * 1.0 + 0.20 * 0.5;
|
||||
assert!((score - expected).abs() < f64::EPSILON);
|
||||
fn compose_score_is_weighted_sum() {
|
||||
let p = WeightProfile::BALANCED;
|
||||
let s = compose_score(&p, 1.0, 1.0, 1.0, 1.0);
|
||||
assert!((s - (p.graph + p.vector + p.keyword + p.freshness)).abs() < 1e-9);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -110,7 +110,7 @@ impl Tool for MemoryHybridSearchTool {
|
||||
));
|
||||
}
|
||||
|
||||
let profile = WeightProfile::from_name(&parsed.mode).unwrap_or(WeightProfile::BALANCED);
|
||||
let profile = WeightProfile::by_name(&parsed.mode);
|
||||
let limit = parsed.limit.clamp(1, 50);
|
||||
|
||||
log::debug!(
|
||||
@@ -152,7 +152,8 @@ impl Tool for MemoryHybridSearchTool {
|
||||
.enumerate()
|
||||
.map(|(i, hit)| {
|
||||
let bd = &hit.score_breakdown;
|
||||
let score = profile.compose_score(
|
||||
let score = crate::openhuman::memory_search::scoring::compose_score(
|
||||
&profile,
|
||||
bd.graph_relevance,
|
||||
bd.vector_similarity,
|
||||
bd.keyword_relevance,
|
||||
|
||||
@@ -1,218 +1,9 @@
|
||||
//! Maximal Marginal Relevance (MMR) selection.
|
||||
//! Maximal Marginal Relevance selection — thin host re-export of
|
||||
//! `tinycortex::memory::retrieval::mmr` (W5).
|
||||
//!
|
||||
//! Given a set of candidate vectors and a query vector, selects a diverse
|
||||
//! subset that balances relevance to the query against redundancy within
|
||||
//! the selected set.
|
||||
//! The MMR algorithm (relevance–diversity tradeoff over embeddings) is the
|
||||
//! crate's, a byte-identical port. Host consumers (`memory_search::tools`) keep
|
||||
//! their `memory_search::vector::mmr::{MmrCandidate, MmrResult, mmr_select}`
|
||||
//! import paths unchanged.
|
||||
|
||||
use crate::openhuman::memory_store::vectors::cosine_similarity;
|
||||
|
||||
/// A candidate for MMR selection.
|
||||
pub struct MmrCandidate<'a> {
|
||||
pub index: usize,
|
||||
pub embedding: &'a [f32],
|
||||
pub relevance: f64,
|
||||
}
|
||||
|
||||
/// Result of MMR selection: the original index and its MMR score.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MmrResult {
|
||||
pub index: usize,
|
||||
pub score: f64,
|
||||
}
|
||||
|
||||
/// Selects up to `limit` items from `candidates` using MMR.
|
||||
///
|
||||
/// `lambda` controls the relevance-diversity tradeoff:
|
||||
/// - 1.0 = pure relevance (no diversity)
|
||||
/// - 0.0 = pure diversity (ignores relevance)
|
||||
/// - 0.7 = recommended default
|
||||
///
|
||||
/// For each selection step:
|
||||
/// mmr(c) = lambda * relevance(c) - (1-lambda) * max_similarity(c, selected)
|
||||
pub fn mmr_select(
|
||||
query_vec: &[f32],
|
||||
candidates: &[MmrCandidate<'_>],
|
||||
limit: usize,
|
||||
lambda: f64,
|
||||
) -> Vec<MmrResult> {
|
||||
if candidates.is_empty() || limit == 0 {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let lambda = lambda.clamp(0.0, 1.0);
|
||||
let limit = limit.min(candidates.len());
|
||||
|
||||
let mut selected: Vec<usize> = Vec::with_capacity(limit);
|
||||
let mut selected_embeddings: Vec<&[f32]> = Vec::with_capacity(limit);
|
||||
let mut results: Vec<MmrResult> = Vec::with_capacity(limit);
|
||||
let mut available: Vec<bool> = vec![true; candidates.len()];
|
||||
|
||||
for _ in 0..limit {
|
||||
let mut best_idx: Option<usize> = None;
|
||||
let mut best_mmr = f64::NEG_INFINITY;
|
||||
|
||||
for (i, candidate) in candidates.iter().enumerate() {
|
||||
if !available[i] {
|
||||
continue;
|
||||
}
|
||||
|
||||
let max_sim_to_selected = if selected_embeddings.is_empty() {
|
||||
0.0
|
||||
} else {
|
||||
selected_embeddings
|
||||
.iter()
|
||||
.map(|sel| cosine_similarity(candidate.embedding, sel))
|
||||
.fold(0.0_f64, f64::max)
|
||||
};
|
||||
|
||||
let mmr_score = lambda * candidate.relevance - (1.0 - lambda) * max_sim_to_selected;
|
||||
|
||||
if mmr_score > best_mmr {
|
||||
best_mmr = mmr_score;
|
||||
best_idx = Some(i);
|
||||
}
|
||||
}
|
||||
|
||||
let Some(idx) = best_idx else { break };
|
||||
|
||||
available[idx] = false;
|
||||
selected.push(idx);
|
||||
selected_embeddings.push(candidates[idx].embedding);
|
||||
results.push(MmrResult {
|
||||
index: candidates[idx].index,
|
||||
score: best_mmr,
|
||||
});
|
||||
}
|
||||
|
||||
let _ = (query_vec, &selected);
|
||||
results
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_vec(vals: &[f32]) -> Vec<f32> {
|
||||
vals.to_vec()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_candidates_returns_empty() {
|
||||
let query = make_vec(&[1.0, 0.0, 0.0]);
|
||||
let result = mmr_select(&query, &[], 5, 0.7);
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_candidate() {
|
||||
let query = make_vec(&[1.0, 0.0, 0.0]);
|
||||
let emb = make_vec(&[1.0, 0.0, 0.0]);
|
||||
let candidates = vec![MmrCandidate {
|
||||
index: 0,
|
||||
embedding: &emb,
|
||||
relevance: 0.95,
|
||||
}];
|
||||
let result = mmr_select(&query, &candidates, 5, 0.7);
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].index, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diversity_selects_distinct_vectors() {
|
||||
let query = make_vec(&[1.0, 0.0, 0.0]);
|
||||
|
||||
// Three near-duplicates (all close to query) + two distinct vectors
|
||||
let dup1 = make_vec(&[0.99, 0.01, 0.0]);
|
||||
let dup2 = make_vec(&[0.98, 0.02, 0.0]);
|
||||
let dup3 = make_vec(&[0.97, 0.03, 0.0]);
|
||||
let distinct1 = make_vec(&[0.0, 1.0, 0.0]);
|
||||
let distinct2 = make_vec(&[0.0, 0.0, 1.0]);
|
||||
|
||||
let candidates = vec![
|
||||
MmrCandidate {
|
||||
index: 0,
|
||||
embedding: &dup1,
|
||||
relevance: 0.99,
|
||||
},
|
||||
MmrCandidate {
|
||||
index: 1,
|
||||
embedding: &dup2,
|
||||
relevance: 0.98,
|
||||
},
|
||||
MmrCandidate {
|
||||
index: 2,
|
||||
embedding: &dup3,
|
||||
relevance: 0.97,
|
||||
},
|
||||
MmrCandidate {
|
||||
index: 3,
|
||||
embedding: &distinct1,
|
||||
relevance: 0.50,
|
||||
},
|
||||
MmrCandidate {
|
||||
index: 4,
|
||||
embedding: &distinct2,
|
||||
relevance: 0.45,
|
||||
},
|
||||
];
|
||||
|
||||
let result = mmr_select(&query, &candidates, 3, 0.5);
|
||||
assert_eq!(result.len(), 3);
|
||||
|
||||
// With lambda=0.5, should pick one from the cluster then diversify
|
||||
let selected_indices: Vec<usize> = result.iter().map(|r| r.index).collect();
|
||||
// At most one duplicate should be selected with strong diversity
|
||||
let dup_count = selected_indices.iter().filter(|&&i| i <= 2).count();
|
||||
assert!(dup_count <= 2, "MMR should diversify away from duplicates");
|
||||
// At least one distinct vector should be picked
|
||||
let distinct_count = selected_indices.iter().filter(|&&i| i >= 3).count();
|
||||
assert!(
|
||||
distinct_count >= 1,
|
||||
"MMR should select at least one distinct vector"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lambda_one_is_pure_relevance() {
|
||||
let query = make_vec(&[1.0, 0.0, 0.0]);
|
||||
let emb1 = make_vec(&[0.99, 0.01, 0.0]);
|
||||
let emb2 = make_vec(&[0.0, 1.0, 0.0]);
|
||||
|
||||
let candidates = vec![
|
||||
MmrCandidate {
|
||||
index: 0,
|
||||
embedding: &emb1,
|
||||
relevance: 0.99,
|
||||
},
|
||||
MmrCandidate {
|
||||
index: 1,
|
||||
embedding: &emb2,
|
||||
relevance: 0.50,
|
||||
},
|
||||
];
|
||||
|
||||
let result = mmr_select(&query, &candidates, 2, 1.0);
|
||||
assert_eq!(result[0].index, 0);
|
||||
assert_eq!(result[1].index, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn limit_caps_output() {
|
||||
let query = make_vec(&[1.0, 0.0]);
|
||||
let embs: Vec<Vec<f32>> = (0..10)
|
||||
.map(|i| make_vec(&[1.0 - i as f32 * 0.1, i as f32 * 0.1]))
|
||||
.collect();
|
||||
let candidates: Vec<MmrCandidate> = embs
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, e)| MmrCandidate {
|
||||
index: i,
|
||||
embedding: e,
|
||||
relevance: 1.0 - i as f64 * 0.1,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let result = mmr_select(&query, &candidates, 3, 0.7);
|
||||
assert_eq!(result.len(), 3);
|
||||
}
|
||||
}
|
||||
pub use tinycortex::memory::retrieval::mmr::{mmr_select, MmrCandidate, MmrResult};
|
||||
|
||||
Reference in New Issue
Block a user