nuvec_io/
safetensors.rs

1//! The safetensors format: an 8-byte little-endian header size, a JSON header mapping
2//! tensor names to their dtype, shape and byte range, then the tensors' bytes.
3//!
4//! Loading maps the file without a copy (see [`MappedFile`]): every tensor is a view of the
5//! mapped bytes, read from disk (or the page cache) when first used. Files whose tensors are
6//! not aligned to their element sizes (MLX writes headers without padding) are copied once,
7//! as a whole, and their tensors are views of the copy.
8
9use std::collections::BTreeMap;
10use std::path::Path;
11
12use nuvec_core::{Array, DType, Error, MappedFile, Result};
13
14use crate::json::{self, Json};
15
16fn invalid(msg: impl Into<String>) -> Error {
17    Error::InvalidArgument(format!("safetensors: {}", msg.into()))
18}
19
20/// The tensors and metadata of a safetensors file.
21#[derive(Debug, Clone, Default)]
22pub struct SafeTensors {
23    /// The tensors by name.
24    pub tensors: BTreeMap<String, Array>,
25    /// The `__metadata__` strings.
26    pub metadata: BTreeMap<String, String>,
27}
28
29/// safetensors' dtype names, with MLX's `C64` for complex64.
30fn dtype_of(name: &str) -> Option<DType> {
31    Some(match name {
32        "BOOL" => DType::Bool,
33        "U8" => DType::U8,
34        "U16" => DType::U16,
35        "U32" => DType::U32,
36        "U64" => DType::U64,
37        "I8" => DType::I8,
38        "I16" => DType::I16,
39        "I32" => DType::I32,
40        "I64" => DType::I64,
41        "F16" => DType::F16,
42        "BF16" => DType::BF16,
43        "F32" => DType::F32,
44        "F64" => DType::F64,
45        "C64" => DType::Complex64,
46        _ => return None,
47    })
48}
49
50fn name_of(dtype: DType) -> &'static str {
51    match dtype {
52        DType::Bool => "BOOL",
53        DType::U8 => "U8",
54        DType::U16 => "U16",
55        DType::U32 => "U32",
56        DType::U64 => "U64",
57        DType::I8 => "I8",
58        DType::I16 => "I16",
59        DType::I32 => "I32",
60        DType::I64 => "I64",
61        DType::F16 => "F16",
62        DType::BF16 => "BF16",
63        DType::F32 => "F32",
64        DType::F64 => "F64",
65        _ => "C64",
66    }
67}
68
69fn usize_of(v: &Json, what: &str) -> Result<usize> {
70    match v {
71        Json::Int(n) if *n >= 0 => Ok(*n as usize),
72        _ => Err(invalid(format!("{what} must be a non-negative integer"))),
73    }
74}
75
76/// The reference implementation's limit on the header size.
77const MAX_HEADER: usize = 100_000_000;
78
79/// Loads a safetensors file. The tensors are views of the memory-mapped file (or of one
80/// copy of its data, if they are not aligned to their element sizes).
81///
82/// Validation follows the reference implementation: a header of at most 100 MB, an
83/// optional (or `null`) `__metadata__` of strings, and tensors whose byte ranges, sorted
84/// by offset, tile the data from its start to the end of the file.
85pub fn load_safetensors(path: impl AsRef<Path>) -> Result<SafeTensors> {
86    let file = MappedFile::open(path)?;
87    let bytes = file.bytes();
88    let size = bytes
89        .get(..8)
90        .map(|b| u64::from_le_bytes(b.try_into().expect("8 bytes")))
91        .ok_or_else(|| invalid("file shorter than its header size"))?;
92    let start = usize::try_from(size)
93        .ok()
94        .filter(|&n| n <= MAX_HEADER)
95        .and_then(|n| n.checked_add(8))
96        .filter(|&end| end <= bytes.len())
97        .ok_or_else(|| invalid("header larger than the file or 100 MB"))?;
98    let Json::Object(entries) = json::parse(&bytes[8..start])? else {
99        return Err(invalid("the header must be a JSON object"));
100    };
101    let mut out = SafeTensors::default();
102    let mut ranges = Vec::new();
103    for (name, entry) in entries {
104        if name == "__metadata__" {
105            match entry {
106                Json::Null => {}
107                Json::Object(fields) => {
108                    for (k, v) in fields {
109                        let Json::Str(s) = v else {
110                            return Err(invalid("metadata values must be strings"));
111                        };
112                        out.metadata.insert(k, s);
113                    }
114                }
115                _ => return Err(invalid("`__metadata__` must be an object")),
116            }
117            continue;
118        }
119        let Json::Object(fields) = entry else {
120            return Err(invalid(format!("entry `{name}` must be an object")));
121        };
122        let field = |key: &str| fields.iter().find(|(k, _)| k == key).map(|(_, v)| v);
123        let dtype = match field("dtype") {
124            Some(Json::Str(d)) => {
125                dtype_of(d).ok_or_else(|| invalid(format!("`{name}`: unsupported dtype {d}")))?
126            }
127            _ => return Err(invalid(format!("`{name}`: missing dtype"))),
128        };
129        let shape = match field("shape") {
130            Some(Json::Array(dims)) => dims
131                .iter()
132                .map(|d| usize_of(d, "a dimension"))
133                .collect::<Result<Vec<_>>>()?,
134            _ => return Err(invalid(format!("`{name}`: missing shape"))),
135        };
136        let (begin, end) = match field("data_offsets") {
137            Some(Json::Array(o)) if o.len() == 2 => {
138                (usize_of(&o[0], "an offset")?, usize_of(&o[1], "an offset")?)
139            }
140            _ => return Err(invalid(format!("`{name}`: missing data_offsets"))),
141        };
142        let expected = shape
143            .iter()
144            .try_fold(dtype.size(), |n, &d| n.checked_mul(d))
145            .ok_or_else(|| invalid(format!("`{name}`: size overflows")))?;
146        if end.checked_sub(begin) != Some(expected) {
147            return Err(invalid(format!(
148                "`{name}`: data_offsets [{begin}, {end}] do not hold {expected} bytes"
149            )));
150        }
151        ranges.push((begin, end, name, dtype, shape));
152    }
153    // The ranges must tile the data exactly.
154    ranges.sort_by_key(|r| (r.0, r.1));
155    let mut at = 0;
156    for (begin, end, name, ..) in &ranges {
157        if *begin != at {
158            return Err(invalid(format!(
159                "`{name}`: data_offsets overlap or leave a gap"
160            )));
161        }
162        at = *end;
163    }
164    if start + at != bytes.len() {
165        return Err(invalid(
166            "the tensors do not cover the data to the end of the file",
167        ));
168    }
169    // A data section not aligned for every tensor (MLX writes headers without padding) is
170    // copied once, as a whole: views of one aligned copy, without keeping the mapping
171    // alive besides per-tensor copies (which would hold the checkpoint twice).
172    let aligned = ranges
173        .iter()
174        .all(|(begin, _, _, dtype, _)| (start + begin).is_multiple_of(dtype.size()));
175    let (data, base) = match aligned {
176        true => (file, start),
177        false => (file.section(start, at)?, 0),
178    };
179    for (begin, _, name, dtype, shape) in ranges {
180        out.tensors
181            .insert(name, data.array(base + begin, dtype, &shape)?);
182    }
183    Ok(out)
184}
185
186/// The array's elements as little-endian bytes, row-major.
187pub(crate) fn array_bytes(a: &Array) -> Result<Vec<u8>> {
188    macro_rules! le {
189        ($t:ty) => {
190            a.to_vec::<$t>()?
191                .iter()
192                .flat_map(|v| v.to_le_bytes())
193                .collect()
194        };
195    }
196    Ok(match a.dtype() {
197        DType::Bool => a.to_vec::<bool>()?.iter().map(|&b| u8::from(b)).collect(),
198        DType::U8 => a.to_vec::<u8>()?,
199        DType::U16 => le!(u16),
200        DType::U32 => le!(u32),
201        DType::U64 => le!(u64),
202        DType::I8 => le!(i8),
203        DType::I16 => le!(i16),
204        DType::I32 => le!(i32),
205        DType::I64 => le!(i64),
206        DType::F16 => le!(nuvec_core::f16),
207        DType::BF16 => le!(nuvec_core::bf16),
208        DType::F32 => le!(f32),
209        DType::F64 => le!(f64),
210        _ => a
211            .to_vec::<nuvec_core::Complex64>()?
212            .iter()
213            .flat_map(|z| [z.re.to_le_bytes(), z.im.to_le_bytes()].concat())
214            .collect(),
215    })
216}
217
218/// Saves `tensors` (evaluating them) and `metadata` as a safetensors file.
219///
220/// As the reference implementation, tensors are laid out by decreasing element size (then
221/// by name) after a header padded to 8 bytes, so every tensor is aligned to its element
222/// size and loads without a copy.
223pub fn save_safetensors(
224    path: impl AsRef<Path>,
225    tensors: &[(&str, &Array)],
226    metadata: &[(&str, &str)],
227) -> Result<()> {
228    let mut sorted: Vec<&(&str, &Array)> = tensors.iter().collect();
229    sorted.sort_by(|(ln, la), (rn, ra)| ra.dtype().size().cmp(&la.dtype().size()).then(ln.cmp(rn)));
230    let mut header = Vec::new();
231    if !metadata.is_empty() {
232        let map: BTreeMap<String, String> = metadata
233            .iter()
234            .map(|(k, v)| (k.to_string(), v.to_string()))
235            .collect();
236        header.push(format!("\"__metadata__\":{}", json::object(&map)));
237    }
238    let mut data = Vec::new();
239    for (name, a) in sorted {
240        let bytes = array_bytes(a)?;
241        let shape: Vec<String> = a.shape().iter().map(usize::to_string).collect();
242        header.push(format!(
243            "{}:{{\"dtype\":\"{}\",\"shape\":[{}],\"data_offsets\":[{},{}]}}",
244            json::quote(name),
245            name_of(a.dtype()),
246            shape.join(","),
247            data.len(),
248            data.len() + bytes.len()
249        ));
250        data.extend(bytes);
251    }
252    let mut header = format!("{{{}}}", header.join(",")).into_bytes();
253    header.resize(header.len().next_multiple_of(8), b' ');
254    let mut out = Vec::with_capacity(8 + header.len() + data.len());
255    out.extend((header.len() as u64).to_le_bytes());
256    out.extend(header);
257    out.extend(data);
258    std::fs::write(path.as_ref(), out)
259        .map_err(|e| Error::Io(format!("{}: {e}", path.as_ref().display())))
260}
nuvec_io/safetensors.rs

nuvec_io/
safetensors.rs
