cherenkov_model_data/
lib.rs1pub 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#[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#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct Inventory {
32 pub format: ContainerFormat,
34 pub metadata: serde_json::Value,
36 pub tensors: Vec<Tensor>,
38}
39
40pub struct Checkpoint {
42 pub inventory: Inventory,
44 pub objects: MappedObjects,
46}
47
48impl Checkpoint {
49 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 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}