uni_xervo/traits/
hybrid.rs

1//! Single-pass hybrid embedding: the [`HybridEmbeddingModel`] trait and its
2//! [`HeadSet`] / [`HybridEmbedResult`] types.
3//!
4//! Some exports — notably BGE-M3 (`aapot/bge-m3-onnx`) — fuse a dense head, a
5//! learned-sparse head, and a multi-vector (ColBERT) head into one graph. The
6//! per-task [`EmbeddingModel`](crate::traits::EmbeddingModel),
7//! [`SparseEmbeddingModel`](crate::traits::SparseEmbeddingModel), and
8//! [`MultiVectorEmbeddingModel`](crate::traits::MultiVectorEmbeddingModel)
9//! handles each load their own session and run their own forward pass, so
10//! serving all three means loading the weights three times and running the graph
11//! three times.
12//!
13//! This trait collapses that to one: a single shared encoder forward pass whose
14//! outputs are post-processed into every requested head. It exists only for
15//! graphs that actually expose multiple heads (declared by a hybrid preset);
16//! single-head models keep using the per-task handles. The caller picks which
17//! heads to materialize with a [`HeadSet`], mirroring how
18//! [`NlpTasks`](crate::traits::NlpTasks) selects NLP heads on a single pass.
19
20use crate::error::Result;
21use crate::traits::{ModelInfo, SparseVector, TokenUsage};
22use async_trait::async_trait;
23use bitflags::bitflags;
24
25bitflags! {
26    /// Selects which embedding heads a [`HybridEmbeddingModel`] should populate.
27    ///
28    /// A head is only post-processed when its flag is set *and* the model
29    /// exposes it (see [`HybridEmbeddingModel::available_heads`]); the
30    /// intersection drives which fields of [`HybridEmbedResult`] are `Some`.
31    /// Skipping a head skips its post-processing, not the (shared) forward pass.
32    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
33    pub struct HeadSet: u32 {
34        /// Pooled dense vector (one per input).
35        const DENSE = 1 << 0;
36        /// Learned-sparse term weights (one [`SparseVector`] per input).
37        const SPARSE = 1 << 1;
38        /// Per-token multi-vector / ColBERT embeddings (ragged, per input).
39        const MULTI_VECTOR = 1 << 2;
40        /// Convenience: all currently-defined heads.
41        const ALL = Self::DENSE.bits() | Self::SPARSE.bits() | Self::MULTI_VECTOR.bits();
42    }
43}
44
45/// Dense, sparse, and multi-vector embeddings from a single forward pass.
46///
47/// Each field is `Some` iff its head was both **requested** (set in the `heads`
48/// argument to [`HybridEmbeddingModel::embed`]) and **available** (set in
49/// [`HybridEmbeddingModel::available_heads`]); otherwise it is `None`. When a
50/// requested-and-available head is computed over an empty input batch the field
51/// is `Some` of an empty `Vec`, never `None`.
52///
53/// `usage` follows the [`EmbedResult`](crate::traits::EmbedResult) convention:
54/// `Some` for remote providers that report token counts, `None` for local ones.
55#[derive(Debug, Clone, Default)]
56#[non_exhaustive]
57pub struct HybridEmbedResult {
58    /// Pooled dense vectors, one per input. `Some` iff the dense head was
59    /// requested and available.
60    pub dense: Option<Vec<Vec<f32>>>,
61    /// Learned-sparse vectors, one per input. `Some` iff the sparse head was
62    /// requested and available.
63    pub sparse: Option<Vec<SparseVector>>,
64    /// Per-token (multi-vector / ColBERT) embeddings, one ragged list per input.
65    /// `Some` iff the multi-vector head was requested and available.
66    pub multi_vector: Option<Vec<Vec<Vec<f32>>>>,
67    /// Token usage reported by the provider, if any.
68    pub usage: Option<TokenUsage>,
69}
70
71/// A model that produces several embedding heads from one shared forward pass.
72///
73/// Backed by multi-output exports such as BGE-M3 (`aapot/bge-m3-onnx`): one
74/// encoder run feeds the dense, sparse, and ColBERT post-processors, so a hybrid
75/// retrieval pipeline pays for one weight load and one forward pass instead of
76/// three. Select the heads to materialize with a [`HeadSet`].
77#[async_trait]
78pub trait HybridEmbeddingModel: ModelInfo {
79    /// Embed a batch of text strings, populating the heads in
80    /// `heads ∩ available_heads()` from a single forward pass.
81    ///
82    /// Heads outside that intersection are left `None` on the result; the others
83    /// are computed and returned in input order.
84    ///
85    /// # Errors
86    /// Returns an error if tokenization fails, the model session errors, or a
87    /// requested head's output cannot be read from the graph.
88    async fn embed(&self, texts: &[&str], heads: HeadSet) -> Result<HybridEmbedResult>;
89
90    /// The heads this model's graph actually exposes.
91    ///
92    /// Parallels [`NlpModel::supported_tasks`](crate::traits::NlpModel::supported_tasks):
93    /// requesting a head outside this set yields `None` for it rather than an error.
94    fn available_heads(&self) -> HeadSet;
95
96    /// Optional warmup hook (e.g. load weights into memory on first access).
97    /// The default is a no-op.
98    ///
99    /// # Errors
100    /// Returns an error if the underlying model fails to initialize.
101    async fn warmup(&self) -> Result<()> {
102        Ok(())
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use std::sync::Arc;
110
111    // Dyn-safety smoke check: fails to compile if the trait becomes non-object-safe.
112    #[test]
113    fn hybrid_embedding_model_is_dyn_safe() {
114        fn _accept(_: Arc<dyn HybridEmbeddingModel>) {}
115    }
116
117    #[test]
118    fn head_set_membership_and_composition() {
119        let ds = HeadSet::DENSE | HeadSet::SPARSE;
120        assert!(ds.contains(HeadSet::DENSE));
121        assert!(ds.contains(HeadSet::SPARSE));
122        assert!(!ds.contains(HeadSet::MULTI_VECTOR));
123        assert_eq!(
124            HeadSet::DENSE | HeadSet::SPARSE | HeadSet::MULTI_VECTOR,
125            HeadSet::ALL
126        );
127    }
128
129    #[test]
130    fn head_set_intersection_disjoint_and_empty() {
131        // Requested ⊄ available collapses to the available subset.
132        let requested = HeadSet::DENSE | HeadSet::MULTI_VECTOR;
133        let available = HeadSet::DENSE | HeadSet::SPARSE;
134        assert_eq!(requested.intersection(available), HeadSet::DENSE);
135
136        // Fully disjoint → empty; empty request → empty.
137        assert!(HeadSet::SPARSE.intersection(HeadSet::DENSE).is_empty());
138        assert!(HeadSet::empty().intersection(HeadSet::ALL).is_empty());
139    }
140}
uni_xervo/traits/hybrid.rs

uni_xervo/traits/
hybrid.rs
