Skip to main content

skippy_model/
hf_checkpoint.rs

1//! Validated SafeTensors checkpoint discovery and tensor access.
2
3use std::collections::{BTreeMap, BTreeSet};
4use std::fs::{self, File};
5use std::io::Write;
6use std::ops::Range;
7use std::path::{Component, Path, PathBuf};
8
9use anyhow::{Context, Result, ensure};
10use memmap2::{Mmap, MmapOptions};
11use safetensors::SafeTensors;
12use serde::{Deserialize, Serialize};
13
14use crate::ConvertOutputType;
15
16#[derive(Debug, Serialize)]
17pub struct HfCheckpointPlan {
18    pub source: PathBuf,
19    pub safetensor_count: usize,
20    pub tensor_count: usize,
21    pub total_tensor_bytes: u64,
22    pub largest_tensor_bytes: u64,
23    pub source_windows: Vec<HfSourceWindow>,
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub stream_verification: Option<HfStreamVerification>,
26}
27
28#[derive(Debug, Serialize)]
29pub struct HfSourceWindow {
30    pub index: u32,
31    pub files: Vec<PathBuf>,
32    pub tensor_count: usize,
33    pub total_tensor_bytes: u64,
34    pub largest_tensor_bytes: u64,
35}
36
37#[derive(Debug, Serialize)]
38pub struct HfStreamVerification {
39    pub safetensor_count: usize,
40    pub tensor_count: usize,
41    pub streamed_bytes: u64,
42    pub buffer_size: usize,
43}
44
45#[derive(Debug)]
46struct SafetensorSummary {
47    path: PathBuf,
48    tensor_count: usize,
49    total_tensor_bytes: u64,
50    largest_tensor_bytes: u64,
51}
52
53#[derive(Debug)]
54pub struct SafetensorFile {
55    path: PathBuf,
56    data_start: u64,
57    mapping: Mmap,
58    tensors: BTreeMap<String, SafetensorTensorInfo>,
59}
60
61impl SafetensorFile {
62    pub fn open(path: &Path) -> Result<Self> {
63        let file = File::open(path).with_context(|| format!("open {}", path.display()))?;
64        // SAFETY: the file is opened read-only and Skippy never mutates or truncates checkpoint
65        // files while a SafetensorFile exists. The mapping is owned by this object and therefore
66        // outlives every tensor slice exposed through its methods.
67        let mapping = unsafe { MmapOptions::new().map(&file) }
68            .with_context(|| format!("mmap {}", path.display()))?;
69        let parsed = SafeTensors::deserialize(&mapping)
70            .with_context(|| format!("validate safetensors container {}", path.display()))?;
71        let data_start = parsed
72            .tensors()
73            .into_iter()
74            .map(|(_, tensor)| tensor.data().as_ptr() as usize - mapping.as_ptr() as usize)
75            .min()
76            .unwrap_or(mapping.len()) as u64;
77        let mut tensors = BTreeMap::new();
78        for (name, tensor) in parsed.tensors() {
79            let absolute_start = tensor.data().as_ptr() as usize - mapping.as_ptr() as usize;
80            let absolute_end = absolute_start
81                .checked_add(tensor.data().len())
82                .with_context(|| format!("data range overflow for tensor {name}"))?;
83            ensure!(
84                absolute_start >= data_start as usize,
85                "tensor {name} precedes data section"
86            );
87            let shape = tensor
88                .shape()
89                .iter()
90                .map(|&dim| u64::try_from(dim).context("tensor dimension does not fit u64"))
91                .collect::<Result<Vec<_>>>()?;
92            tensors.insert(
93                name.to_string(),
94                SafetensorTensorInfo {
95                    name: name.to_string(),
96                    dtype: format!("{:?}", tensor.dtype()),
97                    shape,
98                    relative_data_offsets: [
99                        (absolute_start - data_start as usize) as u64,
100                        (absolute_end - data_start as usize) as u64,
101                    ],
102                    absolute_data_start: absolute_start,
103                    byte_len: tensor.data().len(),
104                },
105            );
106        }
107        Ok(Self {
108            path: path.to_path_buf(),
109            data_start,
110            mapping,
111            tensors,
112        })
113    }
114
115    pub fn path(&self) -> &Path {
116        &self.path
117    }
118
119    pub fn data_start(&self) -> u64 {
120        self.data_start
121    }
122
123    pub fn tensors(&self) -> &BTreeMap<String, SafetensorTensorInfo> {
124        &self.tensors
125    }
126
127    pub fn stream_tensor<W: Write>(
128        &self,
129        name: &str,
130        writer: &mut W,
131        buffer_size: usize,
132    ) -> Result<u64> {
133        self.stream_tensor_chunks(name, buffer_size, |chunk| {
134            writer.write_all(chunk).context("write tensor bytes")
135        })
136    }
137
138    pub fn stream_tensor_chunks<F>(
139        &self,
140        name: &str,
141        buffer_size: usize,
142        mut on_chunk: F,
143    ) -> Result<u64>
144    where
145        F: FnMut(&[u8]) -> Result<()>,
146    {
147        let tensor = self
148            .tensors
149            .get(name)
150            .with_context(|| format!("tensor {name} not found in {}", self.path.display()))?;
151        ensure!(buffer_size > 0, "buffer_size must be greater than zero");
152        let range = tensor.absolute_data_range();
153        let data = self
154            .mapping
155            .get(range.clone())
156            .with_context(|| format!("tensor {name} range is outside {}", self.path.display()))?;
157        for chunk in data.chunks(buffer_size) {
158            on_chunk(chunk)?;
159        }
160        Ok(data.len() as u64)
161    }
162}
163
164#[derive(Debug)]
165pub struct SafetensorTensorInfo {
166    name: String,
167    dtype: String,
168    shape: Vec<u64>,
169    relative_data_offsets: [u64; 2],
170    absolute_data_start: usize,
171    byte_len: usize,
172}
173
174impl SafetensorTensorInfo {
175    pub fn name(&self) -> &str {
176        &self.name
177    }
178
179    pub fn dtype(&self) -> &str {
180        &self.dtype
181    }
182
183    pub fn shape(&self) -> &[u64] {
184        &self.shape
185    }
186
187    pub fn relative_data_offsets(&self) -> [u64; 2] {
188        self.relative_data_offsets
189    }
190
191    pub fn byte_len(&self) -> u64 {
192        self.byte_len as u64
193    }
194
195    fn absolute_data_range(&self) -> Range<usize> {
196        self.absolute_data_start..self.absolute_data_start + self.byte_len
197    }
198}
199
200pub fn inspect_hf_checkpoint(
201    source: &Path,
202    max_memory_bytes: Option<u64>,
203    staging_fraction: f64,
204) -> Result<HfCheckpointPlan> {
205    ensure!(
206        staging_fraction > 0.0 && staging_fraction <= 1.0,
207        "--staging-fraction must be in the range (0, 1]"
208    );
209    let safetensors = discover_safetensors(source)?;
210    ensure!(
211        !safetensors.is_empty(),
212        "no safetensors files found under {}",
213        source.display()
214    );
215    let mut summaries = safetensors
216        .iter()
217        .map(|path| summarize_safetensor(path))
218        .collect::<Result<Vec<_>>>()?;
219    summaries.sort_by(|a, b| a.path.cmp(&b.path));
220    let tensor_count = summaries.iter().map(|summary| summary.tensor_count).sum();
221    let total_tensor_bytes = summaries
222        .iter()
223        .map(|summary| summary.total_tensor_bytes)
224        .sum();
225    let largest_tensor_bytes = summaries
226        .iter()
227        .map(|summary| summary.largest_tensor_bytes)
228        .max()
229        .unwrap_or(0);
230    let source_windows = plan_source_windows(&summaries, max_memory_bytes, staging_fraction)?;
231    Ok(HfCheckpointPlan {
232        source: source.to_path_buf(),
233        safetensor_count: summaries.len(),
234        tensor_count,
235        total_tensor_bytes,
236        largest_tensor_bytes,
237        source_windows,
238        stream_verification: None,
239    })
240}
241
242pub fn verify_hf_checkpoint_tensor_streams(
243    source: &Path,
244    buffer_size: usize,
245) -> Result<HfStreamVerification> {
246    let safetensors = discover_safetensors(source)?;
247    let mut sink = std::io::sink();
248    let mut tensor_count = 0_usize;
249    let mut streamed_bytes = 0_u64;
250    for path in &safetensors {
251        let safetensor = SafetensorFile::open(path)?;
252        ensure!(
253            safetensor.path().is_file(),
254            "safetensor path is not a file: {}",
255            safetensor.path().display()
256        );
257        ensure!(
258            safetensor.data_start() >= 8,
259            "invalid safetensors data start in {}",
260            safetensor.path().display()
261        );
262        for tensor in safetensor.tensors().values() {
263            ensure!(
264                !tensor.name().is_empty(),
265                "safetensor tensor has empty name"
266            );
267            ensure!(
268                dtype_size(tensor.dtype()).is_some(),
269                "unsupported safetensors dtype {}",
270                tensor.dtype()
271            );
272            let offsets = tensor.relative_data_offsets();
273            ensure!(
274                offsets[0] <= offsets[1],
275                "invalid safetensors offsets for {}",
276                tensor.name()
277            );
278            let _rank = tensor.shape().len();
279            streamed_bytes += safetensor.stream_tensor(tensor.name(), &mut sink, buffer_size)?;
280            tensor_count += 1;
281        }
282    }
283    Ok(HfStreamVerification {
284        safetensor_count: safetensors.len(),
285        tensor_count,
286        streamed_bytes,
287        buffer_size,
288    })
289}
290
291pub fn open_safetensor_files(source: &Path) -> Result<Vec<SafetensorFile>> {
292    discover_safetensors(source)?
293        .iter()
294        .map(|path| SafetensorFile::open(path))
295        .collect()
296}
297
298pub fn resolve_auto_output_type(
299    source: &Path,
300    requested: ConvertOutputType,
301) -> Result<ConvertOutputType> {
302    if requested != ConvertOutputType::Auto {
303        return Ok(requested);
304    }
305    let files = open_safetensor_files(source)?;
306    Ok(resolve_auto_output_type_from_files(&files, requested))
307}
308
309pub(crate) fn resolve_auto_output_type_from_files(
310    files: &[SafetensorFile],
311    requested: ConvertOutputType,
312) -> ConvertOutputType {
313    if requested != ConvertOutputType::Auto {
314        return requested;
315    }
316    for safetensor in files {
317        for tensor in safetensor.tensors().values() {
318            if tensor.shape().len() < 2 {
319                continue;
320            }
321            match tensor.dtype() {
322                "BF16" => return ConvertOutputType::Bf16,
323                "F16" => return ConvertOutputType::F16,
324                _ => {}
325            }
326        }
327    }
328    ConvertOutputType::F16
329}
330
331pub fn discover_safetensors(source: &Path) -> Result<Vec<PathBuf>> {
332    ensure!(
333        source.is_dir(),
334        "HF checkpoint source must be a directory: {}",
335        source.display()
336    );
337    let mut indexed = discover_indexed_safetensors(source)?;
338    if !indexed.is_empty() {
339        let mtp_sidecar = source.join("mtp.safetensors");
340        if mtp_sidecar.is_file() && !indexed.contains(&mtp_sidecar) {
341            indexed.push(mtp_sidecar);
342            indexed.sort();
343        }
344        return Ok(indexed);
345    }
346    indexed = fs::read_dir(source)
347        .with_context(|| format!("read checkpoint directory {}", source.display()))?
348        .map(|entry| entry.map(|entry| entry.path()))
349        .collect::<std::io::Result<Vec<_>>>()?
350        .into_iter()
351        .filter(|path| path.extension().is_some_and(|ext| ext == "safetensors"))
352        .collect();
353    indexed.sort();
354    Ok(indexed)
355}
356
357fn discover_indexed_safetensors(source: &Path) -> Result<Vec<PathBuf>> {
358    let index_path = source.join("model.safetensors.index.json");
359    if !index_path.is_file() {
360        return Ok(Vec::new());
361    }
362    let index: SafetensorIndex = serde_json::from_slice(
363        &fs::read(&index_path).with_context(|| format!("read {}", index_path.display()))?,
364    )
365    .with_context(|| format!("parse {}", index_path.display()))?;
366    let mut files = index
367        .weight_map
368        .values()
369        .map(|name| indexed_shard_path(source, name))
370        .collect::<Result<BTreeSet<_>>>()?
371        .into_iter()
372        .collect::<Vec<_>>();
373    files.sort();
374    Ok(files)
375}
376
377fn indexed_shard_path(source: &Path, name: &str) -> Result<PathBuf> {
378    let relative = Path::new(name);
379    ensure!(
380        !name.is_empty()
381            && !relative.is_absolute()
382            && relative
383                .components()
384                .all(|component| matches!(component, Component::Normal(_))),
385        "SafeTensors index shard path must remain within the checkpoint directory: {name:?}"
386    );
387    ensure!(
388        relative
389            .extension()
390            .is_some_and(|extension| extension == "safetensors"),
391        "SafeTensors index shard must have a .safetensors extension: {name:?}"
392    );
393    let path = source.join(relative);
394    ensure!(
395        path.is_file(),
396        "SafeTensors index shard does not exist: {}",
397        path.display()
398    );
399    Ok(path)
400}
401
402#[derive(Debug, Deserialize)]
403struct SafetensorIndex {
404    weight_map: BTreeMap<String, String>,
405}
406
407fn summarize_safetensor(path: &Path) -> Result<SafetensorSummary> {
408    let safetensor = SafetensorFile::open(path)?;
409    let tensor_count = safetensor.tensors.len();
410    let total_tensor_bytes = safetensor
411        .tensors
412        .values()
413        .map(SafetensorTensorInfo::byte_len)
414        .sum();
415    let largest_tensor_bytes = safetensor
416        .tensors
417        .values()
418        .map(SafetensorTensorInfo::byte_len)
419        .max()
420        .unwrap_or(0);
421    Ok(SafetensorSummary {
422        path: path.to_path_buf(),
423        tensor_count,
424        total_tensor_bytes,
425        largest_tensor_bytes,
426    })
427}
428
429fn dtype_size(dtype: &str) -> Option<u64> {
430    match dtype {
431        "BOOL" | "I8" | "U8" | "F8_E4M3" | "F8_E5M2" => Some(1),
432        "I16" | "U16" | "F16" | "BF16" => Some(2),
433        "I32" | "U32" | "F32" => Some(4),
434        "I64" | "U64" | "F64" => Some(8),
435        _ => None,
436    }
437}
438
439fn plan_source_windows(
440    summaries: &[SafetensorSummary],
441    max_memory_bytes: Option<u64>,
442    staging_fraction: f64,
443) -> Result<Vec<HfSourceWindow>> {
444    let budget = max_memory_bytes
445        .map(|memory| ((memory as f64) * staging_fraction).floor() as u64)
446        .unwrap_or(u64::MAX)
447        .max(1);
448    let mut windows = Vec::new();
449    let mut current = SourceWindowBuilder::new(1);
450    for summary in summaries {
451        if !current.is_empty() && current.total_tensor_bytes + summary.total_tensor_bytes > budget {
452            windows.push(current.finish());
453            current = SourceWindowBuilder::new(windows.len() as u32 + 1);
454        }
455        current.push(summary);
456    }
457    if !current.is_empty() {
458        windows.push(current.finish());
459    }
460    Ok(windows)
461}
462
463struct SourceWindowBuilder {
464    index: u32,
465    files: Vec<PathBuf>,
466    tensor_count: usize,
467    total_tensor_bytes: u64,
468    largest_tensor_bytes: u64,
469}
470
471impl SourceWindowBuilder {
472    fn new(index: u32) -> Self {
473        Self {
474            index,
475            files: Vec::new(),
476            tensor_count: 0,
477            total_tensor_bytes: 0,
478            largest_tensor_bytes: 0,
479        }
480    }
481
482    fn is_empty(&self) -> bool {
483        self.files.is_empty()
484    }
485
486    fn push(&mut self, summary: &SafetensorSummary) {
487        self.files.push(summary.path.clone());
488        self.tensor_count += summary.tensor_count;
489        self.total_tensor_bytes += summary.total_tensor_bytes;
490        self.largest_tensor_bytes = self.largest_tensor_bytes.max(summary.largest_tensor_bytes);
491    }
492
493    fn finish(self) -> HfSourceWindow {
494        HfSourceWindow {
495            index: self.index,
496            files: self.files,
497            tensor_count: self.tensor_count,
498            total_tensor_bytes: self.total_tensor_bytes,
499            largest_tensor_bytes: self.largest_tensor_bytes,
500        }
501    }
502}
503
504#[cfg(test)]
505mod tests {
506    use super::*;
507    use std::sync::atomic::{AtomicU64, Ordering};
508
509    static TEMP_DIR_COUNTER: AtomicU64 = AtomicU64::new(0);
510
511    #[test]
512    fn plans_unindexed_safetensors_under_memory_budget() {
513        let root = unique_temp_dir();
514        fs::create_dir_all(&root).unwrap();
515        write_safetensor(
516            &root.join("model-00001-of-00002.safetensors"),
517            &[("a.weight", "F32", &[2], &[1, 2, 3, 4, 5, 6, 7, 8])],
518        );
519        write_safetensor(
520            &root.join("model-00002-of-00002.safetensors"),
521            &[("b.weight", "BF16", &[4], &[1, 2, 3, 4, 5, 6, 7, 8])],
522        );
523
524        let plan = inspect_hf_checkpoint(&root, Some(12), 1.0).unwrap();
525
526        assert_eq!(plan.safetensor_count, 2);
527        assert_eq!(plan.tensor_count, 2);
528        assert_eq!(plan.total_tensor_bytes, 16);
529        assert_eq!(plan.source_windows.len(), 2);
530        assert_eq!(plan.source_windows[0].total_tensor_bytes, 8);
531        assert_eq!(plan.source_windows[1].total_tensor_bytes, 8);
532        fs::remove_dir_all(root).unwrap();
533    }
534
535    #[test]
536    fn uses_index_weight_map_when_present() {
537        let root = unique_temp_dir();
538        fs::create_dir_all(&root).unwrap();
539        write_safetensor(
540            &root.join("shard-b.safetensors"),
541            &[("b.weight", "F32", &[1], &[1, 2, 3, 4])],
542        );
543        write_safetensor(
544            &root.join("shard-a.safetensors"),
545            &[("a.weight", "F32", &[1], &[1, 2, 3, 4])],
546        );
547        fs::write(
548            root.join("model.safetensors.index.json"),
549            r#"{"metadata":{},"weight_map":{"a.weight":"shard-a.safetensors","b.weight":"shard-b.safetensors"}}"#,
550        )
551        .unwrap();
552
553        let plan = inspect_hf_checkpoint(&root, None, 1.0).unwrap();
554
555        assert_eq!(plan.safetensor_count, 2);
556        assert_eq!(plan.source_windows.len(), 1);
557        fs::remove_dir_all(root).unwrap();
558    }
559
560    #[test]
561    fn rejects_index_shards_that_escape_checkpoint_directory() {
562        let root = unique_temp_dir();
563        fs::create_dir_all(&root).unwrap();
564        fs::write(
565            root.join("model.safetensors.index.json"),
566            r#"{"metadata":{},"weight_map":{"a.weight":"../outside.safetensors"}}"#,
567        )
568        .unwrap();
569
570        let error = discover_safetensors(&root).unwrap_err().to_string();
571
572        assert!(error.contains("must remain within the checkpoint directory"));
573        fs::remove_dir_all(root).unwrap();
574    }
575
576    #[test]
577    fn includes_unindexed_mtp_sidecar_with_indexed_checkpoint() {
578        let root = unique_temp_dir();
579        fs::create_dir_all(&root).unwrap();
580        write_safetensor(
581            &root.join("shard-a.safetensors"),
582            &[("a.weight", "F32", &[1], &[1, 2, 3, 4])],
583        );
584        write_safetensor(
585            &root.join("mtp.safetensors"),
586            &[("model.mtp.layers.0.weight", "F32", &[1], &[5, 6, 7, 8])],
587        );
588        fs::write(
589            root.join("model.safetensors.index.json"),
590            r#"{"metadata":{},"weight_map":{"a.weight":"shard-a.safetensors"}}"#,
591        )
592        .unwrap();
593
594        let files = discover_safetensors(&root).unwrap();
595
596        assert_eq!(
597            files,
598            vec![
599                root.join("mtp.safetensors"),
600                root.join("shard-a.safetensors")
601            ]
602        );
603        fs::remove_dir_all(root).unwrap();
604    }
605
606    #[test]
607    fn streams_tensor_bytes_without_reading_neighbor_tensors() {
608        let root = unique_temp_dir();
609        fs::create_dir_all(&root).unwrap();
610        let path = root.join("model.safetensors");
611        write_safetensor(
612            &path,
613            &[
614                ("a.weight", "U8", &[4], &[1, 2, 3, 4]),
615                ("b.weight", "U8", &[3], &[9, 8, 7]),
616            ],
617        );
618
619        let safetensor = SafetensorFile::open(&path).unwrap();
620        let tensor = safetensor.tensors().get("b.weight").unwrap();
621        let mut output = Vec::new();
622        let copied = safetensor
623            .stream_tensor("b.weight", &mut output, 2)
624            .unwrap();
625
626        assert_eq!(safetensor.path(), path);
627        assert!(safetensor.data_start() > 8);
628        assert_eq!(tensor.name(), "b.weight");
629        assert_eq!(tensor.dtype(), "U8");
630        assert_eq!(tensor.shape(), &[3]);
631        assert_eq!(tensor.relative_data_offsets(), [4, 7]);
632        assert_eq!(tensor.byte_len(), 3);
633        assert_eq!(copied, 3);
634        assert_eq!(output, vec![9, 8, 7]);
635        fs::remove_dir_all(root).unwrap();
636    }
637
638    #[test]
639    fn resolves_auto_output_type_from_first_rank_two_float_tensor() {
640        let root = unique_temp_dir();
641        fs::create_dir_all(&root).unwrap();
642        write_safetensor(
643            &root.join("model.safetensors"),
644            &[
645                ("a.bias", "BF16", &[4], &[1, 2, 3, 4, 5, 6, 7, 8]),
646                ("b.weight", "F16", &[2, 2], &[1, 2, 3, 4, 5, 6, 7, 8]),
647            ],
648        );
649
650        let output_type = resolve_auto_output_type(&root, ConvertOutputType::Auto).unwrap();
651
652        assert_eq!(output_type, ConvertOutputType::F16);
653        fs::remove_dir_all(root).unwrap();
654    }
655
656    #[test]
657    fn resolves_auto_output_type_to_bf16_when_rank_two_bf16_appears_first() {
658        let root = unique_temp_dir();
659        fs::create_dir_all(&root).unwrap();
660        write_safetensor(
661            &root.join("model.safetensors"),
662            &[("a.weight", "BF16", &[2, 2], &[1, 2, 3, 4, 5, 6, 7, 8])],
663        );
664
665        let output_type = resolve_auto_output_type(&root, ConvertOutputType::Auto).unwrap();
666
667        assert_eq!(output_type, ConvertOutputType::Bf16);
668        fs::remove_dir_all(root).unwrap();
669    }
670
671    #[test]
672    fn resolves_auto_output_type_to_f16_when_checkpoint_has_no_float_matrix() {
673        let root = unique_temp_dir();
674        fs::create_dir_all(&root).unwrap();
675        write_safetensor(
676            &root.join("model.safetensors"),
677            &[
678                ("a.bias", "BF16", &[4], &[1, 2, 3, 4, 5, 6, 7, 8]),
679                (
680                    "b.count",
681                    "I32",
682                    &[2, 2],
683                    &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
684                ),
685            ],
686        );
687
688        let output_type = resolve_auto_output_type(&root, ConvertOutputType::Auto).unwrap();
689
690        assert_eq!(output_type, ConvertOutputType::F16);
691        fs::remove_dir_all(root).unwrap();
692    }
693
694    fn unique_temp_dir() -> PathBuf {
695        let nanos = std::time::SystemTime::now()
696            .duration_since(std::time::UNIX_EPOCH)
697            .unwrap()
698            .as_nanos();
699        let counter = TEMP_DIR_COUNTER.fetch_add(1, Ordering::Relaxed);
700        std::env::temp_dir().join(format!(
701            "skippy-hf-checkpoint-{}-{nanos}-{counter}",
702            std::process::id()
703        ))
704    }
705
706    fn write_safetensor(path: &Path, tensors: &[(&str, &str, &[u64], &[u8])]) {
707        let mut offset = 0_u64;
708        let mut entries = serde_json::Map::new();
709        for (name, dtype, shape, bytes) in tensors {
710            let end = offset + bytes.len() as u64;
711            entries.insert(
712                (*name).to_string(),
713                serde_json::json!({
714                    "dtype": dtype,
715                    "shape": shape,
716                    "data_offsets": [offset, end],
717                }),
718            );
719            offset = end;
720        }
721        let header = serde_json::Value::Object(entries).to_string();
722        let mut bytes = Vec::new();
723        bytes.extend_from_slice(&(header.len() as u64).to_le_bytes());
724        bytes.extend_from_slice(header.as_bytes());
725        for (_, _, _, tensor_bytes) in tensors {
726            bytes.extend_from_slice(tensor_bytes);
727        }
728        fs::write(path, bytes).unwrap();
729    }
730}