use std::error::Error;
use std::path::Path;
use tokenizers::Tokenizer;
pub struct Model2Vec {
tok: Tokenizer,
mat: Vec<f32>,
rows: usize,
dim: usize,
}
fn l2_normalize(v: &mut [f32]) {
let mut n = 0.0f32;
for &x in v.iter() {
n += x * x;
}
n = n.sqrt() + 1e-9;
for x in v.iter_mut() {
*x /= n;
}
}
impl Model2Vec {
pub fn load(dir: &Path) -> Result<Model2Vec, Box<dyn Error + Send + Sync>> {
let bytes = std::fs::read(dir.join("potion.f32"))?;
if bytes.len() < 8 {
return Err("potion.f32 too short".into());
}
let rows = u32::from_le_bytes(bytes[0..4].try_into().unwrap()) as usize;
let dim = u32::from_le_bytes(bytes[4..8].try_into().unwrap()) as usize;
let need = rows * dim * 4;
if bytes.len() < 8 + need {
return Err(format!("potion.f32 truncated: want {} floats", rows * dim).into());
}
let mut mat = Vec::with_capacity(rows * dim);
for chunk in bytes[8..8 + need].chunks_exact(4) {
mat.push(f32::from_le_bytes(chunk.try_into().unwrap()));
}
let tok = Tokenizer::from_file(dir.join("tokenizer.json"))?;
Ok(Model2Vec { tok, mat, rows, dim })
}
pub fn dim(&self) -> usize {
self.dim
}
pub fn rows(&self) -> usize {
self.rows
}
pub fn embed(&self, text: &str) -> Option<Vec<f32>> {
let enc = self.tok.encode(text, false).ok()?;
let mut acc = vec![0.0f32; self.dim];
let mut n = 0usize;
for &id in enc.get_ids() {
let id = id as usize;
if id >= self.rows {
continue;
}
let off = id * self.dim;
let row = &self.mat[off..off + self.dim];
for i in 0..self.dim {
acc[i] += row[i];
}
n += 1;
}
if n == 0 {
return None;
}
for x in acc.iter_mut() {
*x /= n as f32;
}
l2_normalize(&mut acc);
Some(acc)
}
pub fn token_count(&self, text: &str) -> usize {
self.tok.encode(text, false).map(|e| e.get_ids().len()).unwrap_or(0)
}
}
pub fn cosine(a: &[f32], b: &[f32]) -> f32 {
let mut s = 0.0f32;
for i in 0..a.len().min(b.len()) {
s += a[i] * b[i];
}
s
}