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};
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!("Found {} tensors across {} shard(s).", total_tensors, shard_views.len());
38
39 let pb = ProgressBar::new(total_tensors as u64);
40 pb.set_style(
41 ProgressStyle::with_template(
42 "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} {msg}",
43 )
44 .unwrap()
45 .progress_chars("=>-"),
46 );
47
48 let mut mapped = 0usize;
49 let mut skipped: Vec<String> = Vec::new();
50 let mut fallback_count = 0usize;
51
52 for shard in &shard_views {
53 for (hf_name, tensor_view) in shard.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 let (dst_dtype, gguf_shape, tensor_data) =
74 if opts.output_dtype == GGMLType::Q8_0 && (innermost % 32 != 0 || n_elems % 32 != 0) {
75 fallback_count += 1;
76 let data = cast_data(raw_data, src_dtype, GGMLType::F32)
77 .with_context(|| format!("tensor {hf_name}: cast failed"))?;
78 let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
79 (GGMLType::F32, gs, data)
80 } else {
81 let dst = opts.output_dtype;
82 let data = cast_data(raw_data, src_dtype, dst)
83 .with_context(|| format!("tensor {hf_name}: cast failed"))?;
84 let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
85 (dst, gs, data)
86 };
87
88 writer.add_tensor(gguf_name, gguf_shape, dst_dtype, tensor_data);
89 mapped += 1;
90 pb.inc(1);
91 }
92 }
93
94 pb.finish_with_message("tensors processed");
95
96 if !skipped.is_empty() {
97 eprintln!("\nWarning: {} tensor(s) skipped (not mapped):", skipped.len());
98 for name in &skipped {
99 eprintln!(" {name}");
100 }
101 }
102 if fallback_count > 0 {
103 eprintln!("\nNote: {fallback_count} tensor(s) fell back to F32 (too small for Q8_0 blocks).");
104 }
105
106 println!("Writing {mapped} tensors to {} …", output_path.display());
107 let out_file = File::create(output_path)
108 .with_context(|| format!("create {}", output_path.display()))?;
109 let mut buf_writer = BufWriter::new(out_file);
110 writer.write_to(&mut buf_writer)?;
111 println!("Done.");
112 Ok(())
113}
114
115fn write_metadata(writer: &mut GGUFWriter, model_id: &str, config: &FlowStateConfig) {
116 writer.add_metadata("general.architecture", GGUFMetaValue::String("flowstate".into()));
117 writer.add_metadata("general.name", GGUFMetaValue::String(model_id.into()));
118
119 writer.add_metadata("flowstate.block_count", GGUFMetaValue::Uint32(config.encoder_num_layers));
120 writer.add_metadata("flowstate.embedding_length", GGUFMetaValue::Uint32(config.embedding_feature_dim));
121 writer.add_metadata("flowstate.state_dim", GGUFMetaValue::Uint32(config.encoder_state_dim));
122 writer.add_metadata("flowstate.num_hippo_blocks", GGUFMetaValue::Uint32(config.encoder_num_hippo_blocks));
123 writer.add_metadata("flowstate.context_length", GGUFMetaValue::Uint32(config.context_length));
124 writer.add_metadata("flowstate.min_context", GGUFMetaValue::Uint32(config.min_context));
125 writer.add_metadata("flowstate.decoder_dim", GGUFMetaValue::Uint32(config.decoder_dim));
126 writer.add_metadata("flowstate.decoder_patch_len", GGUFMetaValue::Uint32(config.decoder_patch_len));
127 writer.add_metadata("flowstate.decoder_type", GGUFMetaValue::String(config.decoder_type.clone()));
128 writer.add_metadata("flowstate.quantile_count", GGUFMetaValue::Uint32(config.n_quantiles()));
129 writer.add_metadata("flowstate.quantiles", GGUFMetaValue::ArrayFloat32(config.quantiles.clone()));
130 writer.add_metadata("flowstate.with_missing", GGUFMetaValue::Bool(config.with_missing));
131
132 let range = config.basis_range();
133 writer.add_metadata("flowstate.basis_range_low", GGUFMetaValue::Float32(range[0]));
134 writer.add_metadata("flowstate.basis_range_high", GGUFMetaValue::Float32(range[1]));
135}
136
137fn ggml_type_from_st(dtype: StDtype) -> anyhow::Result<GGMLType> {
138 match dtype {
139 StDtype::F32 => Ok(GGMLType::F32),
140 StDtype::F16 => Ok(GGMLType::F16),
141 StDtype::BF16 => Ok(GGMLType::BF16),
142 other => anyhow::bail!("unsupported safetensors dtype: {other:?}"),
143 }
144}
145
146fn load_shard_bytes(shards: &[std::path::PathBuf]) -> anyhow::Result<Vec<Vec<u8>>> {
147 shards
148 .iter()
149 .map(|p| std::fs::read(p).with_context(|| format!("read shard {}", p.display())))
150 .collect()
151}
152
153fn cast_data(data: &[u8], src: GGMLType, dst: GGMLType) -> anyhow::Result<Vec<u8>> {
154 if src == dst { return Ok(data.to_vec()); }
155 if dst == GGMLType::Q8_0 {
156 let f32_values = decode_to_f32(data, src)?;
157 return quantize_q8_0(&f32_values);
158 }
159 match (src, dst) {
160 (GGMLType::F32, GGMLType::F16) => {
161 let f32_values = parse_f32_le(data)?;
162 let mut out = Vec::with_capacity(f32_values.len() * 2);
163 for v in f32_values { let bits = f32_to_f16_bits(v); out.extend_from_slice(&bits.to_le_bytes()); }
164 Ok(out)
165 }
166 (GGMLType::F32, GGMLType::BF16) => {
167 let f32_values = parse_f32_le(data)?;
168 let mut out = Vec::with_capacity(f32_values.len() * 2);
169 for v in f32_values { let bits = (v.to_bits() >> 16) as u16; out.extend_from_slice(&bits.to_le_bytes()); }
170 Ok(out)
171 }
172 (GGMLType::F16, GGMLType::BF16) => {
173 let mut out = Vec::with_capacity(data.len());
174 for chunk in data.chunks_exact(2) {
175 let f16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
176 let bits = (f16_to_f32(f16_bits).to_bits() >> 16) as u16;
177 out.extend_from_slice(&bits.to_le_bytes());
178 }
179 Ok(out)
180 }
181 (GGMLType::BF16, GGMLType::F32) => {
182 let mut out = Vec::with_capacity(data.len() * 2);
183 for chunk in data.chunks_exact(2) {
184 let bf16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
185 out.extend_from_slice(&((bf16_bits as u32) << 16).to_le_bytes());
186 }
187 Ok(out)
188 }
189 (GGMLType::BF16, GGMLType::F16) => {
190 let mut out = Vec::with_capacity(data.len());
191 for chunk in data.chunks_exact(2) {
192 let bf16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
193 let f32_val = f32::from_bits((bf16_bits as u32) << 16);
194 out.extend_from_slice(&f32_to_f16_bits(f32_val).to_le_bytes());
195 }
196 Ok(out)
197 }
198 (GGMLType::F16, GGMLType::F32) => {
199 let mut out = Vec::with_capacity(data.len() * 2);
200 for chunk in data.chunks_exact(2) {
201 let f16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
202 out.extend_from_slice(&f16_to_f32(f16_bits).to_bits().to_le_bytes());
203 }
204 Ok(out)
205 }
206 _ => anyhow::bail!("unsupported cast: {src:?} → {dst:?}"),
207 }
208}
209
210fn decode_to_f32(data: &[u8], src: GGMLType) -> anyhow::Result<Vec<f32>> {
211 match src {
212 GGMLType::F32 => parse_f32_le(data),
213 GGMLType::F16 => data.chunks_exact(2).map(|c| Ok(f16_to_f32(u16::from_le_bytes([c[0], c[1]])))).collect(),
214 GGMLType::BF16 => data.chunks_exact(2).map(|c| {
215 let bits = u16::from_le_bytes([c[0], c[1]]);
216 Ok(f32::from_bits((bits as u32) << 16))
217 }).collect(),
218 GGMLType::Q8_0 => anyhow::bail!("Q8_0 → Q8_0 re-quantization not supported"),
219 }
220}
221
222fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
223 const BLOCK: usize = 32;
224 if values.len() % BLOCK != 0 {
225 anyhow::bail!("Q8_0 requires element count divisible by {BLOCK}, got {}", values.len());
226 }
227 let n_blocks = values.len() / BLOCK;
228 let mut out = vec![0u8; n_blocks * 34];
229 for b in 0..n_blocks {
230 let blk = &values[b * BLOCK..(b + 1) * BLOCK];
231 let amax = blk.iter().copied().map(f32::abs).fold(0.0f32, f32::max);
232 let d = if amax == 0.0 { 0.0f32 } else { amax / 127.0 };
233 let d_inv = if d == 0.0 { 0.0f32 } else { 1.0 / d };
234 let base = b * 34;
235 out[base..base + 2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes());
236 for i in 0..BLOCK {
237 out[base + 2 + i] = (blk[i] * d_inv).round().clamp(-127.0, 127.0) as i8 as u8;
238 }
239 }
240 Ok(out)
241}
242
243fn parse_f32_le(data: &[u8]) -> anyhow::Result<Vec<f32>> {
244 if data.len() % 4 != 0 { anyhow::bail!("f32 data length not divisible by 4"); }
245 Ok(data.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect())
246}
247
248fn f32_to_f16_bits(v: f32) -> u16 {
249 let bits = v.to_bits();
250 let sign = ((bits >> 16) & 0x8000) as u16;
251 let exp = ((bits >> 23) & 0xFF) as i32;
252 let mantissa = bits & 0x007F_FFFF;
253 if exp == 0xFF { return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 }; }
254 let new_exp = exp - 127 + 15;
255 if new_exp >= 31 { return sign | 0x7C00; }
256 if new_exp <= 0 {
257 if new_exp < -10 { return sign; }
258 let m = (mantissa | 0x0080_0000) >> (1 - new_exp);
259 return sign | (m >> 13) as u16;
260 }
261 sign | ((new_exp as u16) << 10) | (mantissa >> 13) as u16
262}
263
264fn f16_to_f32(bits: u16) -> f32 {
265 let sign = ((bits & 0x8000) as u32) << 16;
266 let exp = ((bits >> 10) & 0x1F) as i32;
267 let mantissa = (bits & 0x03FF) as u32;
268 let f32_bits = if exp == 0 {
269 if mantissa == 0 { sign }
270 else {
271 let mut m = mantissa; let mut e = 0i32;
272 while m & 0x0400 == 0 { m <<= 1; e += 1; }
273 sign | ((127 - 15 - e + 1) as u32) << 23 | (m & 0x03FF) << 13
274 }
275 } else if exp == 31 {
276 sign | 0x7F80_0000 | (mantissa << 13)
277 } else {
278 sign | ((exp + 127 - 15) as u32) << 23 | (mantissa << 13)
279 };
280 f32::from_bits(f32_bits)
281}