cherenkov_model_data/
safetensor.rs

1use crate::*;
2use anyhow::{Context, Result, ensure};
3use safetensors::{Dtype as SafeDtype, tensor::Metadata};
4use serde::Deserialize;
5use serde_json::{Value, json};
6use std::{
7    collections::BTreeMap,
8    path::{Component, Path},
9};
10
11#[derive(Deserialize)]
12struct Index {
13    weight_map: BTreeMap<String, String>,
14}
15
16pub(crate) fn open(dir: &Path) -> Result<Checkpoint> {
17    let index = optional_json(&dir.join("model.safetensors.index.json"))?;
18    let config = optional_json(&dir.join("config.json"))?;
19    let ShardIndex { names, weight_map } = shard_names(&index)?;
20    let paths: Vec<_> = names.iter().map(|name| dir.join(name)).collect();
21    let objects = MappedObjects::open(paths.iter().map(|p| p.as_path()))?;
22    let mut tensors = Vec::new();
23    let mut metadata = Vec::new();
24
25    for (id, name) in names.iter().enumerate() {
26        let (shard, extra) = inspect(&objects, ObjectId(id))?;
27
28        validate_index(&weight_map, &shard, name)?;
29        tensors.extend(shard);
30        metadata.push(extra);
31    }
32
33    if let Some(index) = &weight_map {
34        ensure!(
35            tensors.len() == index.len(),
36            "safetensors index names missing from shards"
37        );
38    }
39
40    tensors.sort_by(|a, b| a.name.cmp(&b.name));
41
42    let checkpoint = Checkpoint {
43        inventory: Inventory {
44            format: ContainerFormat::Safetensors,
45            metadata: json!({"config": config, "index": index, "shards": metadata}),
46            tensors,
47        },
48        objects,
49    };
50
51    checkpoint.validate()?;
52
53    Ok(checkpoint)
54}
55
56/// Read one safetensors header through a byte source. Tensor payloads remain in
57/// that source; successful inspection does not claim their bytes were verified.
58pub fn read(source: &dyn ByteSource, object: ObjectId) -> Result<Inventory> {
59    let (tensors, metadata) = inspect(source, object)?;
60
61    Ok(Inventory {
62        format: ContainerFormat::Safetensors,
63        metadata: json!({"shards": [metadata]}),
64        tensors,
65    })
66}
67
68fn inspect(source: &dyn ByteSource, object: ObjectId) -> Result<(Vec<Tensor>, Value)> {
69    let size = source
70        .objects()
71        .into_iter()
72        .find(|info| info.id == object)
73        .context("unknown safetensors object")?
74        .bytes;
75    let prefix = read_span(source, object, 0, 8)?;
76    let header_len = u64::from_le_bytes(
77        prefix
78            .try_into()
79            .map_err(|_| anyhow::anyhow!("truncated header length"))?,
80    );
81
82    ensure!(
83        header_len <= 64 * 1024 * 1024,
84        "safetensors header exceeds 64 MiB"
85    );
86
87    let base = 8 + header_len;
88
89    ensure!(base <= size, "safetensors header exceeds object size");
90
91    let header = read_span(source, object, 8, header_len)?;
92    // Upstream Metadata validates contiguous offsets, dtype sizes, and shape
93    // overflow without requiring a local mapping of the entire weight file.
94    let metadata: Metadata = serde_json::from_slice(&header).context("safetensors header")?;
95    let end = metadata
96        .tensors()
97        .values()
98        .map(|t| t.data_offsets.1 as u64)
99        .max()
100        .unwrap_or(0);
101
102    ensure!(
103        end.checked_add(base) == Some(size),
104        "safetensors payload size mismatch"
105    );
106
107    let mut tensors = Vec::new();
108
109    for (name, info) in metadata.tensors() {
110        let (start, end) = info.data_offsets;
111        let data = DataSpan {
112            object,
113            offset: base + start as u64,
114            length: (end - start) as u64,
115        };
116        let shape: Vec<_> = info.shape.iter().map(|&n| n as u64).collect();
117        let encoding = match dtype(info.dtype) {
118            Some(dtype) => TensorEncoding::Dense {
119                tensor: StoredTensor::contiguous(dtype, shape.clone(), data)?,
120            },
121            None => TensorEncoding::Opaque {
122                name: info.dtype.to_string(),
123                metadata: json!({"dtype": info.dtype.to_string(), "shape": shape}),
124                data: vec![data],
125            },
126        };
127
128        tensors.push(Tensor::new(
129            name.clone(),
130            TensorRole::Opaque,
131            Some(shape),
132            encoding,
133        ));
134    }
135
136    tensors.sort_by(|a, b| a.name.cmp(&b.name));
137
138    Ok((tensors, serde_json::to_value(metadata.metadata())?))
139}
140
141fn read_span(
142    source: &dyn ByteSource,
143    object: ObjectId,
144    offset: u64,
145    length: u64,
146) -> Result<Vec<u8>> {
147    let span = DataSpan {
148        object,
149        offset,
150        length,
151    };
152    let mut bytes = Vec::new();
153
154    source.read(&span, &mut bytes)?;
155    ensure!(
156        bytes.len() as u64 == length,
157        "byte source returned an incomplete span"
158    );
159
160    Ok(bytes)
161}
162
163fn dtype(dtype: SafeDtype) -> Option<Dtype> {
164    Some(match dtype {
165        SafeDtype::U32 => Dtype::U32,
166        SafeDtype::I64 => Dtype::I64,
167        SafeDtype::F32 => Dtype::F32,
168        SafeDtype::F16 => Dtype::F16,
169        SafeDtype::BF16 => Dtype::Bf16,
170        SafeDtype::BOOL => Dtype::Bool,
171        SafeDtype::U8 => Dtype::U8,
172        SafeDtype::I8 => Dtype::I8,
173        SafeDtype::U16 => Dtype::U16,
174        SafeDtype::I16 => Dtype::I16,
175        SafeDtype::I32 => Dtype::I32,
176        SafeDtype::U64 => Dtype::U64,
177        SafeDtype::F64 => Dtype::F64,
178        _ => return None,
179    })
180}
181
182struct ShardIndex {
183    names: Vec<String>,
184    weight_map: Option<BTreeMap<String, String>>,
185}
186
187fn shard_names(value: &Value) -> Result<ShardIndex> {
188    if value.is_null() {
189        return Ok(ShardIndex {
190            names: vec!["model.safetensors".into()],
191            weight_map: None,
192        });
193    }
194
195    let index: Index = serde_json::from_value(value.clone()).context("safetensors shard index")?;
196
197    ensure!(!index.weight_map.is_empty(), "empty safetensors index");
198
199    let mut names: Vec<_> = index.weight_map.values().cloned().collect();
200
201    names.sort();
202    names.dedup();
203
204    for name in &names {
205        ensure!(
206            !name.is_empty()
207                && Path::new(name)
208                    .components()
209                    .all(|c| matches!(c, Component::Normal(_))),
210            "shard path must stay inside checkpoint: {name}"
211        );
212    }
213
214    Ok(ShardIndex {
215        names,
216        weight_map: Some(index.weight_map),
217    })
218}
219
220fn validate_index(
221    index: &Option<BTreeMap<String, String>>,
222    tensors: &[Tensor],
223    shard: &str,
224) -> Result<()> {
225    let Some(index) = index else {
226        return Ok(());
227    };
228
229    for tensor in tensors {
230        ensure!(
231            index.get(&tensor.name).is_some_and(|file| file == shard),
232            "tensor {} disagrees with shard index",
233            tensor.name
234        );
235    }
236
237    Ok(())
238}
239
240fn optional_json(path: &Path) -> Result<Value> {
241    match std::fs::read(path) {
242        Ok(bytes) => {
243            serde_json::from_slice(&bytes).with_context(|| format!("parsing {}", path.display()))
244        }
245        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(Value::Null),
246        Err(error) => Err(error).with_context(|| format!("reading {}", path.display())),
247    }
248}
cherenkov_model_data/safetensor.rs

cherenkov_model_data/
safetensor.rs
