zsfm/
common.rs

1//! Helpers shared by more than one model's CLI module.
2
3use std::path::Path;
4
5use anyhow::Result;
6use clap::ValueEnum;
7use zsfm_checkpoint::{load_checkpoint, LoadOptions};
8use zsfm_gguf::GGMLType;
9
10/// Shared `--dtype` flag. BF16 is deliberately absent: the in-workspace GGUF
11/// reader used by every model's own `infer` loader handles F32/F16/Q8_0, so a
12/// bf16 file would fail to load back through the same model's `infer` command.
13/// The format-agnostic `zsfm convert` path supports BF16 for other consumers.
14#[derive(Clone, Copy, Debug, ValueEnum)]
15pub enum DtypeArg {
16    F32,
17    F16,
18    Q8,
19}
20
21impl DtypeArg {
22    pub fn name(self) -> &'static str {
23        match self {
24            DtypeArg::F32 => "f32",
25            DtypeArg::F16 => "f16",
26            DtypeArg::Q8 => "q8",
27        }
28    }
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::Q8 => GGMLType::Q8_0,
37        }
38    }
39}
40
41/// Return the canonical cache path, recasting from cache when present.
42///
43/// Returns `true` when the caller can return early (cache hit already
44/// recast to `output`), `false` when the caller must download + convert.
45/// Thin wrapper over [`zsfm_hub::try_recast_from_cache`] so CLI + Python share
46/// one implementation.
47pub fn try_recast_from_cache(
48    canonical: &Path,
49    output: &Path,
50    dtype: GGMLType,
51    redownload: bool,
52) -> Result<bool> {
53    zsfm_hub::try_recast_from_cache(canonical, output, dtype, redownload)
54}
55
56/// Index of the median (q0.5) quantile level, with center fallback.
57pub fn median_index(quantile_levels: &[f32]) -> usize {
58    quantile_levels
59        .iter()
60        .position(|&q| (q - 0.5).abs() < 1e-6)
61        .unwrap_or(quantile_levels.len() / 2)
62}
63
64/// Shared `inspect-tensors` implementation: print every tensor's name, shape,
65/// and dtype from a local checkpoint file. Format is auto-detected — works
66/// whether the model's raw download is safetensors (most models), a raw
67/// PyTorch pickle `.ckpt` (lag_llama, tirex), or a GGUF file, via the same
68/// format-agnostic loader `zsfm convert`/`zsfm inspect` use.
69pub fn inspect_tensors(path: &Path) -> Result<()> {
70    let ckpt = load_checkpoint(&[path.to_path_buf()], &LoadOptions::default())?;
71    eprintln!("Tensors in {}:", path.display());
72    let mut tensors: Vec<_> = ckpt.tensors.iter().collect();
73    tensors.sort_by(|a, b| a.name.cmp(&b.name));
74    for t in tensors {
75        println!("  {:80} {} {:?}", t.name, t.dtype.name(), t.shape);
76    }
77    Ok(())
78}
79
80/// Delete a model's cached files (canonical GGUF + config) and optionally an
81/// output GGUF. Thin wrapper over [`zsfm_hub::delete_cached_model`] so all
82/// CLI modules share one import path (`crate::common`).
83pub fn delete_cached_model(canonical: &Path, output: Option<&Path>) -> Result<()> {
84    zsfm_hub::delete_cached_model(canonical, output)
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn dtype_names() {
93        assert_eq!(DtypeArg::F32.name(), "f32");
94        assert_eq!(DtypeArg::F16.name(), "f16");
95        assert_eq!(DtypeArg::Q8.name(), "q8");
96        assert_eq!(GGMLType::from(DtypeArg::F32), GGMLType::F32);
97        assert_eq!(GGMLType::from(DtypeArg::Q8), GGMLType::Q8_0);
98    }
99
100    #[test]
101    fn median_prefers_q05() {
102        let levels = vec![0.1, 0.5, 0.9];
103        assert_eq!(median_index(&levels), 1);
104        let no_median = vec![0.1, 0.9];
105        assert_eq!(median_index(&no_median), 1);
106        let empty: Vec<f32> = vec![];
107        assert_eq!(median_index(&empty), 0);
108    }
109
110    #[test]
111    fn cache_miss_returns_false() {
112        let dir = std::env::temp_dir().join(format!("zsfm-common-{}", std::process::id()));
113        let canonical = dir.join("owner__repo").join("model-f32.gguf");
114        let out = dir.join("out.gguf");
115        // No file exists → miss, no I/O.
116        assert!(!try_recast_from_cache(&canonical, &out, GGMLType::F16, false).unwrap());
117        // Redownload forces miss even if file existed (still missing here).
118        assert!(!try_recast_from_cache(&canonical, &out, GGMLType::F16, true).unwrap());
119    }
120}
zsfm/common.rs

zsfm/
common.rs
