Skip to main content

cyberbrain_embed/
artefact.rs

1//! Local model artefacts and their verification (SPEC §6: "content-addressed by hash,
2//! verified on every load; a hash mismatch is a hard failure").
3//!
4//! Nothing in this module can reach the network. It takes paths and returns bytes.
5
6use cyberbrain_core::{Error, Result, Slash};
7use serde::{Deserialize, Serialize};
8use std::path::{Path, PathBuf};
9
10/// Length of a blake3 digest rendered as lowercase hex.
11const HEX_LEN: usize = 64;
12
13/// Where the two files of a model2vec-format model live.
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct ModelPaths {
16    /// `model.safetensors`: one tensor named `embeddings`, shape `[vocab, dim]`.
17    pub weights: PathBuf,
18    /// `tokenizer.json` in HuggingFace tokenizers format.
19    pub tokenizer: PathBuf,
20}
21
22impl ModelPaths {
23    /// The conventional layout: `<dir>/model.safetensors` and `<dir>/tokenizer.json`.
24    pub fn in_dir(dir: impl AsRef<Path>) -> Self {
25        let dir = dir.as_ref();
26        Self {
27            weights: dir.join("model.safetensors"),
28            tokenizer: dir.join("tokenizer.json"),
29        }
30    }
31}
32
33/// The expected blake3 digests of both files, lowercase hex. Produced once by whoever
34/// obtained the artefact (the policy crate, after its registered download) and stored in
35/// configuration; checked on every load.
36///
37/// Both files are covered because both change the meaning of a vector: a different
38/// tokenizer maps the same text to different rows of the same matrix.
39#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
40pub struct ArtefactManifest {
41    pub weights_blake3: String,
42    pub tokenizer_blake3: String,
43}
44
45impl ArtefactManifest {
46    /// Rejects anything that is not exactly 64 lowercase hex characters per digest. A
47    /// manifest that cannot match anything is a configuration error, not a load failure.
48    pub fn validate(&self) -> Result<()> {
49        for (what, hex) in [
50            ("weights_blake3", &self.weights_blake3),
51            ("tokenizer_blake3", &self.tokenizer_blake3),
52        ] {
53            if hex.len() != HEX_LEN
54                || !hex
55                    .bytes()
56                    .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
57            {
58                return Err(Error::Config(format!(
59                    "model manifest: {what} must be {HEX_LEN} lowercase hex characters, \
60                     got {:?}",
61                    hex
62                )));
63            }
64        }
65        Ok(())
66    }
67}
68
69/// blake3 of a byte slice, lowercase hex.
70pub fn hash_bytes(bytes: &[u8]) -> String {
71    blake3::hash(bytes).to_hex().to_string()
72}
73
74/// blake3 of a file's contents, lowercase hex. This is what the policy crate calls after a
75/// download to fill an [`ArtefactManifest`].
76pub fn hash_file(path: impl AsRef<Path>) -> Result<String> {
77    let bytes = read(path.as_ref())?;
78    Ok(hash_bytes(&bytes))
79}
80
81pub(crate) fn read(path: &Path) -> Result<Vec<u8>> {
82    std::fs::read(path).map_err(|source| Error::Io {
83        path: path.to_path_buf(),
84        source,
85    })
86}
87
88/// Reads a file and verifies its digest. Returns the bytes only when they match, so the
89/// bytes that were hashed are the bytes that get parsed: no window between check and use.
90pub(crate) fn read_verified(path: &Path, expected_hex: &str, what: &str) -> Result<Vec<u8>> {
91    let bytes = read(path)?;
92    let actual = hash_bytes(&bytes);
93    if actual != expected_hex {
94        return Err(Error::Embed(format!(
95            "{what} artefact {} does not match its manifest: expected blake3 {expected_hex}, \
96             file has {actual}; refusing to load (SPEC §6: a hash mismatch is a hard failure)",
97            Slash(path)
98        )));
99    }
100    Ok(bytes)
101}