uni_xervo/traits/
sparse.rs

1//! Learned-sparse text embedding types and the [`SparseEmbeddingModel`] trait.
2//!
3//! Sparse embeddings represent text as a small set of weighted vocabulary terms
4//! rather than a single dense vector. They power hybrid first-stage retrieval:
5//! a learned-sparse model (SPLADE family, the BGE-M3 sparse head) keeps the
6//! inverted-index compatibility of BM25 while learning term weights — and term
7//! expansion, for SPLADE — that beat BM25 in and out of domain.
8//!
9//! This trait sits alongside [`EmbeddingModel`](crate::traits::EmbeddingModel):
10//! it is a separate, narrow capability so a provider opts into it only when it
11//! can produce term-weight vectors. The downstream inverted index and impact
12//! scoring are deliberately out of scope — this is the *producer* side only.
13
14use crate::error::Result;
15use crate::traits::{ModelInfo, TokenUsage};
16use async_trait::async_trait;
17
18/// A single learned-sparse vector as `(term_id, weight)` pairs.
19///
20/// Each `term_id` indexes into the model's term space (see
21/// [`SparseEmbeddingModel::vocab_size`]); `weight` is the non-negative learned
22/// importance of that term. Entries are sparse: only non-zero terms appear.
23/// Term ids are only comparable within the *same* model — two models with
24/// different vocabularies produce incompatible sparse vectors.
25pub type SparseVector = Vec<(u32, f32)>;
26
27/// Result of a sparse-embedding call, carrying one [`SparseVector`] per input.
28///
29/// Mirrors [`EmbedResult`](crate::traits::EmbedResult): `usage` is `Some` when a
30/// remote provider reports token counts and `None` for local providers where the
31/// concept does not apply.
32#[derive(Debug, Clone)]
33pub struct SparseEmbedResult {
34    /// One sparse vector per input, in input order.
35    pub vectors: Vec<SparseVector>,
36    /// Token usage reported by the provider, if any.
37    pub usage: Option<TokenUsage>,
38}
39
40/// A model that produces learned-sparse term-weight vectors from text.
41///
42/// Use cases include SPLADE-style retrieval and the BGE-M3 sparse head for
43/// hybrid dense + sparse search. One [`SparseVector`] is returned per input
44/// text. For scoring two sparse vectors, see
45/// [`sparse_dot`](crate::score::sparse_dot).
46#[async_trait]
47pub trait SparseEmbeddingModel: ModelInfo {
48    /// Embed a batch of text strings into learned-sparse vectors.
49    ///
50    /// Returns one [`SparseVector`] per input, in input order.
51    ///
52    /// # Errors
53    /// Returns an error if tokenization fails, the model session errors, or the
54    /// upstream API rejects the request.
55    async fn embed(&self, texts: &[&str]) -> Result<SparseEmbedResult>;
56
57    /// The size of the term space the produced `term_id`s index into.
58    ///
59    /// For SPLADE this is the model's full vocabulary (term expansion); for the
60    /// BGE-M3 lexical head it is the tokenizer vocabulary the input tokens map to.
61    fn vocab_size(&self) -> u32;
62
63    /// Optional warmup hook (e.g. load weights into memory on first access).
64    /// The default is a no-op.
65    ///
66    /// # Errors
67    /// Returns an error if the underlying model fails to initialize.
68    async fn warmup(&self) -> Result<()> {
69        Ok(())
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use std::sync::Arc;
77
78    // Dyn-safety smoke check: fails to compile if the trait becomes non-object-safe.
79    #[test]
80    fn sparse_embedding_model_is_dyn_safe() {
81        fn _accept(_: Arc<dyn SparseEmbeddingModel>) {}
82    }
83}
uni_xervo/traits/sparse.rs

uni_xervo/traits/
sparse.rs
