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

pocket_encoder/
lib.rs
