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