pocket_encoder/
lib.rs

1//! Bundled sentence-transformer embeddings, offline, with native Rust CPU inference.
2//!
3//! This runs all three attention layers of MiniLM. The 4-bit **storage** format is
4//! dequantized to F32 once at load; inference is F32, not an integer kernel.
5//! No model downloads, Python, ONNX shared library, cache or build script.
6//!
7//! ```no_run
8//! let encoder = pocket_encoder::Encoder::new()?;
9//! let vectors = encoder.embed(&["Send a document", "Transfer a file"])?;
10//! assert_eq!(vectors[0].len(), 384);
11//! # Ok::<(), pocket_encoder::Error>(())
12//! ```
13//!
14//! Enable the `search` feature for an exact cosine index with compact binary
15//! storage. Embedding-only users do not compile the index API.
16#[cfg(feature = "search")]
17mod index;
18#[cfg(feature = "search")]
19pub use index::{Index, Neighbor};
20
21use candle_core::{DType, Device, Tensor};
22use candle_nn::VarBuilder;
23use candle_transformers::models::bert::{BertModel, Config};
24use serde::{Deserialize, Serialize};
25use std::{
26    collections::HashMap,
27    sync::{Arc, OnceLock},
28};
29use tokenizers::{Tokenizer, TruncationParams};
30
31/// Embedding coordinate-space identity. Persist this with indexed vectors.
32/// Exact equality, not just equal dimensions, is required before comparison.
33#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
34pub struct ModelInfo {
35    /// Identifies the model and preprocessing contract.
36    pub id: String,
37    /// SHA-256 of the bundled quantized weights.
38    pub revision: String,
39    /// Upstream model repository.
40    pub upstream_model: String,
41    /// Immutable upstream revision.
42    pub upstream_revision: String,
43    /// SHA-256 of the unchanged tokenizer JSON.
44    pub tokenizer_sha256: String,
45    /// Number of normalized F32 coordinates per input.
46    pub dimensions: usize,
47    /// Includes special tokens. Longer inputs are right-truncated.
48    pub max_tokens: usize,
49    /// Pooling and normalization contract.
50    pub pooling: String,
51}
52
53/// Metadata is available without loading weights or doing inference.
54pub fn model_info() -> ModelInfo {
55    let value: serde_json::Value = serde_json::from_str(include_str!("../assets/provenance.json"))
56        .expect("bundled provenance");
57    ModelInfo {
58        id: "pocket-encoder/minilm-l3-q4-g32-v1".into(),
59        revision: format!("sha256:{}", value["bundled_sha256"].as_str().unwrap()),
60        upstream_model: value["upstream_model"].as_str().unwrap().into(),
61        upstream_revision: value["upstream_revision"].as_str().unwrap().into(),
62        tokenizer_sha256: value["tokenizer_sha256"].as_str().unwrap().into(),
63        dimensions: 384,
64        max_tokens: 128,
65        pooling: "attention-masked mean; L2 normalized; no query prefix; right truncation".into(),
66    }
67}
68
69/// Validation or native inference failure. Input text is never included.
70#[derive(Debug, thiserror::Error)]
71pub enum Error {
72    /// Input bounds were exceeded.
73    #[error("{0}")]
74    Input(&'static str),
75    /// Tokenizer/model initialization or inference failed.
76    #[error("encoder failed: {0}")]
77    Inference(String),
78}
79
80struct Inner {
81    model: BertModel,
82    tokenizer: Tokenizer,
83}
84
85/// Cheap clones share immutable weights and tokenizer; CPU only.
86#[derive(Clone)]
87pub struct Encoder {
88    inner: Arc<Inner>,
89}
90
91impl Encoder {
92    /// Load compiled-in weights once per process. No filesystem or network I/O.
93    pub fn new() -> Result<Self, Error> {
94        static MODEL: OnceLock<Result<Arc<Inner>, String>> = OnceLock::new();
95        MODEL
96            .get_or_init(|| load_bundled().map(Arc::new).map_err(|e| e.to_string()))
97            .clone()
98            .map(|inner| Self { inner })
99            .map_err(Error::Inference)
100    }
101
102    /// Encode at most 4096 inputs, each at most 64 KiB, at most 4 MiB combined.
103    /// Inputs are independently right-truncated to 128 WordPiece tokens, including
104    /// special tokens. Use smaller chunks for long documents. Empty text is rejected.
105    /// Returns unit-length vectors in the same order; no input text is retained.
106    pub fn embed(&self, texts: &[impl AsRef<str>]) -> Result<Vec<Vec<f32>>, Error> {
107        if texts.len() > 4096 {
108            return Err(Error::Input("at most 4096 texts per call"));
109        }
110        let mut bytes = 0usize;
111        for text in texts {
112            let text = text.as_ref();
113            if text.trim().is_empty() {
114                return Err(Error::Input("text must not be empty"));
115            }
116            if text.len() > 65536 {
117                return Err(Error::Input("text exceeds 64 KiB"));
118            }
119            bytes = bytes.saturating_add(text.len());
120        }
121        if bytes > 4 * 1024 * 1024 {
122            return Err(Error::Input("texts exceed 4 MiB"));
123        }
124        texts
125            .iter()
126            .map(|text| {
127                self.one(text.as_ref())
128                    .map_err(|e| Error::Inference(e.to_string()))
129            })
130            .collect()
131    }
132
133    fn one(&self, text: &str) -> candle_core::Result<Vec<f32>> {
134        let encoded = self
135            .inner
136            .tokenizer
137            .encode(text, true)
138            .map_err(|e| candle_core::Error::Msg(e.to_string()))?;
139        let n = encoded.len();
140        let ids = Tensor::new(encoded.get_ids(), &Device::Cpu)?.unsqueeze(0)?;
141        let types = Tensor::new(encoded.get_type_ids(), &Device::Cpu)?.unsqueeze(0)?;
142        let mask = Tensor::new(encoded.get_attention_mask(), &Device::Cpu)?.unsqueeze(0)?;
143        let hidden = self.inner.model.forward(&ids, &types, Some(&mask))?;
144        let weights = mask.to_dtype(DType::F32)?.unsqueeze(2)?;
145        let pooled = hidden
146            .broadcast_mul(&weights)?
147            .sum(1)?
148            .squeeze(0)?
149            .to_vec1::<f32>()?;
150        let norm = pooled
151            .iter()
152            .map(|v| (*v as f64).powi(2))
153            .sum::<f64>()
154            .sqrt();
155        if n == 0 || !norm.is_finite() || norm <= 0.0 {
156            candle_core::bail!("non-finite or zero embedding");
157        }
158        Ok(pooled
159            .into_iter()
160            .map(|v| (v as f64 / norm) as f32)
161            .collect())
162    }
163}
164
165fn tokenizer() -> candle_core::Result<Tokenizer> {
166    let mut tokenizer = Tokenizer::from_bytes(include_bytes!("../assets/tokenizer.json"))
167        .map_err(|e| candle_core::Error::Msg(e.to_string()))?;
168    tokenizer.with_padding(None);
169    tokenizer
170        .with_truncation(Some(TruncationParams {
171            max_length: 128,
172            ..Default::default()
173        }))
174        .map_err(|e| candle_core::Error::Msg(e.to_string()))?;
175    Ok(tokenizer)
176}
177
178fn load_bundled() -> candle_core::Result<Inner> {
179    let file = safetensors::SafeTensors::deserialize(include_bytes!("../assets/model.safetensors"))
180        .map_err(|e| candle_core::Error::Msg(e.to_string()))?;
181    let shapes: HashMap<String, Vec<usize>> =
182        serde_json::from_str(include_str!("../assets/shapes.json"))
183            .map_err(|e| candle_core::Error::Msg(e.to_string()))?;
184    let floats = |bytes: &[u8]| {
185        bytes
186            .as_chunks::<4>()
187            .0
188            .iter()
189            .map(|b| f32::from_le_bytes(*b))
190            .collect::<Vec<_>>()
191    };
192    let mut tensors = HashMap::new();
193    for (name, shape) in shapes {
194        let data = if shape.len() < 2 {
195            floats(
196                file.tensor(&name)
197                    .map_err(|e| candle_core::Error::Msg(e.to_string()))?
198                    .data(),
199            )
200        } else {
201            let packed = file
202                .tensor(&format!("{name}.q4"))
203                .map_err(|e| candle_core::Error::Msg(e.to_string()))?;
204            let scales = file
205                .tensor(&format!("{name}.scales"))
206                .map_err(|e| candle_core::Error::Msg(e.to_string()))?;
207            let scales = scales
208                .data()
209                .as_chunks::<2>()
210                .0
211                .iter()
212                .map(|b| half::f16::from_bits(u16::from_le_bytes(*b)).to_f32())
213                .collect::<Vec<_>>();
214            (0..shape.iter().product())
215                .map(|i| {
216                    let nibble = (packed.data()[i / 2] >> ((i % 2) * 4)) & 15;
217                    (nibble as i16 - 8) as f32 * scales[i / 32]
218                })
219                .collect()
220        };
221        tensors.insert(name, Tensor::from_vec(data, shape, &Device::Cpu)?);
222    }
223    let config: Config = serde_json::from_str(include_str!("../assets/config.json"))
224        .map_err(|e| candle_core::Error::Msg(e.to_string()))?;
225    let model = BertModel::load(
226        VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu),
227        &config,
228    )?;
229    Ok(Inner {
230        model,
231        tokenizer: tokenizer()?,
232    })
233}
pocket_encoder/lib.rs

pocket_encoder/
lib.rs
