hypersteeldb 0.2.4

//! model2vec static embedder — the MULTILINGUAL (EN/JA/KO) span embedder, ported from the TS
//! `Model2VecEmbedder`. No forward pass: tokenize (no special tokens), look up each subword id in the
//! packed `potion.f32` matrix, mean-pool, L2-normalise. CPU-only, language-agnostic — the same
//! artifact the ingest pipeline already ships.
//!
//! `potion.f32` layout: u32 LE `rows`, u32 LE `dim`, then `rows*dim` little-endian f32 (row-major).

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 {
    /// `dir` holds `potion.f32` and `tokenizer.json` (e.g. step0_bundle/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
    }

    /// Mean-pooled, L2-normalised embedding of `text`; None if no in-vocab tokens.
    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)
    }

    /// Number of subword tokens `text` produces (diagnostic — how well the tokenizer covers a script).
    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
}