uni_xervo/traits/
multivector.rs

1//! Multi-vector (late-interaction / ColBERT) embedding types and the
2//! [`MultiVectorEmbeddingModel`] trait.
3//!
4//! Where [`EmbeddingModel`](crate::traits::EmbeddingModel) pools a sequence into
5//! one dense vector, a multi-vector model emits *one vector per token*. Scoring
6//! is then done with MaxSim (sum over query tokens of the max similarity to any
7//! document token) — the late-interaction approach popularized by ColBERT and
8//! the strongest known method for layout-rich document retrieval (ColPali /
9//! ColQwen2).
10//!
11//! This trait is a separate, narrow capability so a provider opts into it only
12//! when it can skip pooling and surface per-token vectors. The native
13//! multi-vector index is deliberately out of scope — this is the *producer*
14//! side, immediately usable via host-side MaxSim (see
15//! [`max_sim`](crate::score::max_sim)).
16
17use crate::error::Result;
18use crate::traits::{ModelInfo, TokenUsage};
19use async_trait::async_trait;
20
21/// Result of a multi-vector embedding call: per-token vectors for each input.
22///
23/// `vectors` is indexed input → token → dimension. It is *ragged*: padding and
24/// (optionally) special tokens are stripped, so the token count varies per
25/// input. Each per-token vector has [`MultiVectorEmbeddingModel::dimensions`]
26/// elements. `usage` follows the [`EmbedResult`](crate::traits::EmbedResult)
27/// convention — `Some` for remote providers, `None` for local ones.
28#[derive(Debug, Clone)]
29pub struct MultiVectorEmbedResult {
30    /// Per input, per retained token, one dense vector.
31    pub vectors: Vec<Vec<Vec<f32>>>,
32    /// Token usage reported by the provider, if any.
33    pub usage: Option<TokenUsage>,
34}
35
36/// A model that produces per-token (multi-vector / ColBERT) embeddings from text.
37///
38/// Use cases include late-interaction reranking with BGE-M3's ColBERT head and
39/// `answerai-colbert-small`. Each input yields a variable-length list of
40/// per-token vectors; score them against a query with
41/// [`max_sim`](crate::score::max_sim).
42#[async_trait]
43pub trait MultiVectorEmbeddingModel: ModelInfo {
44    /// Embed a batch of text strings into per-token vectors.
45    ///
46    /// Returns one ragged list of per-token vectors per input, in input order.
47    ///
48    /// # Errors
49    /// Returns an error if tokenization fails, the model session errors, or the
50    /// upstream API rejects the request.
51    async fn embed(&self, texts: &[&str]) -> Result<MultiVectorEmbedResult>;
52
53    /// The dimensionality of each per-token vector produced by this model.
54    fn dimensions(&self) -> u32;
55
56    /// Optional warmup hook (e.g. load weights into memory on first access).
57    /// The default is a no-op.
58    ///
59    /// # Errors
60    /// Returns an error if the underlying model fails to initialize.
61    async fn warmup(&self) -> Result<()> {
62        Ok(())
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69    use std::sync::Arc;
70
71    // Dyn-safety smoke check: fails to compile if the trait becomes non-object-safe.
72    #[test]
73    fn multi_vector_embedding_model_is_dyn_safe() {
74        fn _accept(_: Arc<dyn MultiVectorEmbeddingModel>) {}
75    }
76}
uni_xervo/traits/multivector.rs

uni_xervo/traits/
multivector.rs
