steeldb/text/
model2vec.rs

1//! model2vec static embedder — the MULTILINGUAL (EN/JA/KO) span embedder, ported from the TS
2//! `Model2VecEmbedder`. No forward pass: tokenize (no special tokens), look up each subword id in the
3//! packed `potion.f32` matrix, mean-pool, L2-normalise. CPU-only, language-agnostic — the same
4//! artifact the ingest pipeline already ships.
5//!
6//! `potion.f32` layout: u32 LE `rows`, u32 LE `dim`, then `rows*dim` little-endian f32 (row-major).
7
8use std::error::Error;
9use std::path::Path;
10use tokenizers::Tokenizer;
11
12pub struct Model2Vec {
13    tok: Tokenizer,
14    mat: Vec<f32>,
15    rows: usize,
16    dim: usize,
17}
18
19fn l2_normalize(v: &mut [f32]) {
20    let mut n = 0.0f32;
21    for &x in v.iter() {
22        n += x * x;
23    }
24    n = n.sqrt() + 1e-9;
25    for x in v.iter_mut() {
26        *x /= n;
27    }
28}
29
30impl Model2Vec {
31    /// `dir` holds `potion.f32` and `tokenizer.json` (e.g. step0_bundle/model2vec).
32    pub fn load(dir: &Path) -> Result<Model2Vec, Box<dyn Error + Send + Sync>> {
33        let bytes = std::fs::read(dir.join("potion.f32"))?;
34        if bytes.len() < 8 {
35            return Err("potion.f32 too short".into());
36        }
37        let rows = u32::from_le_bytes(bytes[0..4].try_into().unwrap()) as usize;
38        let dim = u32::from_le_bytes(bytes[4..8].try_into().unwrap()) as usize;
39        let need = rows * dim * 4;
40        if bytes.len() < 8 + need {
41            return Err(format!("potion.f32 truncated: want {} floats", rows * dim).into());
42        }
43        let mut mat = Vec::with_capacity(rows * dim);
44        for chunk in bytes[8..8 + need].chunks_exact(4) {
45            mat.push(f32::from_le_bytes(chunk.try_into().unwrap()));
46        }
47        let tok = Tokenizer::from_file(dir.join("tokenizer.json"))?;
48        Ok(Model2Vec { tok, mat, rows, dim })
49    }
50
51    pub fn dim(&self) -> usize {
52        self.dim
53    }
54    pub fn rows(&self) -> usize {
55        self.rows
56    }
57
58    /// Mean-pooled, L2-normalised embedding of `text`; None if no in-vocab tokens.
59    pub fn embed(&self, text: &str) -> Option<Vec<f32>> {
60        let enc = self.tok.encode(text, false).ok()?;
61        let mut acc = vec![0.0f32; self.dim];
62        let mut n = 0usize;
63        for &id in enc.get_ids() {
64            let id = id as usize;
65            if id >= self.rows {
66                continue;
67            }
68            let off = id * self.dim;
69            let row = &self.mat[off..off + self.dim];
70            for i in 0..self.dim {
71                acc[i] += row[i];
72            }
73            n += 1;
74        }
75        if n == 0 {
76            return None;
77        }
78        for x in acc.iter_mut() {
79            *x /= n as f32;
80        }
81        l2_normalize(&mut acc);
82        Some(acc)
83    }
84
85    /// Number of subword tokens `text` produces (diagnostic — how well the tokenizer covers a script).
86    pub fn token_count(&self, text: &str) -> usize {
87        self.tok.encode(text, false).map(|e| e.get_ids().len()).unwrap_or(0)
88    }
89}
90
91pub fn cosine(a: &[f32], b: &[f32]) -> f32 {
92    let mut s = 0.0f32;
93    for i in 0..a.len().min(b.len()) {
94        s += a[i] * b[i];
95    }
96    s
97}
steeldb/text/model2vec.rs

steeldb/text/
model2vec.rs
