DarkRoom/core/dr-face/src/embedding.rs

//! What an embedder produces, and how it is stored (docs/faces.md §6).
//!
//! Deliberately **model-free**: the vector, its identity, its comparison and
//! its storage encoding are arithmetic, and `calibrate` and `cluster` are built
//! on them. Keeping them out of the `inference` feature is what lets the part
//! of this subsystem most likely to be subtly wrong be tested on a machine with
//! no weights on it.
//!
//! [`crate::embed::Embedder`] is the thing that needs a model, and it lives
//! behind the feature.

/// Embedding dimensionality. Fixed by the model family, not a parameter.
pub const EMBEDDING_DIM: usize = 512;

/// Which model produced an embedding.
///
/// Embeddings from different models are not comparable, and this is the one
/// mistake that produces plausible-looking garbage rather than an error — so
/// the id travels *with* the vector rather than beside it, and
/// [`Embedding::cosine`] refuses a cross-model comparison.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ModelId(pub std::sync::Arc<str>);

impl ModelId {
    pub fn new(s: impl Into<std::sync::Arc<str>>) -> Self {
        Self(s.into())
    }
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl std::fmt::Display for ModelId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0)
    }
}

/// A 512-d L2-normalised face embedding.
#[derive(Debug, Clone, PartialEq)]
pub struct Embedding {
    pub model: ModelId,
    pub v: Box<[f32; EMBEDDING_DIM]>,
}

impl Embedding {
    /// Cosine similarity, which for unit vectors is the plain dot product.
    ///
    /// `None` when the two came from different models. That is a real
    /// possibility in a library indexed across a model upgrade, and the
    /// alternative — returning a number — is the failure mode
    /// `faces.model_id` exists to prevent.
    pub fn cosine(&self, other: &Embedding) -> Option<f32> {
        if self.model != other.model {
            return None;
        }
        Some(dot(&self.v, &other.v))
    }

    /// Storage form: `512 × f16`, 1 KB per face (catalog.md §10.1).
    pub fn to_f16_bytes(&self) -> Vec<u8> {
        let mut out = Vec::with_capacity(EMBEDDING_DIM * 2);
        for &x in self.v.iter() {
            out.extend_from_slice(&f32_to_f16_bits(x).to_le_bytes());
        }
        out
    }

    /// Read back from storage, re-normalising.
    ///
    /// The f16 round-trip perturbs a unit vector by ~1e-3 in cosine — three
    /// orders below the separation between a match and a non-match — but the
    /// drift is free to remove and invisible if left, so it is removed here
    /// rather than remembered at every call site.
    pub fn from_f16_bytes(model: ModelId, bytes: &[u8]) -> Option<Self> {
        if bytes.len() != EMBEDDING_DIM * 2 {
            return None;
        }
        let mut v = Box::new([0.0_f32; EMBEDDING_DIM]);
        for (i, chunk) in bytes.chunks_exact(2).enumerate() {
            v[i] = f16_bits_to_f32(u16::from_le_bytes([chunk[0], chunk[1]]));
        }
        normalise(&mut v);
        Some(Self { model, v })
    }
}

fn dot(a: &[f32; EMBEDDING_DIM], b: &[f32; EMBEDDING_DIM]) -> f32 {
    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
}

pub(crate) fn normalise(v: &mut [f32; EMBEDDING_DIM]) {
    // Clamped rather than checked: a zero-norm embedding is a broken model,
    // not a runtime condition worth an error path, and dividing by 1e-6 keeps
    // the NaN out of the catalog.
    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt().max(1e-6);
    for x in v.iter_mut() {
        *x /= norm;
    }
}

// ── f16 ───────────────────────────────────────────────────────────────────
//
// Hand-rolled rather than pulling in `half`: two functions over a format that
// has not changed since 2008, used at exactly one boundary. The dependency
// policy (D13, D1) makes the bar for a new crate high, and this is well under
// it.

fn f32_to_f16_bits(x: f32) -> u16 {
    let bits = x.to_bits();
    let sign = ((bits >> 16) & 0x8000) as u16;
    let exp = ((bits >> 23) & 0xff) as i32 - 127 + 15;
    let mant = bits & 0x007f_ffff;

    if exp >= 0x1f {
        // Overflow, inf, or NaN. Embeddings are unit-norm so this is the
        // broken-model path; infinity is the honest answer, not a clamp that
        // hides it.
        return sign
            | 0x7c00
            | if mant != 0 && exp == 0x1f + 112 {
                0x200
            } else {
                0
            };
    }
    if exp <= 0 {
        // Subnormal or underflow. A component of a unit 512-vector is ~0.04,
        // nowhere near here, so this branch exists for correctness rather than
        // for traffic.
        if exp < -10 {
            return sign;
        }
        let mant = mant | 0x0080_0000;
        let shift = (14 - exp) as u32;
        let half = (mant >> shift) as u16;
        // Round to nearest, ties to even.
        let rem = mant & ((1 << shift) - 1);
        let tie = 1 << (shift - 1);
        let round = u16::from(rem > tie || (rem == tie && (half & 1) == 1));
        return sign | (half + round);
    }

    let half = ((exp as u16) << 10) | (mant >> 13) as u16;
    let rem = mant & 0x1fff;
    let round = u16::from(rem > 0x1000 || (rem == 0x1000 && (half & 1) == 1));
    sign | (half + round)
}

fn f16_bits_to_f32(h: u16) -> f32 {
    let sign = ((h & 0x8000) as u32) << 16;
    let exp = ((h >> 10) & 0x1f) as u32;
    let mant = (h & 0x03ff) as u32;

    if exp == 0 {
        if mant == 0 {
            return f32::from_bits(sign);
        }
        // Subnormal: renormalise into f32's range.
        let mut e = -1_i32;
        let mut m = mant;
        while m & 0x0400 == 0 {
            m <<= 1;
            e -= 1;
        }
        let m = m & 0x03ff;
        return f32::from_bits(sign | (((127 - 15 + 1 + e) as u32) << 23) | (m << 13));
    }
    if exp == 0x1f {
        return f32::from_bits(sign | 0x7f80_0000 | (mant << 13));
    }
    f32::from_bits(sign | ((exp + 127 - 15) << 23) | (mant << 13))
}

#[cfg(test)]
mod tests {
    use super::*;

    fn unit(seed: u32) -> Embedding {
        let mut v = Box::new([0.0_f32; EMBEDDING_DIM]);
        let mut s = seed.wrapping_mul(2_654_435_761).wrapping_add(1);
        for x in v.iter_mut() {
            s = s.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
            *x = (s >> 8) as f32 / (1u32 << 23) as f32 - 0.5;
        }
        normalise(&mut v);
        Embedding {
            model: ModelId::new("test"),
            v,
        }
    }

    #[test]
    fn a_normalised_embedding_has_cosine_one_with_itself() {
        let e = unit(7);
        assert!((e.cosine(&e).unwrap() - 1.0).abs() < 1e-5);
    }

    #[test]
    fn embeddings_from_different_models_do_not_compare() {
        let a = unit(1);
        let mut b = unit(1);
        b.model = ModelId::new("other");
        assert_eq!(
            a.cosine(&b),
            None,
            "a cross-model cosine must not be a number"
        );
    }

    /// The claim docs/faces.md §6 makes about the storage format: the f16
    /// round-trip costs ~1e-3 of cosine, three orders below the separation
    /// between a match and a non-match.
    #[test]
    fn f16_round_trip_preserves_the_embedding() {
        for seed in 0..16 {
            let e = unit(seed);
            let back = Embedding::from_f16_bytes(e.model.clone(), &e.to_f16_bytes()).unwrap();
            let cos = e.cosine(&back).unwrap();
            assert!(cos > 0.9999, "seed {seed}: round-trip cosine {cos}");
        }
    }

    #[test]
    fn f16_round_trip_rejects_a_wrong_length_blob() {
        assert!(Embedding::from_f16_bytes(ModelId::new("m"), &[0u8; 100]).is_none());
    }

    #[test]
    fn f16_handles_the_values_an_embedding_actually_contains() {
        // Components of a unit 512-vector cluster around ±1/sqrt(512) ≈ 0.044.
        for &x in &[0.0_f32, 1.0, -1.0, 0.044_194_17, -0.044_194_17, 1e-3, -7e-4] {
            let back = f16_bits_to_f32(f32_to_f16_bits(x));
            assert!(
                (back - x).abs() <= 1e-3 * x.abs().max(1e-3),
                "{x} round-tripped to {back}"
            );
        }
    }
}