spatialrust_semantic/
search.rs

1//! Multimodal fusion scoring and brute-force semantic search.
2
3use crate::{
4    cosine_similarity, Embedding, EntityId, SemanticEntity, SemanticError, SemanticResult,
5};
6
7/// Weighted fusion of embedding similarity and label confidence.
8#[derive(Clone, Copy, Debug, PartialEq)]
9pub struct FusionScore {
10    /// Final fused score.
11    pub score: f32,
12    /// Embedding similarity contribution.
13    pub embedding: f32,
14    /// Best label confidence contribution.
15    pub label: f32,
16}
17
18/// Multimodal fusion weights.
19#[derive(Clone, Copy, Debug, PartialEq)]
20pub struct MultimodalFusion {
21    /// Weight for embedding cosine.
22    pub embedding_weight: f32,
23    /// Weight for best open-vocab confidence.
24    pub label_weight: f32,
25}
26
27impl Default for MultimodalFusion {
28    fn default() -> Self {
29        Self { embedding_weight: 0.7, label_weight: 0.3 }
30    }
31}
32
33impl MultimodalFusion {
34    /// Scores one entity against a query embedding.
35    pub fn score(&self, entity: &SemanticEntity, query: &Embedding) -> SemanticResult<FusionScore> {
36        let embedding = match &entity.embedding {
37            Some(values) => cosine_similarity(values, query)?,
38            None => 0.0,
39        };
40        let label = entity.labels.iter().map(|l| l.confidence).fold(0.0_f32, f32::max);
41        let score = self.embedding_weight * embedding + self.label_weight * label;
42        Ok(FusionScore { score, embedding, label })
43    }
44}
45
46/// In-memory nearest-neighbor index over semantic entities.
47#[derive(Clone, Debug, Default)]
48pub struct SemanticSearchIndex {
49    entities: Vec<SemanticEntity>,
50}
51
52impl SemanticSearchIndex {
53    /// Creates an empty index.
54    #[must_use]
55    pub fn new() -> Self {
56        Self::default()
57    }
58
59    /// Inserts an entity.
60    pub fn insert(&mut self, entity: SemanticEntity) {
61        self.entities.push(entity);
62    }
63
64    /// Returns top-k entity ids by fused score.
65    pub fn search(
66        &self,
67        query: &Embedding,
68        fusion: MultimodalFusion,
69        k: usize,
70    ) -> SemanticResult<Vec<(EntityId, FusionScore)>> {
71        if k == 0 {
72            return Err(SemanticError::InvalidConfiguration("k must be positive".into()));
73        }
74        let mut scored = self
75            .entities
76            .iter()
77            .map(|entity| fusion.score(entity, query).map(|score| (entity.id.clone(), score)))
78            .collect::<SemanticResult<Vec<_>>>()?;
79        scored
80            .sort_by(|a, b| b.1.score.partial_cmp(&a.1.score).unwrap_or(std::cmp::Ordering::Equal));
81        scored.truncate(k);
82        Ok(scored)
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::{MultimodalFusion, SemanticSearchIndex};
89    use crate::{Embedding, EntityId, OpenVocabLabel, SemanticEntity};
90
91    #[test]
92    fn ranks_entities_by_fusion() {
93        let mut index = SemanticSearchIndex::new();
94        index.insert(SemanticEntity {
95            id: EntityId::new("a"),
96            centroid: None,
97            labels: vec![OpenVocabLabel { text: "chair".into(), confidence: 0.2 }],
98            embedding: Some(Embedding::try_new(vec![1.0, 0.0]).unwrap()),
99        });
100        index.insert(SemanticEntity {
101            id: EntityId::new("b"),
102            centroid: None,
103            labels: vec![OpenVocabLabel { text: "table".into(), confidence: 0.9 }],
104            embedding: Some(Embedding::try_new(vec![0.0, 1.0]).unwrap()),
105        });
106        let hits = index
107            .search(&Embedding::try_new(vec![1.0, 0.0]).unwrap(), MultimodalFusion::default(), 1)
108            .unwrap();
109        assert_eq!(hits[0].0, EntityId::new("a"));
110    }
111}
spatialrust_semantic/search.rs

spatialrust_semantic/
search.rs
