seyk_embed/quant/
mod.rs

1//! Offline Photo Anchor 4-bit codec for already normalized photo embeddings.
2//!
3//! Tables: Pastel-Cloud OÜ, CC BY 4.0; see the bundled attribution.
4//! This codec does not change Seyk's default f32 storage or search behavior.
5use crate::Embedding;
6
7mod sha256;
8#[cfg(test)]
9mod tests;
10
11/// Exact upstream revision of all bundled Photo Anchor assets.
12pub const REVISION: &str = "be9e066984345c92424c86bfcce3b17f3e3e8484";
13/// Attribution must accompany packages and binaries that bundle these tables.
14pub const ATTRIBUTION: &str = include_str!("../../assets/photo-anchor/NOTICE.md");
15const MANIFEST: &[u8] = include_bytes!("../../assets/photo-anchor/manifest.json");
16const MANIFEST_SHA: &str = "4dbde655fdd44a1feb4b010685fe065ffbe9c1f996331b4f4a40c9a5b2580c1e";
17
18/// Invalid inputs or corrupted/unrecognized calibration assets fail closed.
19#[derive(Debug, thiserror::Error, PartialEq, Eq)]
20pub enum Error {
21    #[error("Photo Anchor tables are not bundled for dimension {0}")]
22    UnsupportedDimension(usize),
23    #[error("Photo Anchor asset integrity failure: {0}")]
24    AssetIntegrity(&'static str),
25    #[error("expected {expected} values/bytes, got {actual}")]
26    Length { expected: usize, actual: usize },
27    #[error("Photo Anchor requires a finite unit embedding")]
28    NonUnitEmbedding,
29}
30
31/// Dimension-specific photo mean and 4-bit calibration, loaded without I/O.
32pub struct PhotoAnchorQ4 {
33    anchor: Vec<f32>,
34    offset: Vec<f32>,
35    scale: Vec<f32>,
36}
37
38impl PhotoAnchorQ4 {
39    /// Load and validate the hash-pinned 256- or 768-dimensional tables.
40    pub fn new(dim: usize) -> Result<Self, Error> {
41        let (anchor, table, anchor_sha, table_sha) = match dim {
42            256 => (
43                include_bytes!("../../assets/photo-anchor/photo_anchor_256.f32.bin").as_slice(),
44                include_bytes!("../../assets/photo-anchor/quant4_256.f32.bin").as_slice(),
45                "4e30ed052129ee745bb810fc1673e360857557e4a97daffdb2734e2b8dc5dc28",
46                "fe029c4cf510f6a628a24bcdc3db097ce8c387e948080d5065ef81ff904a2353",
47            ),
48            768 => (
49                include_bytes!("../../assets/photo-anchor/photo_anchor_768.f32.bin").as_slice(),
50                include_bytes!("../../assets/photo-anchor/quant4_768.f32.bin").as_slice(),
51                "a3978221be716a22f0d738efdc93df0de90fb8c490339011422ab15c2e8a3bb6",
52                "5088d57dbb8ae53da9f45f94ce20aa9177160fcbc499d41a8ecb755270ba5560",
53            ),
54            _ => return Err(Error::UnsupportedDimension(dim)),
55        };
56        Self::load(dim, MANIFEST, anchor, table, anchor_sha, table_sha)
57    }
58
59    fn load(
60        dim: usize,
61        manifest: &[u8],
62        anchor: &[u8],
63        table: &[u8],
64        anchor_sha: &str,
65        table_sha: &str,
66    ) -> Result<Self, Error> {
67        // Exact manifest identity pins the model revision, methods and bit order,
68        // without a permissive JSON parser or a runtime download fallback.
69        check_hash(manifest, MANIFEST_SHA, "manifest")?;
70        let manifest =
71            std::str::from_utf8(manifest).map_err(|_| Error::AssetIntegrity("manifest"))?;
72        for expected in [
73            "google/embeddinggemma-2",
74            "914f7f89142e33e77833254d9c9b90c3cef7303b",
75            anchor_sha,
76            table_sha,
77        ] {
78            if !manifest.contains(expected) {
79                return Err(Error::AssetIntegrity("manifest metadata"));
80            }
81        }
82        check_hash(anchor, anchor_sha, "anchor")?;
83        check_hash(table, table_sha, "calibration")?;
84        check_len(anchor.len(), dim * 4)?;
85        check_len(table.len(), dim * 8)?;
86        let floats = |bytes: &[u8]| {
87            bytes
88                .as_chunks::<4>()
89                .0
90                .iter()
91                .map(|b| f32::from_le_bytes(*b))
92                .collect::<Vec<_>>()
93        };
94        let anchor = floats(anchor);
95        let offset = floats(&table[..dim * 4]);
96        let scale = floats(&table[dim * 4..]);
97        if anchor
98            .iter()
99            .chain(&offset)
100            .chain(&scale)
101            .any(|x| !x.is_finite())
102            || scale.iter().any(|&x| x <= 0.0)
103        {
104            return Err(Error::AssetIntegrity(
105                "non-finite values or nonpositive scale",
106            ));
107        }
108        Ok(Self {
109            anchor,
110            offset,
111            scale,
112        })
113    }
114
115    pub fn dim(&self) -> usize {
116        self.anchor.len()
117    }
118
119    /// Subtract the anchor, clip, round ties to even, and pack high nibble first.
120    /// Rejects wrong dimensions and zero/non-unit deserialized embeddings.
121    pub fn encode(&self, embedding: &Embedding) -> Result<Vec<u8>, Error> {
122        check_embedding(embedding, self.dim())?;
123        Ok(embedding
124            .as_slice()
125            .as_chunks::<2>()
126            .0
127            .iter()
128            .enumerate()
129            .map(|(j, pair)| {
130                let code = |i: usize, x: f32| {
131                    let centered = (x - self.anchor[i])
132                        .clamp(self.offset[i], self.offset[i] + 15.0 * self.scale[i]);
133                    ((centered - self.offset[i]) / self.scale[i])
134                        .round_ties_even()
135                        .clamp(0.0, 15.0) as u8
136                };
137                (code(2 * j, pair[0]) << 4) | code(2 * j + 1, pair[1])
138            })
139            .collect())
140    }
141}
142
143/// Asymmetric query scoring: the query remains f32 and is never centered.
144/// Reconstruction is not renormalized, so scores are approximate dot products.
145pub struct QueryLut {
146    bias: f32,
147    pairs: Vec<[f32; 256]>,
148}
149
150impl QueryLut {
151    /// Validate the query and build one 256-entry table for each packed byte.
152    /// Build once per query; reuse across all photo codes of the same dimension.
153    pub fn new(query: &Embedding) -> Result<Self, Error> {
154        let codec = PhotoAnchorQ4::new(query.dim())?;
155        check_embedding(query, codec.dim())?;
156        let q = query.as_slice();
157        let bias = (0..codec.dim())
158            .map(|i| (codec.anchor[i] + codec.offset[i]) * q[i])
159            .sum();
160        let pairs = (0..codec.dim())
161            .step_by(2)
162            .map(|i| {
163                let (hi, lo) = (codec.scale[i] * q[i], codec.scale[i + 1] * q[i + 1]);
164                std::array::from_fn(|byte| (byte >> 4) as f32 * hi + (byte & 15) as f32 * lo)
165            })
166            .collect();
167        Ok(Self { bias, pairs })
168    }
169
170    /// Score exactly one packed vector, rejecting truncated or oversized codes.
171    pub fn score(&self, codes: &[u8]) -> Result<f32, Error> {
172        check_len(codes.len(), self.pairs.len())?;
173        Ok(self.bias
174            + self
175                .pairs
176                .iter()
177                .zip(codes)
178                .map(|(table, &code)| table[code as usize])
179                .sum::<f32>())
180    }
181}
182
183fn check_len(actual: usize, expected: usize) -> Result<(), Error> {
184    if actual == expected {
185        Ok(())
186    } else {
187        Err(Error::Length { expected, actual })
188    }
189}
190
191fn check_embedding(embedding: &Embedding, dim: usize) -> Result<(), Error> {
192    check_len(embedding.dim(), dim)?;
193    if embedding.normalized() {
194        Ok(())
195    } else {
196        Err(Error::NonUnitEmbedding)
197    }
198}
199
200fn check_hash(bytes: &[u8], expected: &str, name: &'static str) -> Result<(), Error> {
201    if sha256::hex(bytes) == expected {
202        Ok(())
203    } else {
204        Err(Error::AssetIntegrity(name))
205    }
206}
seyk_embed/quant/mod.rs

seyk_embed/quant/
mod.rs
