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::FlowStateConfig;
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: &FlowStateConfig,
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
81 let (dst_dtype, gguf_shape, tensor_data) = if opts.output_dtype == GGMLType::Q8_0
82 && (innermost % 32 != 0 || n_elems % 32 != 0)
83 {
84 fallback_count += 1;
85 let data = cast_data(raw_data, src_dtype, GGMLType::F32)
86 .with_context(|| format!("tensor {hf_name}: cast failed"))?;
87 let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
88 (GGMLType::F32, gs, data)
89 } else {
90 let dst = opts.output_dtype;
91 let data = cast_data(raw_data, src_dtype, dst)
92 .with_context(|| format!("tensor {hf_name}: cast failed"))?;
93 let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
94 (dst, gs, data)
95 };
96
97 writer.add_tensor(gguf_name, gguf_shape, dst_dtype, tensor_data);
98 mapped += 1;
99 pb.inc(1);
100 }
101 }
102
103 pb.finish_with_message("tensors processed");
104
105 if !skipped.is_empty() {
106 eprintln!(
107 "\nWarning: {} tensor(s) skipped (not mapped):",
108 skipped.len()
109 );
110 for name in &skipped {
111 eprintln!(" {name}");
112 }
113 }
114 if fallback_count > 0 {
115 eprintln!(
116 "\nNote: {fallback_count} tensor(s) fell back to F32 (too small for Q8_0 blocks)."
117 );
118 }
119
120 println!("Writing {mapped} tensors to {} …", output_path.display());
121 let out_file =
122 File::create(output_path).with_context(|| format!("create {}", output_path.display()))?;
123 let mut buf_writer = BufWriter::new(out_file);
124 writer.write_to(&mut buf_writer)?;
125 println!("Done.");
126 Ok(())
127}
128
129fn write_metadata(writer: &mut GGUFWriter, model_id: &str, config: &FlowStateConfig) {
130 writer.add_metadata(
131 "general.architecture",
132 GGUFMetaValue::String("flowstate".into()),
133 );
134 writer.add_metadata("general.name", GGUFMetaValue::String(model_id.into()));
135
136 writer.add_metadata(
137 "flowstate.block_count",
138 GGUFMetaValue::Uint32(config.encoder_num_layers),
139 );
140 writer.add_metadata(
141 "flowstate.embedding_length",
142 GGUFMetaValue::Uint32(config.embedding_feature_dim),
143 );
144 writer.add_metadata(
145 "flowstate.state_dim",
146 GGUFMetaValue::Uint32(config.encoder_state_dim),
147 );
148 writer.add_metadata(
149 "flowstate.num_hippo_blocks",
150 GGUFMetaValue::Uint32(config.encoder_num_hippo_blocks),
151 );
152 writer.add_metadata(
153 "flowstate.context_length",
154 GGUFMetaValue::Uint32(config.context_length),
155 );
156 writer.add_metadata(
157 "flowstate.min_context",
158 GGUFMetaValue::Uint32(config.min_context),
159 );
160 writer.add_metadata(
161 "flowstate.decoder_dim",
162 GGUFMetaValue::Uint32(config.decoder_dim),
163 );
164 writer.add_metadata(
165 "flowstate.decoder_patch_len",
166 GGUFMetaValue::Uint32(config.decoder_patch_len),
167 );
168 writer.add_metadata(
169 "flowstate.decoder_type",
170 GGUFMetaValue::String(config.decoder_type.clone()),
171 );
172 writer.add_metadata(
173 "flowstate.quantile_count",
174 GGUFMetaValue::Uint32(config.n_quantiles()),
175 );
176 writer.add_metadata(
177 "flowstate.quantiles",
178 GGUFMetaValue::ArrayFloat32(config.quantiles.clone()),
179 );
180 writer.add_metadata(
181 "flowstate.with_missing",
182 GGUFMetaValue::Bool(config.with_missing),
183 );
184
185 let range = config.basis_range();
186 writer.add_metadata(
187 "flowstate.basis_range_low",
188 GGUFMetaValue::Float32(range[0]),
189 );
190 writer.add_metadata(
191 "flowstate.basis_range_high",
192 GGUFMetaValue::Float32(range[1]),
193 );
194}
195
196fn ggml_type_from_st(dtype: StDtype) -> anyhow::Result<GGMLType> {
197 match dtype {
198 StDtype::F32 => Ok(GGMLType::F32),
199 StDtype::F16 => Ok(GGMLType::F16),
200 StDtype::BF16 => Ok(GGMLType::BF16),
201 other => anyhow::bail!("unsupported safetensors dtype: {other:?}"),
202 }
203}
204
205fn load_shard_bytes(shards: &[std::path::PathBuf]) -> anyhow::Result<Vec<Vec<u8>>> {
206 shards
207 .iter()
208 .map(|p| std::fs::read(p).with_context(|| format!("read shard {}", p.display())))
209 .collect()
210}
211
212fn cast_data(data: &[u8], src: GGMLType, dst: GGMLType) -> anyhow::Result<Vec<u8>> {
213 if src == dst {
214 return Ok(data.to_vec());
215 }
216 if dst == GGMLType::Q8_0 {
217 let f32_values = decode_to_f32(data, src)?;
218 return quantize_q8_0(&f32_values);
219 }
220 match (src, dst) {
221 (GGMLType::F32, GGMLType::F16) => {
222 let f32_values = parse_f32_le(data)?;
223 let mut out = Vec::with_capacity(f32_values.len() * 2);
224 for v in f32_values {
225 let bits = f32_to_f16_bits(v);
226 out.extend_from_slice(&bits.to_le_bytes());
227 }
228 Ok(out)
229 }
230 (GGMLType::F32, GGMLType::BF16) => {
231 let f32_values = parse_f32_le(data)?;
232 let mut out = Vec::with_capacity(f32_values.len() * 2);
233 for v in f32_values {
234 let bits = (v.to_bits() >> 16) as u16;
235 out.extend_from_slice(&bits.to_le_bytes());
236 }
237 Ok(out)
238 }
239 (GGMLType::F16, GGMLType::BF16) => {
240 let mut out = Vec::with_capacity(data.len());
241 for chunk in data.chunks_exact(2) {
242 let f16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
243 let bits = (f16_to_f32(f16_bits).to_bits() >> 16) as u16;
244 out.extend_from_slice(&bits.to_le_bytes());
245 }
246 Ok(out)
247 }
248 (GGMLType::BF16, GGMLType::F32) => {
249 let mut out = Vec::with_capacity(data.len() * 2);
250 for chunk in data.chunks_exact(2) {
251 let bf16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
252 out.extend_from_slice(&((bf16_bits as u32) << 16).to_le_bytes());
253 }
254 Ok(out)
255 }
256 (GGMLType::BF16, GGMLType::F16) => {
257 let mut out = Vec::with_capacity(data.len());
258 for chunk in data.chunks_exact(2) {
259 let bf16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
260 let f32_val = f32::from_bits((bf16_bits as u32) << 16);
261 out.extend_from_slice(&f32_to_f16_bits(f32_val).to_le_bytes());
262 }
263 Ok(out)
264 }
265 (GGMLType::F16, GGMLType::F32) => {
266 let mut out = Vec::with_capacity(data.len() * 2);
267 for chunk in data.chunks_exact(2) {
268 let f16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
269 out.extend_from_slice(&f16_to_f32(f16_bits).to_bits().to_le_bytes());
270 }
271 Ok(out)
272 }
273 _ => anyhow::bail!("unsupported cast: {src:?} → {dst:?}"),
274 }
275}
276
277fn decode_to_f32(data: &[u8], src: GGMLType) -> anyhow::Result<Vec<f32>> {
278 match src {
279 GGMLType::F32 => parse_f32_le(data),
280 GGMLType::F16 => data
281 .chunks_exact(2)
282 .map(|c| Ok(f16_to_f32(u16::from_le_bytes([c[0], c[1]]))))
283 .collect(),
284 GGMLType::BF16 => data
285 .chunks_exact(2)
286 .map(|c| {
287 let bits = u16::from_le_bytes([c[0], c[1]]);
288 Ok(f32::from_bits((bits as u32) << 16))
289 })
290 .collect(),
291 GGMLType::Q8_0 => anyhow::bail!("Q8_0 → Q8_0 re-quantization not supported"),
292 }
293}
294
295fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
296 const BLOCK: usize = 32;
297 if values.len() % BLOCK != 0 {
298 anyhow::bail!(
299 "Q8_0 requires element count divisible by {BLOCK}, got {}",
300 values.len()
301 );
302 }
303 let n_blocks = values.len() / BLOCK;
304 let mut out = vec![0u8; n_blocks * 34];
305 for b in 0..n_blocks {
306 let blk = &values[b * BLOCK..(b + 1) * BLOCK];
307 let amax = blk.iter().copied().map(f32::abs).fold(0.0f32, f32::max);
308 let d = if amax == 0.0 { 0.0f32 } else { amax / 127.0 };
309 let d_inv = if d == 0.0 { 0.0f32 } else { 1.0 / d };
310 let base = b * 34;
311 out[base..base + 2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes());
312 for i in 0..BLOCK {
313 out[base + 2 + i] = (blk[i] * d_inv).round().clamp(-127.0, 127.0) as i8 as u8;
314 }
315 }
316 Ok(out)
317}
318
319fn parse_f32_le(data: &[u8]) -> anyhow::Result<Vec<f32>> {
320 if data.len() % 4 != 0 {
321 anyhow::bail!("f32 data length not divisible by 4");
322 }
323 Ok(data
324 .chunks_exact(4)
325 .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
326 .collect())
327}
328
329fn f32_to_f16_bits(v: f32) -> u16 {
330 let bits = v.to_bits();
331 let sign = ((bits >> 16) & 0x8000) as u16;
332 let exp = ((bits >> 23) & 0xFF) as i32;
333 let mantissa = bits & 0x007F_FFFF;
334 if exp == 0xFF {
335 return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 };
336 }
337 let new_exp = exp - 127 + 15;
338 if new_exp >= 31 {
339 return sign | 0x7C00;
340 }
341 if new_exp <= 0 {
342 if new_exp < -10 {
343 return sign;
344 }
345 let m = (mantissa | 0x0080_0000) >> (1 - new_exp);
346 return sign | (m >> 13) as u16;
347 }
348 sign | ((new_exp as u16) << 10) | (mantissa >> 13) as u16
349}
350
351fn f16_to_f32(bits: u16) -> f32 {
352 let sign = ((bits & 0x8000) as u32) << 16;
353 let exp = ((bits >> 10) & 0x1F) as i32;
354 let mantissa = (bits & 0x03FF) as u32;
355 let f32_bits = if exp == 0 {
356 if mantissa == 0 {
357 sign
358 } else {
359 let mut m = mantissa;
360 let mut e = 0i32;
361 while m & 0x0400 == 0 {
362 m <<= 1;
363 e += 1;
364 }
365 sign | ((127 - 15 - e + 1) as u32) << 23 | (m & 0x03FF) << 13
366 }
367 } else if exp == 31 {
368 sign | 0x7F80_0000 | (mantissa << 13)
369 } else {
370 sign | ((exp + 127 - 15) as u32) << 23 | (mantissa << 13)
371 };
372 f32::from_bits(f32_bits)
373}