zsfm/
generic.rs

1//! Model-agnostic `zsfm convert` / `zsfm inspect` / `zsfm pull`: turn ANY supported
2//! checkpoint (safetensors, sharded-safetensors index, PyTorch pickle, npy/npz,
3//! ONNX, HDF5/Keras, GGUF) into a GGUF file with tensor names passed through
4//! unchanged, list its contents, or just fetch a HuggingFace repo's checkpoint in
5//! whatever format it happens to be published in.
6//!
7//! Unlike the per-model subcommands (which map tensor names to each inference
8//! engine's schema and assume `config.json` + `model.safetensors`), this path
9//! does no renaming beyond an optional `--strip-prefix` and downloads whatever
10//! format a repo actually has (via `zsfm_hub::download_any_format`), so it works
11//! for arbitrary models.
12
13use std::fs::File;
14use std::io::BufWriter;
15use std::path::PathBuf;
16
17use anyhow::Context;
18use clap::{Args, ValueEnum};
19
20use zsfm_checkpoint::{cast, load_checkpoint, LoadOptions};
21use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
22
23#[derive(Clone, ValueEnum)]
24pub enum DtypeArg {
25    F32,
26    F16,
27    Bf16,
28    Q8,
29}
30
31impl From<DtypeArg> for GGMLType {
32    fn from(d: DtypeArg) -> Self {
33        match d {
34            DtypeArg::F32 => GGMLType::F32,
35            DtypeArg::F16 => GGMLType::F16,
36            DtypeArg::Bf16 => GGMLType::BF16,
37            DtypeArg::Q8 => GGMLType::Q8_0,
38        }
39    }
40}
41
42#[derive(Args)]
43pub struct ConvertArgs {
44    /// Input checkpoint file(s): .safetensors (one or more shards),
45    /// model.safetensors.index.json, .pt/.pth/.bin/.ckpt, .npy/.npz, .onnx,
46    /// .h5/.hdf5, .keras, or .gguf (GGUF input re-quantizes and carries
47    /// metadata through). TF SavedModel/checkpoints: see tools/export-weights.py.
48    /// Omit if using --repo.
49    #[arg(required_unless_present = "repo")]
50    pub inputs: Vec<PathBuf>,
51    /// Download and convert directly from a HuggingFace repo (owner/name) instead
52    /// of local files — whatever checkpoint format it's published in.
53    #[arg(long, conflicts_with = "inputs", value_name = "OWNER/NAME")]
54    pub repo: Option<String>,
55    /// With --repo: which checkpoint format to fetch (safetensors, bin, pt, pth,
56    /// ckpt, onnx, h5, hdf5, keras, npz, npy, gguf). Default: the highest-priority
57    /// format actually present in the repo.
58    #[arg(long, requires = "repo")]
59    pub format: Option<String>,
60    /// With --repo: case-insensitive substring to pick among multiple files of the
61    /// chosen format — e.g. a specific GGUF quant variant, or a safetensors file
62    /// when the repo has no sharding index and ships several unrelated ones.
63    #[arg(long, requires = "repo")]
64    pub file: Option<String>,
65    /// With --repo: branch, tag, or commit to fetch from.
66    #[arg(long, default_value = "main", requires = "repo")]
67    pub revision: String,
68    /// With --repo: HuggingFace API token, for gated/private repos.
69    #[arg(long, env = "HF_TOKEN", requires = "repo")]
70    pub token: Option<String>,
71    /// With --repo: directory to download into. Default: models/`<owner>__<name>`.
72    #[arg(long, requires = "repo")]
73    pub model_dir: Option<PathBuf>,
74    /// Output GGUF path.
75    #[arg(short, long)]
76    pub output: PathBuf,
77    /// Output dtype. Tensors too small for Q8_0 blocks fall back to F32.
78    #[arg(long, default_value = "f16")]
79    pub dtype: DtypeArg,
80    /// Value for the `general.architecture` metadata key.
81    #[arg(long)]
82    pub arch: Option<String>,
83    /// Value for the `general.name` metadata key (default: first input's file stem).
84    #[arg(long)]
85    pub name: Option<String>,
86    /// Extra metadata entries, `key=value` (repeatable). Values are parsed as
87    /// bool / u32 / i32 / f32 when they look like one, else stored as strings.
88    #[arg(long = "metadata", value_name = "KEY=VALUE")]
89    pub metadata: Vec<String>,
90    /// For pickle inputs: descend into this top-level dict (e.g. `state_dict`).
91    #[arg(long)]
92    pub pickle_key: Option<String>,
93    /// Strip this prefix from tensor names that carry it (e.g. `model.`).
94    #[arg(long)]
95    pub strip_prefix: Option<String>,
96}
97
98#[derive(Args)]
99pub struct PullArgs {
100    /// HuggingFace repo to download from (owner/name).
101    #[arg(value_name = "OWNER/NAME")]
102    pub repo: String,
103    /// Which checkpoint format to fetch (safetensors, bin, pt, pth, ckpt, onnx,
104    /// h5, hdf5, keras, npz, npy, gguf). Default: the highest-priority format
105    /// actually present in the repo.
106    #[arg(long)]
107    pub format: Option<String>,
108    /// Case-insensitive substring to pick among multiple files of the chosen
109    /// format — e.g. a specific GGUF quant variant.
110    #[arg(long)]
111    pub file: Option<String>,
112    /// Branch, tag, or commit to fetch from.
113    #[arg(long, default_value = "main")]
114    pub revision: String,
115    /// HuggingFace API token, for gated/private repos.
116    #[arg(long, env = "HF_TOKEN")]
117    pub token: Option<String>,
118    /// Directory to download into. Default: models/`<owner>__<name>`.
119    #[arg(long)]
120    pub model_dir: Option<PathBuf>,
121}
122
123#[derive(Args)]
124pub struct InspectArgs {
125    /// Checkpoint file to inspect (any supported format, including .gguf).
126    pub path: PathBuf,
127    /// Print the first N values of --tensor (or of every tensor if --tensor is
128    /// omitted and N is given).
129    #[arg(long, value_name = "N")]
130    pub values: Option<usize>,
131    /// Restrict --values output to this tensor.
132    #[arg(long)]
133    pub tensor: Option<String>,
134    /// For pickle inputs: descend into this top-level dict (e.g. `state_dict`).
135    #[arg(long)]
136    pub pickle_key: Option<String>,
137}
138
139pub async fn run_convert(args: ConvertArgs) -> anyhow::Result<()> {
140    // --repo downloads first, then falls through to the same local-file path
141    // below using the downloaded checkpoint files as `inputs`.
142    let (inputs, repo_default_name) = if let Some(repo) = &args.repo {
143        let model_dir = args
144            .model_dir
145            .clone()
146            .unwrap_or_else(|| PathBuf::from("models").join(repo.replace('/', "__")));
147        println!("Looking up files in {repo}@{} …", args.revision);
148        let downloaded = zsfm_hub::download_any_format(
149            repo,
150            args.format.as_deref(),
151            args.file.as_deref(),
152            &args.revision,
153            args.token.as_deref(),
154            &model_dir,
155        )
156        .await
157        .with_context(|| format!("download {repo}"))?;
158        println!(
159            "Downloaded {} {} file(s) into {} …",
160            downloaded.checkpoint_files.len(),
161            downloaded.format,
162            model_dir.display()
163        );
164        let default_name = repo.rsplit('/').next().unwrap_or(repo).to_string();
165        (downloaded.checkpoint_files, Some(default_name))
166    } else {
167        (args.inputs.clone(), None)
168    };
169
170    let opts = LoadOptions {
171        pickle_key: args.pickle_key.clone(),
172        strip_prefix: args.strip_prefix.clone(),
173    };
174    println!("Loading {} input file(s) …", inputs.len());
175    let ckpt = load_checkpoint(&inputs, &opts)?;
176    println!("Loaded {} tensors.", ckpt.tensors.len());
177    anyhow::ensure!(!ckpt.tensors.is_empty(), "checkpoint contains no tensors");
178
179    let mut writer = GGUFWriter::new();
180
181    // Metadata: carried-over (GGUF inputs) first, then general.* defaults /
182    // overrides, then explicit --metadata pairs. Later entries replace earlier
183    // ones with the same key.
184    let mut meta: Vec<(String, GGUFMetaValue)> = ckpt.metadata;
185    let default_name = repo_default_name.unwrap_or_else(|| {
186        inputs[0]
187            .file_stem()
188            .and_then(|s| s.to_str())
189            .unwrap_or("model")
190            .to_string()
191    });
192    let set = |meta: &mut Vec<(String, GGUFMetaValue)>, key: &str, value: GGUFMetaValue| {
193        meta.retain(|(k, _)| k != key);
194        meta.push((key.to_string(), value));
195    };
196    if args.arch.is_some() || !meta.iter().any(|(k, _)| k == "general.architecture") {
197        let arch = args.arch.clone().unwrap_or_else(|| "unknown".to_string());
198        set(
199            &mut meta,
200            "general.architecture",
201            GGUFMetaValue::String(arch),
202        );
203    }
204    if args.name.is_some() || !meta.iter().any(|(k, _)| k == "general.name") {
205        let name = args.name.clone().unwrap_or(default_name);
206        set(&mut meta, "general.name", GGUFMetaValue::String(name));
207    }
208    for pair in &args.metadata {
209        let (key, value) = pair
210            .split_once('=')
211            .with_context(|| format!("--metadata {pair}: expected key=value"))?;
212        set(&mut meta, key, parse_meta_value(value));
213    }
214    for (key, value) in meta {
215        writer.add_metadata(key, value);
216    }
217
218    let out_dtype: GGMLType = args.dtype.into();
219    let mut fallback_count = 0usize;
220    for t in &ckpt.tensors {
221        let (data, dst, fell_back) =
222            cast::cast_with_q8_fallback(&t.data, t.dtype, &t.shape, out_dtype)
223                .with_context(|| format!("tensor {}: cast failed", t.name))?;
224        fallback_count += fell_back as usize;
225        let gguf_shape: Vec<u64> = t.shape.iter().rev().copied().collect();
226        writer.add_tensor(t.name.clone(), gguf_shape, dst, data);
227    }
228    if fallback_count > 0 {
229        eprintln!("note: {fallback_count} tensor(s) fell back to F32 (too small for Q8_0 blocks)");
230    }
231
232    if let Some(parent) = args.output.parent() {
233        if !parent.as_os_str().is_empty() {
234            std::fs::create_dir_all(parent)?;
235        }
236    }
237    println!(
238        "Writing {} tensors to {} …",
239        ckpt.tensors.len(),
240        args.output.display()
241    );
242    let out_file =
243        File::create(&args.output).with_context(|| format!("create {}", args.output.display()))?;
244    writer.write_to(&mut BufWriter::new(out_file))?;
245    println!("Done.");
246    Ok(())
247}
248
249/// Download a repo's checkpoint in whatever format it's published in, without
250/// converting — for inspecting raw files, or when the target format is one
251/// `zsfm convert` doesn't need to touch (already GGUF).
252pub async fn run_pull(args: PullArgs) -> anyhow::Result<()> {
253    let model_dir = args
254        .model_dir
255        .clone()
256        .unwrap_or_else(|| PathBuf::from("models").join(args.repo.replace('/', "__")));
257    println!("Looking up files in {}@{} …", args.repo, args.revision);
258    let downloaded = zsfm_hub::download_any_format(
259        &args.repo,
260        args.format.as_deref(),
261        args.file.as_deref(),
262        &args.revision,
263        args.token.as_deref(),
264        &model_dir,
265    )
266    .await
267    .with_context(|| format!("download {}", args.repo))?;
268
269    println!(
270        "Downloaded {} file(s) in {} format:",
271        downloaded.checkpoint_files.len(),
272        downloaded.format
273    );
274    for f in &downloaded.checkpoint_files {
275        println!("  {}", f.display());
276    }
277    if let Some(cfg) = &downloaded.config_json {
278        println!("  {} (config.json)", cfg.display());
279    }
280    println!(
281        "\nTo convert: zsfm convert {} -o model.gguf --dtype f16",
282        downloaded
283            .checkpoint_files
284            .iter()
285            .map(|p| p.display().to_string())
286            .collect::<Vec<_>>()
287            .join(" ")
288    );
289    Ok(())
290}
291
292fn parse_meta_value(s: &str) -> GGUFMetaValue {
293    match s {
294        "true" => return GGUFMetaValue::Bool(true),
295        "false" => return GGUFMetaValue::Bool(false),
296        _ => {}
297    }
298    if let Ok(v) = s.parse::<u32>() {
299        return GGUFMetaValue::Uint32(v);
300    }
301    if let Ok(v) = s.parse::<i32>() {
302        return GGUFMetaValue::Int32(v);
303    }
304    if let Ok(v) = s.parse::<f32>() {
305        return GGUFMetaValue::Float32(v);
306    }
307    GGUFMetaValue::String(s.to_string())
308}
309
310pub fn run_inspect(args: InspectArgs) -> anyhow::Result<()> {
311    let ext = args
312        .path
313        .extension()
314        .and_then(|e| e.to_str())
315        .unwrap_or_default()
316        .to_ascii_lowercase();
317
318    // GGUF: show stored (possibly quantized) dtypes and metadata directly.
319    if ext == "gguf" {
320        let mut file = std::fs::File::open(&args.path)
321            .with_context(|| format!("open {}", args.path.display()))?;
322        // The in-workspace reader covers the dtypes this toolchain writes
323        // (F32/F16/BF16/Q8_0); anything else errors out below.
324        {
325            use std::io::Seek;
326            if let Ok(gguf) = zsfm_gguf::GGUFFile::read(&mut file) {
327                println!("Metadata in {}:", args.path.display());
328                for (key, value) in &gguf.metadata {
329                    println!("  {key:48} {value:?}");
330                }
331                println!("Tensors:");
332                let mut infos: Vec<&zsfm_gguf::GGUFTensorInfo> = gguf.tensors.iter().collect();
333                infos.sort_by(|a, b| a.name.cmp(&b.name));
334                for info in &infos {
335                    let py_shape: Vec<u64> = info.shape.iter().rev().copied().collect();
336                    println!("  {:64} {:?} {py_shape:?}", info.name, info.dtype);
337                }
338                if let Some(n) = args.values {
339                    for info in infos {
340                        if let Some(only) = &args.tensor {
341                            if info.name != *only {
342                                continue;
343                            }
344                        }
345                        let vals = gguf.tensor_f32(&mut file, info)?;
346                        let shown: Vec<String> =
347                            vals.iter().take(n).map(|v| v.to_string()).collect();
348                        println!(
349                            "  {} values[..{}]: [{}]",
350                            info.name,
351                            shown.len(),
352                            shown.join(", ")
353                        );
354                    }
355                }
356                return Ok(());
357            }
358            file.seek(std::io::SeekFrom::Start(0))?;
359        }
360        anyhow::bail!(
361            "cannot inspect {}: GGUF uses a quantization format outside F32/F16/BF16/Q8_0",
362            args.path.display()
363        );
364    }
365
366    let opts = LoadOptions {
367        pickle_key: args.pickle_key.clone(),
368        strip_prefix: None,
369    };
370    let ckpt = load_checkpoint(&[args.path.clone()], &opts)?;
371    println!("Tensors in {}:", args.path.display());
372    let mut tensors: Vec<_> = ckpt.tensors.iter().collect();
373    tensors.sort_by(|a, b| a.name.cmp(&b.name));
374    for t in &tensors {
375        println!("  {:64} {} {:?}", t.name, t.dtype.name(), t.shape);
376    }
377    if let Some(n) = args.values {
378        for t in tensors {
379            if let Some(only) = &args.tensor {
380                if &t.name != only {
381                    continue;
382                }
383            }
384            let vals = cast::decode_to_f32(&t.data, t.dtype)?;
385            let shown: Vec<String> = vals.iter().take(n).map(|v| v.to_string()).collect();
386            println!(
387                "  {} values[..{}]: [{}]",
388                t.name,
389                shown.len(),
390                shown.join(", ")
391            );
392        }
393    }
394    Ok(())
395}
zsfm/generic.rs

zsfm/
generic.rs
