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    /// The width of a single head, if this model knows it.
97    ///
98    /// What "width" means is per head, matching the accessor each head
99    /// corresponds to on the single-head traits:
100    ///
101    /// - [`HeadSet::DENSE`] — the dense vector dimensionality
102    ///   ([`EmbeddingModel::dimensions`](crate::traits::EmbeddingModel::dimensions))
103    /// - [`HeadSet::SPARSE`] — the term-space size
104    ///   ([`SparseEmbeddingModel::vocab_size`](crate::traits::SparseEmbeddingModel::vocab_size))
105    /// - [`HeadSet::MULTI_VECTOR`] — the per-token vector dimensionality
106    ///   ([`MultiVectorEmbeddingModel::dimensions`](crate::traits::MultiVectorEmbeddingModel::dimensions))
107    ///
108    /// Returns `None` for a head this model does not expose, for a `head`
109    /// naming more than one flag, and for implementations that simply don't
110    /// track it — the default. Without it, the runtime cannot serve this model
111    /// through the single-head accessors, since those traits must report a
112    /// width; see [`ModelRuntime::sparse_embedder`](crate::runtime::ModelRuntime::sparse_embedder).
113    fn head_width(&self, head: HeadSet) -> Option<u32> {
114        let _ = head;
115        None
116    }
117
118    /// Optional warmup hook (e.g. load weights into memory on first access).
119    /// The default is a no-op.
120    ///
121    /// # Errors
122    /// Returns an error if the underlying model fails to initialize.
123    async fn warmup(&self) -> Result<()> {
124        Ok(())
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131    use std::sync::Arc;
132
133    // Dyn-safety smoke check: fails to compile if the trait becomes non-object-safe.
134    #[test]
135    fn hybrid_embedding_model_is_dyn_safe() {
136        fn _accept(_: Arc<dyn HybridEmbeddingModel>) {}
137    }
138
139    #[test]
140    fn head_set_membership_and_composition() {
141        let ds = HeadSet::DENSE | HeadSet::SPARSE;
142        assert!(ds.contains(HeadSet::DENSE));
143        assert!(ds.contains(HeadSet::SPARSE));
144        assert!(!ds.contains(HeadSet::MULTI_VECTOR));
145        assert_eq!(
146            HeadSet::DENSE | HeadSet::SPARSE | HeadSet::MULTI_VECTOR,
147            HeadSet::ALL
148        );
149    }
150
151    #[test]
152    fn head_set_intersection_disjoint_and_empty() {
153        // Requested ⊄ available collapses to the available subset.
154        let requested = HeadSet::DENSE | HeadSet::MULTI_VECTOR;
155        let available = HeadSet::DENSE | HeadSet::SPARSE;
156        assert_eq!(requested.intersection(available), HeadSet::DENSE);
157
158        // Fully disjoint → empty; empty request → empty.
159        assert!(HeadSet::SPARSE.intersection(HeadSet::DENSE).is_empty());
160        assert!(HeadSet::empty().intersection(HeadSet::ALL).is_empty());
161    }
162}
uni_xervo/traits/hybrid.rs

uni_xervo/traits/
hybrid.rs
