1use std::fs::File;
2use std::io::BufWriter;
3use std::path::{Path, PathBuf};
4
5use anyhow::Context;
6use indicatif::{ProgressBar, ProgressStyle};
7use safetensors::Dtype as StDtype;
8use safetensors::SafeTensors;
9
10use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
11
12use crate::config::TabDptConfig;
13use crate::tensor_map::map_tensor_name;
14
15pub struct ConvertOptions {
16 pub output_dtype: GGMLType,
17}
18
19pub fn convert(
20 shard_paths: &[PathBuf],
21 config: &TabDptConfig,
22 opts: &ConvertOptions,
23 output_path: &Path,
24) -> anyhow::Result<()> {
25 let mut writer = GGUFWriter::new();
26 write_metadata(&mut writer, config);
27
28 let shard_bytes: Vec<Vec<u8>> = shard_paths
29 .iter()
30 .map(|p| std::fs::read(p).with_context(|| format!("read {}", p.display())))
31 .collect::<anyhow::Result<_>>()?;
32 let shard_views: Vec<SafeTensors> = shard_bytes
33 .iter()
34 .map(|b| SafeTensors::deserialize(b).context("deserialize shard"))
35 .collect::<anyhow::Result<_>>()?;
36
37 let total: usize = shard_views.iter().map(|s| s.len()).sum();
38 let pb = ProgressBar::new(total as u64);
39 pb.set_style(
40 ProgressStyle::with_template(
41 "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} {msg}",
42 )
43 .unwrap()
44 .progress_chars("=>-"),
45 );
46
47 let mut mapped = 0usize;
48 let mut skipped: Vec<String> = Vec::new();
49 let mut fallback_count = 0usize;
50
51 for shard in &shard_views {
52 for (hf_name, tensor_view) in shard.tensors() {
53 pb.set_message(hf_name.to_string());
54
55 let gguf_name = match map_tensor_name(&hf_name) {
56 Some(n) => n,
57 None => {
58 skipped.push(hf_name.to_string());
59 pb.inc(1);
60 continue;
61 }
62 };
63
64 let src_dtype = ggml_type_from_st(tensor_view.dtype()).with_context(|| {
65 format!(
66 "tensor {hf_name}: unsupported dtype {:?}",
67 tensor_view.dtype()
68 )
69 })?;
70
71 let raw_data = tensor_view.data();
72 let py_shape = tensor_view.shape();
73 let n_elems: usize = py_shape.iter().product();
74 let innermost = py_shape.last().copied().unwrap_or(1);
75
76 let (dst_dtype, gguf_shape, tensor_data) = if opts.output_dtype == GGMLType::Q8_0
77 && (innermost % 32 != 0 || n_elems % 32 != 0)
78 {
79 fallback_count += 1;
80 let data = cast_data(raw_data, src_dtype, GGMLType::F32)
81 .with_context(|| format!("tensor {hf_name}: cast failed"))?;
82 let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
83 (GGMLType::F32, gs, data)
84 } else {
85 let dst = opts.output_dtype;
86 let data = cast_data(raw_data, src_dtype, dst)
87 .with_context(|| format!("tensor {hf_name}: cast failed"))?;
88 let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
89 (dst, gs, data)
90 };
91
92 writer.add_tensor(gguf_name, gguf_shape, dst_dtype, tensor_data);
93 mapped += 1;
94 pb.inc(1);
95 }
96 }
97
98 pb.finish_with_message("tensors processed");
99
100 if !skipped.is_empty() {
101 eprintln!("\nNote: {} tensor(s) skipped (scalar attention-scaling buffers, recomputed from config):", skipped.len());
102 for name in &skipped {
103 eprintln!(" {name}");
104 }
105 }
106 if fallback_count > 0 {
107 eprintln!("\nNote: {fallback_count} tensor(s) fell back to F32.");
108 }
109
110 println!("Writing {mapped} tensors to {} …", output_path.display());
111 let out_file =
112 File::create(output_path).with_context(|| format!("create {}", output_path.display()))?;
113 let mut buf_writer = BufWriter::new(out_file);
114 writer.write_to(&mut buf_writer)?;
115 println!("Done.");
116 Ok(())
117}
118
119fn write_metadata(writer: &mut GGUFWriter, config: &TabDptConfig) {
120 writer.add_metadata(
121 "general.architecture",
122 GGUFMetaValue::String("tabdpt".into()),
123 );
124 writer.add_metadata(
125 "general.name",
126 GGUFMetaValue::String("Layer6/TabDPT".into()),
127 );
128 writer.add_metadata("tabdpt.dim", GGUFMetaValue::Uint32(config.dim as u32));
129 writer.add_metadata(
130 "tabdpt.n_layers",
131 GGUFMetaValue::Uint32(config.n_layers as u32),
132 );
133 writer.add_metadata(
134 "tabdpt.n_heads",
135 GGUFMetaValue::Uint32(config.n_heads as u32),
136 );
137 writer.add_metadata("tabdpt.ff_dim", GGUFMetaValue::Uint32(config.ff_dim as u32));
138 writer.add_metadata(
139 "tabdpt.y_encoder_dim",
140 GGUFMetaValue::Uint32(config.y_encoder_dim as u32),
141 );
142 writer.add_metadata(
143 "tabdpt.max_num_classes",
144 GGUFMetaValue::Uint32(config.max_num_classes as u32),
145 );
146 writer.add_metadata(
147 "tabdpt.regression_bin_count",
148 GGUFMetaValue::Uint32(config.regression_bin_count as u32),
149 );
150 writer.add_metadata(
151 "tabdpt.regression_bin_min",
152 GGUFMetaValue::Float32(config.regression_bin_min),
153 );
154 writer.add_metadata(
155 "tabdpt.regression_bin_max",
156 GGUFMetaValue::Float32(config.regression_bin_max),
157 );
158 writer.add_metadata(
159 "tabdpt.max_num_features",
160 GGUFMetaValue::Uint32(config.max_num_features as u32),
161 );
162 writer.add_metadata(
163 "tabdpt.base_len",
164 GGUFMetaValue::Uint32(config.base_len as u32),
165 );
166 writer.add_metadata(
167 "tabdpt.max_len",
168 GGUFMetaValue::Uint32(config.max_len as u32),
169 );
170 writer.add_metadata(
171 "tabdpt.n_thinking_rows",
172 GGUFMetaValue::Uint32(config.n_thinking_rows as u32),
173 );
174}
175
176fn ggml_type_from_st(dtype: StDtype) -> anyhow::Result<GGMLType> {
177 match dtype {
178 StDtype::F32 => Ok(GGMLType::F32),
179 StDtype::F16 => Ok(GGMLType::F16),
180 StDtype::BF16 => Ok(GGMLType::BF16),
181 other => anyhow::bail!("unsupported safetensors dtype: {other:?}"),
182 }
183}
184
185fn cast_data(data: &[u8], src: GGMLType, dst: GGMLType) -> anyhow::Result<Vec<u8>> {
186 if src == dst {
187 return Ok(data.to_vec());
188 }
189 if dst == GGMLType::Q8_0 {
190 let f32_values = decode_to_f32(data, src)?;
191 return quantize_q8_0(&f32_values);
192 }
193 match (src, dst) {
194 (GGMLType::F32, GGMLType::F16) => {
195 let vals = parse_f32_le(data)?;
196 let mut out = Vec::with_capacity(vals.len() * 2);
197 for v in vals {
198 out.extend_from_slice(&f32_to_f16_bits(v).to_le_bytes());
199 }
200 Ok(out)
201 }
202 (GGMLType::F32, GGMLType::BF16) => {
203 let vals = parse_f32_le(data)?;
204 let mut out = Vec::with_capacity(vals.len() * 2);
205 for v in vals {
206 out.extend_from_slice(&((v.to_bits() >> 16) as u16).to_le_bytes());
207 }
208 Ok(out)
209 }
210 (GGMLType::F16, GGMLType::BF16) => {
211 let mut out = Vec::with_capacity(data.len());
212 for c in data.chunks_exact(2) {
213 let f32_val = f16_to_f32(u16::from_le_bytes([c[0], c[1]]));
214 out.extend_from_slice(&((f32_val.to_bits() >> 16) as u16).to_le_bytes());
215 }
216 Ok(out)
217 }
218 (GGMLType::BF16, GGMLType::F32) => {
219 let mut out = Vec::with_capacity(data.len() * 2);
220 for c in data.chunks_exact(2) {
221 let bits = (u16::from_le_bytes([c[0], c[1]]) as u32) << 16;
222 out.extend_from_slice(&bits.to_le_bytes());
223 }
224 Ok(out)
225 }
226 (GGMLType::BF16, GGMLType::F16) => {
227 let mut out = Vec::with_capacity(data.len());
228 for c in data.chunks_exact(2) {
229 let bits = (u16::from_le_bytes([c[0], c[1]]) as u32) << 16;
230 out.extend_from_slice(&f32_to_f16_bits(f32::from_bits(bits)).to_le_bytes());
231 }
232 Ok(out)
233 }
234 (GGMLType::F16, GGMLType::F32) => {
235 let mut out = Vec::with_capacity(data.len() * 2);
236 for c in data.chunks_exact(2) {
237 out.extend_from_slice(
238 &f16_to_f32(u16::from_le_bytes([c[0], c[1]]))
239 .to_bits()
240 .to_le_bytes(),
241 );
242 }
243 Ok(out)
244 }
245 _ => anyhow::bail!("unsupported cast: {src:?} → {dst:?}"),
246 }
247}
248
249fn decode_to_f32(data: &[u8], src: GGMLType) -> anyhow::Result<Vec<f32>> {
250 match src {
251 GGMLType::F32 => parse_f32_le(data),
252 GGMLType::F16 => data
253 .chunks_exact(2)
254 .map(|c| Ok(f16_to_f32(u16::from_le_bytes([c[0], c[1]]))))
255 .collect(),
256 GGMLType::BF16 => data
257 .chunks_exact(2)
258 .map(|c| {
259 Ok(f32::from_bits(
260 (u16::from_le_bytes([c[0], c[1]]) as u32) << 16,
261 ))
262 })
263 .collect(),
264 GGMLType::Q8_0 => anyhow::bail!("Q8_0 as source not supported"),
265 }
266}
267
268fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
269 const BLOCK: usize = 32;
270 if values.len() % BLOCK != 0 {
271 anyhow::bail!("Q8_0 requires count divisible by {BLOCK}");
272 }
273 let n_blocks = values.len() / BLOCK;
274 let mut out = vec![0u8; n_blocks * 34];
275 for b in 0..n_blocks {
276 let blk = &values[b * BLOCK..(b + 1) * BLOCK];
277 let amax = blk.iter().copied().map(f32::abs).fold(0.0f32, f32::max);
278 let d = if amax == 0.0 { 0.0f32 } else { amax / 127.0 };
279 let d_inv = if d == 0.0 { 0.0f32 } else { 1.0 / d };
280 let base = b * 34;
281 out[base..base + 2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes());
282 for i in 0..BLOCK {
283 out[base + 2 + i] = (blk[i] * d_inv).round().clamp(-127.0, 127.0) as i8 as u8;
284 }
285 }
286 Ok(out)
287}
288
289fn parse_f32_le(data: &[u8]) -> anyhow::Result<Vec<f32>> {
290 if data.len() % 4 != 0 {
291 anyhow::bail!("f32 data length not divisible by 4");
292 }
293 Ok(data
294 .chunks_exact(4)
295 .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
296 .collect())
297}
298
299fn f32_to_f16_bits(v: f32) -> u16 {
300 let bits = v.to_bits();
301 let sign = ((bits >> 16) & 0x8000) as u16;
302 let exp = ((bits >> 23) & 0xFF) as i32;
303 let mantissa = bits & 0x007F_FFFF;
304 if exp == 0xFF {
305 return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 };
306 }
307 let new_exp = exp - 127 + 15;
308 if new_exp >= 31 {
309 return sign | 0x7C00;
310 }
311 if new_exp <= 0 {
312 if new_exp < -10 {
313 return sign;
314 }
315 let m = (mantissa | 0x0080_0000) >> (1 - new_exp);
316 return sign | (m >> 13) as u16;
317 }
318 sign | ((new_exp as u16) << 10) | (mantissa >> 13) as u16
319}
320
321fn f16_to_f32(bits: u16) -> f32 {
322 let sign = ((bits & 0x8000) as u32) << 16;
323 let exp = ((bits >> 10) & 0x1F) as i32;
324 let mantissa = (bits & 0x03FF) as u32;
325 let f32_bits = if exp == 0 {
326 if mantissa == 0 {
327 sign
328 } else {
329 let mut m = mantissa;
330 let mut e = 0i32;
331 while m & 0x0400 == 0 {
332 m <<= 1;
333 e += 1;
334 }
335 sign | ((127 - 15 - e + 1) as u32) << 23 | (m & 0x03FF) << 13
336 }
337 } else if exp == 31 {
338 sign | 0x7F80_0000 | (mantissa << 13)
339 } else {
340 sign | ((exp + 127 - 15) as u32) << 23 | (mantissa << 13)
341 };
342 f32::from_bits(f32_bits)
343}