logicaffeine_compile/model/
safetensors.rs

1//! Safetensors files as a checkpoint directory publishes them: one file, or shards named by
2//! `model.safetensors.index.json`. A file is an 8-byte little-endian header length, a JSON header naming every tensor's
3//! dtype, shape and byte range within the data section, then the data. Tensors are read by position, never mapped or
4//! read whole, so a checkpoint of any size opens in the time it takes to read its headers.
5
6use std::collections::BTreeMap;
7use std::fs::File;
8use std::io::{self, Read};
9use std::path::{Path, PathBuf};
10
11/// An element type a safetensors header names.
12#[derive(Clone, Copy, Debug, PartialEq, Eq)]
13pub enum Dtype {
14    F64,
15    F32,
16    F16,
17    Bf16,
18    F8E4m3,
19    F8E5m2,
20    I64,
21    I32,
22    I16,
23    I8,
24    U8,
25    Bool,
26}
27
28impl Dtype {
29    fn parse(s: &str) -> Option<Dtype> {
30        Some(match s {
31            "F64" => Dtype::F64,
32            "F32" => Dtype::F32,
33            "F16" => Dtype::F16,
34            "BF16" => Dtype::Bf16,
35            "F8_E4M3" => Dtype::F8E4m3,
36            "F8_E5M2" => Dtype::F8E5m2,
37            "I64" => Dtype::I64,
38            "I32" => Dtype::I32,
39            "I16" => Dtype::I16,
40            "I8" => Dtype::I8,
41            "U8" => Dtype::U8,
42            "BOOL" => Dtype::Bool,
43            _ => return None,
44        })
45    }
46
47    /// Bytes an element takes.
48    pub fn bytes(self) -> usize {
49        match self {
50            Dtype::F64 | Dtype::I64 => 8,
51            Dtype::F32 | Dtype::I32 => 4,
52            Dtype::F16 | Dtype::Bf16 | Dtype::I16 => 2,
53            Dtype::F8E4m3 | Dtype::F8E5m2 | Dtype::I8 | Dtype::U8 | Dtype::Bool => 1,
54        }
55    }
56}
57
58/// One tensor's header entry: its dtype, its shape (outermost first, as PyTorch states it) and its byte range within
59/// its file's data section.
60#[derive(Clone, Debug, PartialEq)]
61pub struct Entry {
62    pub dtype: Dtype,
63    pub shape: Vec<u64>,
64    pub begin: u64,
65    pub end: u64,
66}
67
68impl Entry {
69    pub fn elements(&self) -> u64 {
70        self.shape.iter().product()
71    }
72}
73
74/// One safetensors file.
75pub struct Shard {
76    file: File,
77    path: PathBuf,
78    data_start: u64,
79    pub tensors: BTreeMap<String, Entry>,
80}
81
82fn invalid(msg: impl Into<String>) -> io::Error {
83    io::Error::new(io::ErrorKind::InvalidData, msg.into())
84}
85
86impl Shard {
87    pub fn open(path: &Path) -> io::Result<Shard> {
88        let mut file = File::open(path).map_err(|e| io::Error::new(e.kind(), format!("{}: {e}", path.display())))?;
89        let len = file.metadata()?.len();
90        let mut head = [0u8; 8];
91        file.read_exact(&mut head)?;
92        let header_len = u64::from_le_bytes(head);
93        if header_len > len.saturating_sub(8) {
94            return Err(invalid(format!("{}: a header of {header_len} bytes in a file of {len}", path.display())));
95        }
96        let mut header = vec![0u8; header_len as usize];
97        file.read_exact(&mut header)?;
98        let json: serde_json::Value = serde_json::from_slice(&header)
99            .map_err(|e| invalid(format!("{}: the header is not JSON: {e}", path.display())))?;
100        let object = json.as_object().ok_or_else(|| invalid(format!("{}: the header is not an object", path.display())))?;
101        let data_start = 8 + header_len;
102        let data_len = len - data_start;
103        let mut tensors = BTreeMap::new();
104        for (name, v) in object {
105            if name == "__metadata__" {
106                continue;
107            }
108            let dtype_name = v["dtype"].as_str().unwrap_or_default();
109            let dtype = Dtype::parse(dtype_name)
110                .ok_or_else(|| invalid(format!("{}: {name} has dtype `{dtype_name}`, which is not read", path.display())))?;
111            let shape: Vec<u64> = v["shape"]
112                .as_array()
113                .ok_or_else(|| invalid(format!("{}: {name} has no shape", path.display())))?
114                .iter()
115                .map(|d| d.as_u64().ok_or_else(|| invalid(format!("{}: {name}'s shape is not integers", path.display()))))
116                .collect::<io::Result<_>>()?;
117            let offsets = v["data_offsets"]
118                .as_array()
119                .filter(|a| a.len() == 2)
120                .ok_or_else(|| invalid(format!("{}: {name} has no data offsets", path.display())))?;
121            let (begin, end) = (offsets[0].as_u64().unwrap_or(u64::MAX), offsets[1].as_u64().unwrap_or(u64::MAX));
122            let entry = Entry { dtype, shape, begin, end };
123            if begin > end || end > data_len || end - begin != entry.elements() * dtype.bytes() as u64 {
124                return Err(invalid(format!(
125                    "{}: {name} claims bytes {begin}..{end} for {} elements of {dtype_name} in a data section of {data_len}",
126                    path.display(),
127                    entry.elements()
128                )));
129            }
130            tensors.insert(name.clone(), entry);
131        }
132        Ok(Shard { file, path: path.to_path_buf(), data_start, tensors })
133    }
134
135    pub fn path(&self) -> &Path {
136        &self.path
137    }
138
139    /// `len` bytes of tensor `name` starting `offset` bytes into it.
140    pub fn read(&self, name: &str, offset: u64, len: usize) -> io::Result<Vec<u8>> {
141        let e = self
142            .tensors
143            .get(name)
144            .ok_or_else(|| invalid(format!("{}: no tensor {name}", self.path.display())))?;
145        if offset + len as u64 > e.end - e.begin {
146            return Err(invalid(format!("{}: {name} read past its end", self.path.display())));
147        }
148        let mut out = vec![0u8; len];
149        self.read_at(&mut out, self.data_start + e.begin + offset)?;
150        Ok(out)
151    }
152
153    /// `buf.len()` bytes of the file from `offset`, moving no cursor another reader of the shard shares: a positional
154    /// read where the platform has one, and elsewhere a seek and a read on a handle of the read's own.
155    #[cfg(unix)]
156    fn read_at(&self, buf: &mut [u8], offset: u64) -> io::Result<()> {
157        use std::os::unix::fs::FileExt;
158        self.file.read_exact_at(buf, offset)
159    }
160
161    #[cfg(windows)]
162    fn read_at(&self, buf: &mut [u8], offset: u64) -> io::Result<()> {
163        use std::os::windows::fs::FileExt;
164        let mut done = 0;
165        while done < buf.len() {
166            match self.file.seek_read(&mut buf[done..], offset + done as u64)? {
167                0 => return Err(io::Error::new(io::ErrorKind::UnexpectedEof, format!("{}: read past its end", self.path.display()))),
168                n => done += n,
169            }
170        }
171        Ok(())
172    }
173
174    #[cfg(not(any(unix, windows)))]
175    fn read_at(&self, buf: &mut [u8], offset: u64) -> io::Result<()> {
176        use std::io::{Seek, SeekFrom};
177        let mut file = File::open(&self.path)?;
178        file.seek(SeekFrom::Start(offset))?;
179        file.read_exact(buf)
180    }
181}
182
183/// The shard files a checkpoint directory holds: the files its index names, else its one `model.safetensors`, else
184/// every `*.safetensors` file in it. A file the index names that is not there is refused by name.
185pub fn shard_paths(dir: &Path) -> io::Result<Vec<PathBuf>> {
186    let index = dir.join("model.safetensors.index.json");
187    if index.is_file() {
188        let json: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&index)?)
189            .map_err(|e| invalid(format!("{}: {e}", index.display())))?;
190        let map = json["weight_map"]
191            .as_object()
192            .ok_or_else(|| invalid(format!("{}: no weight_map", index.display())))?;
193        let mut files: Vec<String> = map.values().filter_map(|v| v.as_str().map(str::to_string)).collect();
194        files.sort();
195        files.dedup();
196        let mut out = Vec::new();
197        for f in files {
198            let p = dir.join(&f);
199            if !p.is_file() {
200                return Err(io::Error::new(
201                    io::ErrorKind::NotFound,
202                    format!("{}: the index names shard {f}, which the directory does not hold", dir.display()),
203                ));
204            }
205            out.push(p);
206        }
207        return Ok(out);
208    }
209    let single = dir.join("model.safetensors");
210    if single.is_file() {
211        return Ok(vec![single]);
212    }
213    let mut out: Vec<PathBuf> = std::fs::read_dir(dir)?
214        .filter_map(|e| e.ok().map(|e| e.path()))
215        .filter(|p| p.extension().is_some_and(|x| x == "safetensors"))
216        .collect();
217    out.sort();
218    if out.is_empty() {
219        return Err(io::Error::new(io::ErrorKind::NotFound, format!("{}: no safetensors file", dir.display())));
220    }
221    Ok(out)
222}
logicaffeine_compile/model/safetensors.rs

logicaffeine_compile/model/
safetensors.rs
