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!(!ckpt.tensors.is_empty(), "checkpoint {} contains no tensors", input.display());
22
23    let mut writer = GGUFWriter::new();
24    for (k, v) in ckpt.metadata {
25        writer.add_metadata(k, v);
26    }
27
28    let mut fallback_count = 0usize;
29    for t in &ckpt.tensors {
30        let n_elems: u64 = t.shape.iter().product();
31        let innermost = t.shape.last().copied().unwrap_or(1);
32        // Tensors too small for Q8_0's 32-element blocks fall back to F32, matching
33        // the generic `zsfm convert` command's behavior.
34        let dst = if dtype == GGMLType::Q8_0 && (innermost % 32 != 0 || n_elems % 32 != 0) {
35            fallback_count += 1;
36            GGMLType::F32
37        } else {
38            dtype
39        };
40        let data =
41            cast::cast_data(&t.data, t.dtype, dst).with_context(|| format!("tensor {}: cast failed", t.name))?;
42        let gguf_shape: Vec<u64> = t.shape.iter().rev().copied().collect();
43        writer.add_tensor(t.name.clone(), gguf_shape, dst, data);
44    }
45    if fallback_count > 0 {
46        eprintln!("note: {fallback_count} tensor(s) fell back to F32 (too small for Q8_0 blocks)");
47    }
48
49    if let Some(parent) = output.parent() {
50        if !parent.as_os_str().is_empty() {
51            std::fs::create_dir_all(parent)?;
52        }
53    }
54    let out_file = File::create(output).with_context(|| format!("create {}", output.display()))?;
55    writer.write_to(&mut BufWriter::new(out_file))?;
56    Ok(())
57}
zsfm_checkpoint/recast.rs

zsfm_checkpoint/
recast.rs
