Skip to main content

cherenkov_model_data/
lib.rs

1//! Container metadata and lazy byte access, independent of model execution.
2//!
3//! Readers preserve names and metadata. Architecture adapters assign tensor
4//! roles and interpret quantization conventions after reading the container.
5
6pub mod discovery;
7mod gguf;
8pub mod mlx;
9mod safetensor;
10pub use safetensor::read as read_safetensors;
11mod source;
12mod tensor;
13
14use anyhow::{Result, ensure};
15use serde::{Deserialize, Serialize};
16pub use source::{ByteSource, MappedBytes, MappedObjects, ObjectInfo};
17use std::path::Path;
18pub use tensor::*;
19
20/// Container format identified before applying model-specific conventions.
21#[derive(Debug, Clone, Serialize, Deserialize)]
22#[serde(tag = "kind", rename_all = "snake_case")]
23pub enum ContainerFormat {
24    Safetensors,
25    Gguf { version: u32 },
26    Opaque { name: String, version: Option<u32> },
27}
28
29/// Container metadata and tensor descriptors whose bytes remain in a separate source.
30#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct Inventory {
32    /// Container format and any recognized version.
33    pub format: ContainerFormat,
34    /// Container metadata, including configuration when available.
35    pub metadata: serde_json::Value,
36    /// Tensor descriptors whose object IDs refer to the accompanying source.
37    pub tensors: Vec<Tensor>,
38}
39
40/// A local container inventory with the mappings backing its tensor data.
41pub struct Checkpoint {
42    /// Metadata read from the container, before architecture adaptation.
43    pub inventory: Inventory,
44    /// Immutable backing files for the inventory's data spans.
45    pub objects: MappedObjects,
46}
47
48impl Checkpoint {
49    /// Open a safetensors folder/file or a GGUF file. Files must remain
50    /// unchanged while the checkpoint or any mapped view is alive.
51    pub fn open(path: &Path) -> Result<Self> {
52        if path.is_dir() {
53            return safetensor::open(path);
54        }
55
56        let objects = MappedObjects::open([path])?;
57        let bytes = objects.mapping(ObjectId(0))?;
58        let inventory = if bytes.starts_with(b"GGUF") {
59            gguf::inspect(bytes)?
60        } else {
61            safetensor::read(&objects, ObjectId(0))?
62        };
63        let checkpoint = Self { inventory, objects };
64
65        checkpoint.validate()?;
66
67        Ok(checkpoint)
68    }
69
70    /// Check unique tensor names, tensor layouts, and bounds within backing files.
71    /// This does not verify tensor values or architecture compatibility.
72    pub fn validate(&self) -> Result<()> {
73        let mut names = std::collections::HashSet::new();
74
75        for tensor in &self.inventory.tensors {
76            ensure!(
77                names.insert(&tensor.name),
78                "duplicate tensor {}",
79                tensor.name
80            );
81            tensor.validate()?;
82
83            for span in tensor.encoding.data() {
84                self.objects.bytes(span)?;
85            }
86        }
87
88        Ok(())
89    }
90}