zsfm_checkpoint/
recast.rs

1//! Read any supported checkpoint (including an existing GGUF) and rewrite it at a
2//! different dtype, carrying metadata through. Shared by the generic `zsfm convert`
3//! command and by every per-model CLI's cached-GGUF fast path: once a repo has been
4//! downloaded and converted to a canonical F32 GGUF once, later requests for a
5//! different dtype recast from that cache instead of re-downloading from HuggingFace.
6
7use std::fs::File;
8use std::io::BufWriter;
9use std::path::Path;
10
11use anyhow::{Context, Result};
12
13use zsfm_gguf::{GGMLType, GGUFWriter};
14
15use crate::cast;
16use crate::read::{load_checkpoint, LoadOptions};
17
18pub fn recast(input: &Path, output: &Path, dtype: GGMLType) -> Result<()> {
19    let ckpt = load_checkpoint(&[input.to_path_buf()], &LoadOptions::default())
20        .with_context(|| format!("load checkpoint {}", input.display()))?;
21    anyhow::ensure!(
22        !ckpt.tensors.is_empty(),
23        "checkpoint {} contains no tensors",
24        input.display()
25    );
26
27    let mut writer = GGUFWriter::new();
28    for (k, v) in ckpt.metadata {
29        writer.add_metadata(k, v);
30    }
31
32    let mut fallback_count = 0usize;
33    for t in &ckpt.tensors {
34        let n_elems: u64 = t.shape.iter().product();
35        let innermost = t.shape.last().copied().unwrap_or(1);
36        // Tensors too small for Q8_0's 32-element blocks fall back to F32, matching
37        // the generic `zsfm convert` command's behavior.
38        let dst = if dtype == GGMLType::Q8_0 && (innermost % 32 != 0 || n_elems % 32 != 0) {
39            fallback_count += 1;
40            GGMLType::F32
41        } else {
42            dtype
43        };
44        let data = cast::cast_data(&t.data, t.dtype, dst)
45            .with_context(|| format!("tensor {}: cast failed", t.name))?;
46        let gguf_shape: Vec<u64> = t.shape.iter().rev().copied().collect();
47        writer.add_tensor(t.name.clone(), gguf_shape, dst, data);
48    }
49    if fallback_count > 0 {
50        eprintln!("note: {fallback_count} tensor(s) fell back to F32 (too small for Q8_0 blocks)");
51    }
52
53    if let Some(parent) = output.parent() {
54        if !parent.as_os_str().is_empty() {
55            std::fs::create_dir_all(parent)?;
56        }
57    }
58    let out_file = File::create(output).with_context(|| format!("create {}", output.display()))?;
59    writer.write_to(&mut BufWriter::new(out_file))?;
60    Ok(())
61}
zsfm_checkpoint/recast.rs

zsfm_checkpoint/
recast.rs
