zsfm_tabfm/
convert.rs

1use std::fs::File;
2use std::io::BufWriter;
3use std::path::Path;
4
5use anyhow::Context;
6use indicatif::{ProgressBar, ProgressStyle};
7use safetensors::SafeTensors;
8use safetensors::Dtype as StDtype;
9
10use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
11
12use crate::config::TabFMConfig;
13use crate::tensor_map::map_tensor_name;
14
15/// Options that control how tensors are stored in the output GGUF.
16pub struct ConvertOptions {
17    pub output_dtype: GGMLType,
18}
19
20/// Read a single (already `pytorch_model.bin` -> safetensors converted) checkpoint, convert
21/// every tensor, and write a GGUF file to `output_path`.
22pub fn convert(
23    model_id: &str,
24    safetensors_path: &Path,
25    config: &TabFMConfig,
26    opts: &ConvertOptions,
27    output_path: &Path,
28) -> anyhow::Result<()> {
29    let mut writer = GGUFWriter::new();
30
31    write_metadata(&mut writer, model_id, config);
32
33    let bytes = std::fs::read(safetensors_path)
34        .with_context(|| format!("read {}", safetensors_path.display()))?;
35    let view = SafeTensors::deserialize(&bytes).context("deserialize safetensors")?;
36
37    let total_tensors = view.len();
38    println!("Found {total_tensors} tensors.");
39
40    let pb = ProgressBar::new(total_tensors as u64);
41    pb.set_style(
42        ProgressStyle::with_template(
43            "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} {msg}",
44        )
45        .unwrap()
46        .progress_chars("=>-"),
47    );
48
49    let mut mapped = 0usize;
50    let mut skipped: Vec<String> = Vec::new();
51    let mut fallback_count = 0usize;
52
53    for (hf_name, tensor_view) in view.tensors() {
54        pb.set_message(hf_name.to_string());
55
56        let gguf_name = match map_tensor_name(&hf_name) {
57            Some(n) => n,
58            None => {
59                skipped.push(hf_name.to_string());
60                pb.inc(1);
61                continue;
62            }
63        };
64
65        let src_dtype = ggml_type_from_st(tensor_view.dtype())
66            .with_context(|| format!("tensor {hf_name}: unsupported dtype {:?}", tensor_view.dtype()))?;
67
68        let raw_data = tensor_view.data();
69        let py_shape = tensor_view.shape();
70        let n_elems: usize = py_shape.iter().product();
71        let innermost = py_shape.last().copied().unwrap_or(1);
72
73        // Q8_0: candle requires the innermost dim (last Python dim, first GGUF dim) to be
74        // divisible by 32. When it isn't, fall back to F32 for that tensor (biases, 1-D
75        // buffers like fourier frequencies / RoPE freqs / per_dim_scale, etc).
76        let (dst_dtype, gguf_shape, tensor_data) =
77            if opts.output_dtype == GGMLType::Q8_0 && (innermost % 32 != 0 || n_elems % 32 != 0) {
78                fallback_count += 1;
79                let data = cast_data(raw_data, src_dtype, GGMLType::F32)
80                    .with_context(|| format!("tensor {hf_name}: cast failed"))?;
81                let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
82                (GGMLType::F32, gs, data)
83            } else {
84                let dst = opts.output_dtype;
85                let data = cast_data(raw_data, src_dtype, dst)
86                    .with_context(|| format!("tensor {hf_name}: cast failed"))?;
87                let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
88                (dst, gs, data)
89            };
90
91        writer.add_tensor(gguf_name, gguf_shape, dst_dtype, tensor_data);
92        mapped += 1;
93        pb.inc(1);
94    }
95
96    pb.finish_with_message("tensors processed");
97
98    if !skipped.is_empty() {
99        eprintln!(
100            "\nWarning: {} tensor(s) had unrecognised names and were skipped:",
101            skipped.len()
102        );
103        for name in &skipped {
104            eprintln!("  {name}");
105        }
106        eprintln!("Update tensor_map.rs to include these if needed.");
107    }
108
109    if fallback_count > 0 {
110        eprintln!(
111            "\nNote: {fallback_count} tensor(s) fell back to F32 (scalars/biases/buffers too \
112             small for Q8_0 blocks)."
113        );
114    }
115    println!("Writing {mapped} tensors to {} …", output_path.display());
116    let out_file = File::create(output_path)
117        .with_context(|| format!("create output file {}", output_path.display()))?;
118    let mut buf_writer = BufWriter::new(out_file);
119    writer.write_to(&mut buf_writer)?;
120    println!("Done.");
121
122    Ok(())
123}
124
125fn ggml_type_from_st(dtype: StDtype) -> anyhow::Result<GGMLType> {
126    match dtype {
127        StDtype::F32 => Ok(GGMLType::F32),
128        StDtype::F16 => Ok(GGMLType::F16),
129        StDtype::BF16 => Ok(GGMLType::BF16),
130        other => anyhow::bail!("unsupported safetensors dtype: {other:?}"),
131    }
132}
133
134/// Build GGUF metadata section from parsed config.
135fn write_metadata(writer: &mut GGUFWriter, model_id: &str, config: &TabFMConfig) {
136    writer.add_metadata("general.architecture", GGUFMetaValue::String("tabfm".into()));
137    writer.add_metadata("general.name", GGUFMetaValue::String(model_id.into()));
138    writer.add_metadata("tabfm.is_classifier", GGUFMetaValue::Bool(config.is_classifier));
139    writer.add_metadata("tabfm.embedding_length", GGUFMetaValue::Uint32(config.embed_dim));
140    writer.add_metadata("tabfm.max_classes", GGUFMetaValue::Uint32(config.max_classes));
141    writer.add_metadata("tabfm.col_block_count", GGUFMetaValue::Uint32(config.col_num_blocks));
142    writer.add_metadata("tabfm.col_head_count", GGUFMetaValue::Uint32(config.col_nhead));
143    writer.add_metadata("tabfm.col_num_inds", GGUFMetaValue::Uint32(config.col_num_inds));
144    writer.add_metadata("tabfm.row_block_count", GGUFMetaValue::Uint32(config.row_num_blocks));
145    writer.add_metadata("tabfm.row_head_count", GGUFMetaValue::Uint32(config.row_nhead));
146    writer.add_metadata("tabfm.row_num_cls", GGUFMetaValue::Uint32(config.row_num_cls));
147    writer.add_metadata("tabfm.icl_block_count", GGUFMetaValue::Uint32(config.icl_num_blocks));
148    writer.add_metadata("tabfm.icl_head_count", GGUFMetaValue::Uint32(config.icl_nhead));
149    writer.add_metadata("tabfm.ff_factor", GGUFMetaValue::Uint32(config.ff_factor));
150    writer.add_metadata("tabfm.feature_group_size", GGUFMetaValue::Uint32(config.feature_group_size));
151    writer.add_metadata("tabfm.num_freq", GGUFMetaValue::Uint32(config.num_freq));
152    writer.add_metadata("tabfm.norm_eps", GGUFMetaValue::Float64(config.norm_eps));
153}
154
155/// Cast tensor bytes from `src_dtype` to `dst_dtype`.
156fn cast_data(data: &[u8], src: GGMLType, dst: GGMLType) -> anyhow::Result<Vec<u8>> {
157    if src == dst {
158        return Ok(data.to_vec());
159    }
160    if dst == GGMLType::Q8_0 {
161        let f32_values = decode_to_f32(data, src)?;
162        return quantize_q8_0(&f32_values);
163    }
164    match (src, dst) {
165        (GGMLType::F32, GGMLType::F16) => {
166            let f32_values = parse_f32_le(data)?;
167            let mut out = Vec::with_capacity(f32_values.len() * 2);
168            for v in f32_values {
169                let bits = f32_to_f16_bits(v);
170                out.extend_from_slice(&bits.to_le_bytes());
171            }
172            Ok(out)
173        }
174        (GGMLType::F32, GGMLType::BF16) => {
175            let f32_values = parse_f32_le(data)?;
176            let mut out = Vec::with_capacity(f32_values.len() * 2);
177            for v in f32_values {
178                let bits = (v.to_bits() >> 16) as u16;
179                out.extend_from_slice(&bits.to_le_bytes());
180            }
181            Ok(out)
182        }
183        (GGMLType::F16, GGMLType::BF16) => {
184            let mut out = Vec::with_capacity(data.len());
185            for chunk in data.chunks_exact(2) {
186                let f16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
187                let f32_val = f16_to_f32(f16_bits);
188                let bits = (f32_val.to_bits() >> 16) as u16;
189                out.extend_from_slice(&bits.to_le_bytes());
190            }
191            Ok(out)
192        }
193        (GGMLType::BF16, GGMLType::F32) => {
194            let mut out = Vec::with_capacity(data.len() * 2);
195            for chunk in data.chunks_exact(2) {
196                let bf16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
197                let f32_bits = (bf16_bits as u32) << 16;
198                out.extend_from_slice(&f32_bits.to_le_bytes());
199            }
200            Ok(out)
201        }
202        (GGMLType::BF16, GGMLType::F16) => {
203            let mut out = Vec::with_capacity(data.len());
204            for chunk in data.chunks_exact(2) {
205                let bf16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
206                let f32_bits = (bf16_bits as u32) << 16;
207                let f32_val = f32::from_bits(f32_bits);
208                let f16_bits = f32_to_f16_bits(f32_val);
209                out.extend_from_slice(&f16_bits.to_le_bytes());
210            }
211            Ok(out)
212        }
213        (GGMLType::F16, GGMLType::F32) => {
214            let mut out = Vec::with_capacity(data.len() * 2);
215            for chunk in data.chunks_exact(2) {
216                let f16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
217                let f32_val = f16_to_f32(f16_bits);
218                out.extend_from_slice(&f32_val.to_bits().to_le_bytes());
219            }
220            Ok(out)
221        }
222        _ => anyhow::bail!("unsupported cast: {src:?} → {dst:?}"),
223    }
224}
225
226fn decode_to_f32(data: &[u8], src: GGMLType) -> anyhow::Result<Vec<f32>> {
227    match src {
228        GGMLType::F32 => parse_f32_le(data),
229        GGMLType::F16 => data
230            .chunks_exact(2)
231            .map(|c| Ok(f16_to_f32(u16::from_le_bytes([c[0], c[1]]))))
232            .collect(),
233        GGMLType::BF16 => data
234            .chunks_exact(2)
235            .map(|c| {
236                let bf16_bits = u16::from_le_bytes([c[0], c[1]]);
237                Ok(f32::from_bits((bf16_bits as u32) << 16))
238            })
239            .collect(),
240        GGMLType::Q8_0 => anyhow::bail!("Q8_0 → Q8_0 re-quantization not supported as source"),
241    }
242}
243
244/// Quantize a slice of f32 values to GGML Q8_0 block format.
245/// Each block: `[f16 scale (2 bytes)][32 × i8 (32 bytes)]` = 34 bytes.
246fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
247    const BLOCK: usize = 32;
248    if values.len() % BLOCK != 0 {
249        anyhow::bail!(
250            "Q8_0 requires element count divisible by {BLOCK}, got {}",
251            values.len()
252        );
253    }
254    let n_blocks = values.len() / BLOCK;
255    let mut out = vec![0u8; n_blocks * 34];
256
257    for b in 0..n_blocks {
258        let blk = &values[b * BLOCK..(b + 1) * BLOCK];
259        let amax = blk.iter().copied().map(f32::abs).fold(0.0f32, f32::max);
260        let d = if amax == 0.0 { 0.0f32 } else { amax / 127.0 };
261        let d_inv = if d == 0.0 { 0.0f32 } else { 1.0 / d };
262
263        let base = b * 34;
264        let d_f16 = f32_to_f16_bits(d);
265        out[base..base + 2].copy_from_slice(&d_f16.to_le_bytes());
266        for i in 0..BLOCK {
267            let q = (blk[i] * d_inv).round().clamp(-127.0, 127.0) as i8;
268            out[base + 2 + i] = q as u8;
269        }
270    }
271    Ok(out)
272}
273
274fn parse_f32_le(data: &[u8]) -> anyhow::Result<Vec<f32>> {
275    if data.len() % 4 != 0 {
276        anyhow::bail!("f32 data length not divisible by 4");
277    }
278    Ok(data
279        .chunks_exact(4)
280        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
281        .collect())
282}
283
284fn f32_to_f16_bits(v: f32) -> u16 {
285    let bits = v.to_bits();
286    let sign = ((bits >> 16) & 0x8000) as u16;
287    let exp = ((bits >> 23) & 0xFF) as i32;
288    let mantissa = bits & 0x007F_FFFF;
289
290    if exp == 0xFF {
291        return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 };
292    }
293    let new_exp = exp - 127 + 15;
294    if new_exp >= 31 {
295        return sign | 0x7C00;
296    }
297    if new_exp <= 0 {
298        if new_exp < -10 {
299            return sign;
300        }
301        let m = (mantissa | 0x0080_0000) >> (1 - new_exp);
302        return sign | (m >> 13) as u16;
303    }
304    sign | ((new_exp as u16) << 10) | (mantissa >> 13) as u16
305}
306
307fn f16_to_f32(bits: u16) -> f32 {
308    let sign = ((bits & 0x8000) as u32) << 16;
309    let exp = ((bits >> 10) & 0x1F) as i32;
310    let mantissa = (bits & 0x03FF) as u32;
311
312    let f32_bits = if exp == 0 {
313        if mantissa == 0 {
314            sign
315        } else {
316            let mut m = mantissa;
317            let mut e = 0i32;
318            while m & 0x0400 == 0 {
319                m <<= 1;
320                e += 1;
321            }
322            sign | ((127 - 15 - e + 1) as u32) << 23 | (m & 0x03FF) << 13
323        }
324    } else if exp == 31 {
325        sign | 0x7F80_0000 | (mantissa << 13)
326    } else {
327        sign | ((exp + 127 - 15) as u32) << 23 | (mantissa << 13)
328    };
329    f32::from_bits(f32_bits)
330}
zsfm_tabfm/convert.rs

zsfm_tabfm/
convert.rs
