| 1 | //! GGUF writer module. |
| 2 | //! |
| 3 | //! Converts a merged safetensors model directory (containing `config.json`, |
| 4 | //! `tokenizer.json`, and `model.safetensors`) into a single `.gguf` file |
| 5 | //! compatible with mistral.rs / llama.cpp. |
| 6 | |
| 7 | use std::collections::HashMap; |
| 8 | use std::io::{BufWriter, Write}; |
| 9 | use std::path::PathBuf; |
| 10 | |
| 11 | use anyhow::{Context, Result}; |
| 12 | use candle_core::Device; |
| 13 | use candle_core::safetensors::MmapedSafetensors; |
| 14 | use half::f16; |
| 15 | |
| 16 | // --------------------------------------------------------------------------- |
| 17 | // Public types |
| 18 | // --------------------------------------------------------------------------- |
| 19 | |
| 20 | /// Target data type for the exported GGUF tensors. |
| 21 | #[allow(dead_code)] |
| 22 | pub enum GgufDtype { |
| 23 | F16, |
| 24 | Q8_0, |
| 25 | } |
| 26 | |
| 27 | /// Configuration for a GGUF export run. |
| 28 | pub struct GgufConfig { |
| 29 | /// Directory containing `config.json`, `tokenizer.json`, and |
| 30 | /// `model.safetensors`. |
| 31 | pub model_dir: PathBuf, |
| 32 | /// Path where the output `.gguf` file will be written. |
| 33 | pub output_path: PathBuf, |
| 34 | /// Target quantisation type for 2-D weight tensors. |
| 35 | pub dtype: GgufDtype, |
| 36 | } |
| 37 | |
| 38 | /// Progress events emitted by the background export thread. |
| 39 | pub enum GgufProgress { |
| 40 | ReadingModel, |
| 41 | WritingTensor { |
| 42 | index: usize, |
| 43 | total: usize, |
| 44 | name: String, |
| 45 | }, |
| 46 | Done { |
| 47 | output_path: PathBuf, |
| 48 | size_bytes: u64, |
| 49 | }, |
| 50 | Failed(String), |
| 51 | } |
| 52 | |
| 53 | // --------------------------------------------------------------------------- |
| 54 | // Public entry point |
| 55 | // --------------------------------------------------------------------------- |
| 56 | |
| 57 | pub fn start_gguf_export(config: GgufConfig, tx: tokio::sync::mpsc::UnboundedSender<GgufProgress>) { |
| 58 | std::thread::spawn(move || { |
| 59 | if let Err(e) = run_gguf_export(&config, &tx) { |
| 60 | eprintln!("[gguf] error: {e:#}"); |
| 61 | let _ = tx.send(GgufProgress::Failed(format!("{e:#}"))); |
| 62 | } |
| 63 | }); |
| 64 | } |
| 65 | |
| 66 | // --------------------------------------------------------------------------- |
| 67 | // Private — model / tokenizer config structs |
| 68 | // --------------------------------------------------------------------------- |
| 69 | |
| 70 | #[derive(serde::Deserialize)] |
| 71 | struct ModelConfig { |
| 72 | hidden_size: usize, |
| 73 | num_hidden_layers: usize, |
| 74 | num_attention_heads: usize, |
| 75 | num_key_value_heads: Option<usize>, |
| 76 | /// Per-head attention dimension. Decoupled from `hidden_size / num_attention_heads` |
| 77 | /// in Qwen3 (e.g. Qwen3-0.6B: hidden_size 1024, 16 heads, but head_dim 128). |
| 78 | /// Absent in Qwen2/2.5 configs, where it equals `hidden_size / num_attention_heads`. |
| 79 | #[serde(default)] |
| 80 | head_dim: Option<usize>, |
| 81 | intermediate_size: usize, |
| 82 | vocab_size: usize, |
| 83 | #[serde(default = "default_rms_norm_eps")] |
| 84 | rms_norm_eps: f64, |
| 85 | #[serde(default = "default_rope_theta")] |
| 86 | rope_theta: f64, |
| 87 | #[serde(default = "default_max_position_embeddings")] |
| 88 | max_position_embeddings: usize, |
| 89 | #[serde(default)] |
| 90 | tie_word_embeddings: bool, |
| 91 | /// HuggingFace model type, e.g. "qwen2" or "qwen3". |
| 92 | /// Kept for future architecture-specific logic. |
| 93 | #[allow(dead_code)] |
| 94 | #[serde(default = "default_model_type")] |
| 95 | model_type: String, |
| 96 | } |
| 97 | |
| 98 | fn default_rms_norm_eps() -> f64 { |
| 99 | 1e-6 |
| 100 | } |
| 101 | fn default_rope_theta() -> f64 { |
| 102 | 10000.0 |
| 103 | } |
| 104 | fn default_max_position_embeddings() -> usize { |
| 105 | 32768 |
| 106 | } |
| 107 | fn default_model_type() -> String { |
| 108 | "qwen2".to_string() |
| 109 | } |
| 110 | |
| 111 | #[derive(serde::Deserialize)] |
| 112 | struct TokenizerJson { |
| 113 | model: TokenizerModel, |
| 114 | #[serde(default)] |
| 115 | added_tokens: Vec<AddedToken>, |
| 116 | } |
| 117 | |
| 118 | #[derive(serde::Deserialize)] |
| 119 | struct TokenizerModel { |
| 120 | vocab: HashMap<String, u32>, |
| 121 | #[serde(default, deserialize_with = "deserialize_merges")] |
| 122 | merges: Vec<String>, |
| 123 | } |
| 124 | |
| 125 | /// A single merge entry — either `"a b"` (Qwen2.5) or `["a","b"]` (Qwen3). |
| 126 | #[derive(serde::Deserialize)] |
| 127 | #[serde(untagged)] |
| 128 | enum MergeEntry { |
| 129 | Str(String), |
| 130 | Pair(Vec<String>), |
| 131 | } |
| 132 | |
| 133 | /// Deserialize merges from either format into a uniform `Vec<String>` of |
| 134 | /// space-separated pairs (the format GGUF expects). |
| 135 | fn deserialize_merges<'de, D>(deserializer: D) -> Result<Vec<String>, D::Error> |
| 136 | where |
| 137 | D: serde::Deserializer<'de>, |
| 138 | { |
| 139 | use serde::Deserialize; |
| 140 | let entries: Vec<MergeEntry> = Vec::deserialize(deserializer)?; |
| 141 | Ok(entries |
| 142 | .into_iter() |
| 143 | .map(|e| match e { |
| 144 | MergeEntry::Str(s) => s, |
| 145 | MergeEntry::Pair(parts) => parts.join(" "), |
| 146 | }) |
| 147 | .collect()) |
| 148 | } |
| 149 | |
| 150 | #[derive(serde::Deserialize)] |
| 151 | struct AddedToken { |
| 152 | id: u32, |
| 153 | content: String, |
| 154 | special: bool, |
| 155 | } |
| 156 | |
| 157 | // --------------------------------------------------------------------------- |
| 158 | // GGUF constants |
| 159 | // --------------------------------------------------------------------------- |
| 160 | |
| 161 | const GGUF_MAGIC: u32 = 0x46554747; |
| 162 | const GGUF_VERSION: u32 = 3; |
| 163 | const ALIGNMENT: u64 = 32; |
| 164 | |
| 165 | /// Standard Qwen ChatML template used as a fallback when `tokenizer_config.json` |
| 166 | /// does not provide one. Matches the template shipped in Bartowski's Qwen GGUFs. |
| 167 | const QWEN_CHATML_TEMPLATE: &str = concat!( |
| 168 | "{%- for message in messages %}", |
| 169 | "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}", |
| 170 | "{%- endfor %}", |
| 171 | "{%- if add_generation_prompt %}", |
| 172 | "{{'<|im_start|>assistant\n'}}", |
| 173 | "{%- endif %}", |
| 174 | ); |
| 175 | |
| 176 | // GGUF tensor type IDs |
| 177 | const GGUF_TYPE_F32: u32 = 0; |
| 178 | const GGUF_TYPE_F16: u32 = 1; |
| 179 | const GGUF_TYPE_Q8_0: u32 = 8; |
| 180 | |
| 181 | // GGUF metadata value type IDs |
| 182 | const GGUF_META_U32: u32 = 4; |
| 183 | const GGUF_META_INT32: u32 = 5; |
| 184 | const GGUF_META_F32: u32 = 6; |
| 185 | const GGUF_META_STRING: u32 = 8; |
| 186 | const GGUF_META_ARRAY: u32 = 9; |
| 187 | |
| 188 | // --------------------------------------------------------------------------- |
| 189 | // Private — GGUF write helpers |
| 190 | // --------------------------------------------------------------------------- |
| 191 | |
| 192 | fn write_u32(w: &mut impl Write, v: u32) -> std::io::Result<()> { |
| 193 | w.write_all(&v.to_le_bytes()) |
| 194 | } |
| 195 | |
| 196 | fn write_u64(w: &mut impl Write, v: u64) -> std::io::Result<()> { |
| 197 | w.write_all(&v.to_le_bytes()) |
| 198 | } |
| 199 | |
| 200 | fn write_f32(w: &mut impl Write, v: f32) -> std::io::Result<()> { |
| 201 | w.write_all(&v.to_le_bytes()) |
| 202 | } |
| 203 | |
| 204 | fn write_gguf_string(w: &mut impl Write, s: &str) -> std::io::Result<()> { |
| 205 | write_u64(w, s.len() as u64)?; |
| 206 | w.write_all(s.as_bytes()) |
| 207 | } |
| 208 | |
| 209 | fn write_kv_string(w: &mut impl Write, key: &str, val: &str) -> std::io::Result<()> { |
| 210 | write_gguf_string(w, key)?; |
| 211 | write_u32(w, GGUF_META_STRING)?; |
| 212 | write_gguf_string(w, val) |
| 213 | } |
| 214 | |
| 215 | fn write_kv_u32(w: &mut impl Write, key: &str, val: u32) -> std::io::Result<()> { |
| 216 | write_gguf_string(w, key)?; |
| 217 | write_u32(w, GGUF_META_U32)?; |
| 218 | write_u32(w, val) |
| 219 | } |
| 220 | |
| 221 | fn write_kv_f32(w: &mut impl Write, key: &str, val: f32) -> std::io::Result<()> { |
| 222 | write_gguf_string(w, key)?; |
| 223 | write_u32(w, GGUF_META_F32)?; |
| 224 | write_f32(w, val) |
| 225 | } |
| 226 | |
| 227 | fn write_kv_string_array(w: &mut impl Write, key: &str, vals: &[String]) -> std::io::Result<()> { |
| 228 | write_gguf_string(w, key)?; |
| 229 | write_u32(w, GGUF_META_ARRAY)?; |
| 230 | write_u32(w, GGUF_META_STRING)?; |
| 231 | write_u64(w, vals.len() as u64)?; |
| 232 | for s in vals { |
| 233 | write_gguf_string(w, s)?; |
| 234 | } |
| 235 | Ok(()) |
| 236 | } |
| 237 | |
| 238 | /// Write an array of signed 32-bit integers. GGUF declares `token_type` as an |
| 239 | /// INT32 array; llama.cpp / mistral.rs read it with `to_i32()` and reject a |
| 240 | /// UINT32-typed array, so the element type tag must be INT32 (not U32). |
| 241 | fn write_kv_i32_array(w: &mut impl Write, key: &str, vals: &[i32]) -> std::io::Result<()> { |
| 242 | write_gguf_string(w, key)?; |
| 243 | write_u32(w, GGUF_META_ARRAY)?; |
| 244 | write_u32(w, GGUF_META_INT32)?; |
| 245 | write_u64(w, vals.len() as u64)?; |
| 246 | for &v in vals { |
| 247 | w.write_all(&v.to_le_bytes())?; |
| 248 | } |
| 249 | Ok(()) |
| 250 | } |
| 251 | |
| 252 | fn pad_to_alignment(w: &mut impl Write, pos: u64, alignment: u64) -> std::io::Result<u64> { |
| 253 | let rem = pos % alignment; |
| 254 | if rem != 0 { |
| 255 | let pad = alignment - rem; |
| 256 | for _ in 0..pad { |
| 257 | w.write_all(&[0u8])?; |
| 258 | } |
| 259 | Ok(pos + pad) |
| 260 | } else { |
| 261 | Ok(pos) |
| 262 | } |
| 263 | } |
| 264 | |
| 265 | // --------------------------------------------------------------------------- |
| 266 | // Private — tensor name mapping (HuggingFace → GGUF) |
| 267 | // --------------------------------------------------------------------------- |
| 268 | |
| 269 | fn map_tensor_name(hf_name: &str) -> Option<String> { |
| 270 | // Non-layer names |
| 271 | if hf_name == "model.embed_tokens.weight" { |
| 272 | return Some("token_embd.weight".to_string()); |
| 273 | } |
| 274 | if hf_name == "model.norm.weight" { |
| 275 | return Some("output_norm.weight".to_string()); |
| 276 | } |
| 277 | if hf_name == "lm_head.weight" { |
| 278 | return Some("output.weight".to_string()); |
| 279 | } |
| 280 | |
| 281 | // Layer names: model.layers.{i}.suffix → blk.{i}.gguf_suffix |
| 282 | let prefix = "model.layers."; |
| 283 | if !hf_name.starts_with(prefix) { |
| 284 | return None; |
| 285 | } |
| 286 | |
| 287 | let rest = &hf_name[prefix.len()..]; |
| 288 | let dot_pos = rest.find('.')?; |
| 289 | let layer_idx = &rest[..dot_pos]; |
| 290 | let suffix = &rest[dot_pos + 1..]; |
| 291 | |
| 292 | let gguf_suffix = match suffix { |
| 293 | "self_attn.q_proj.weight" => "attn_q.weight", |
| 294 | "self_attn.k_proj.weight" => "attn_k.weight", |
| 295 | "self_attn.v_proj.weight" => "attn_v.weight", |
| 296 | "self_attn.o_proj.weight" => "attn_output.weight", |
| 297 | // Bias tensors (present in Qwen2/2.5, absent in Qwen3) |
| 298 | "self_attn.q_proj.bias" => "attn_q.bias", |
| 299 | "self_attn.k_proj.bias" => "attn_k.bias", |
| 300 | "self_attn.v_proj.bias" => "attn_v.bias", |
| 301 | // QK-Norm (Qwen3) |
| 302 | "self_attn.q_norm.weight" => "attn_q_norm.weight", |
| 303 | "self_attn.k_norm.weight" => "attn_k_norm.weight", |
| 304 | "mlp.gate_proj.weight" => "ffn_gate.weight", |
| 305 | "mlp.up_proj.weight" => "ffn_up.weight", |
| 306 | "mlp.down_proj.weight" => "ffn_down.weight", |
| 307 | "input_layernorm.weight" => "attn_norm.weight", |
| 308 | "post_attention_layernorm.weight" => "ffn_norm.weight", |
| 309 | _ => return None, |
| 310 | }; |
| 311 | |
| 312 | Some(format!("blk.{layer_idx}.{gguf_suffix}")) |
| 313 | } |
| 314 | |
| 315 | // --------------------------------------------------------------------------- |
| 316 | // Private — tensor data size calculation |
| 317 | // --------------------------------------------------------------------------- |
| 318 | |
| 319 | fn tensor_data_size(num_elements: usize, gguf_type: u32) -> usize { |
| 320 | match gguf_type { |
| 321 | GGUF_TYPE_F32 => num_elements * 4, |
| 322 | GGUF_TYPE_F16 => num_elements * 2, |
| 323 | GGUF_TYPE_Q8_0 => num_elements.div_ceil(32) * 34, |
| 324 | _ => num_elements * 4, |
| 325 | } |
| 326 | } |
| 327 | |
| 328 | // --------------------------------------------------------------------------- |
| 329 | // Private — determine GGUF type for a tensor |
| 330 | // --------------------------------------------------------------------------- |
| 331 | |
| 332 | fn choose_gguf_type(gguf_name: &str, ndim: usize, dtype: &GgufDtype) -> u32 { |
| 333 | // 1-D tensors (norms, biases) → always F32 |
| 334 | if ndim <= 1 { |
| 335 | return GGUF_TYPE_F32; |
| 336 | } |
| 337 | |
| 338 | // Embedding and output head → always F16 |
| 339 | if gguf_name == "token_embd.weight" || gguf_name == "output.weight" { |
| 340 | return GGUF_TYPE_F16; |
| 341 | } |
| 342 | |
| 343 | // 2-D weight tensors → follow requested dtype |
| 344 | match dtype { |
| 345 | GgufDtype::F16 => GGUF_TYPE_F16, |
| 346 | GgufDtype::Q8_0 => GGUF_TYPE_Q8_0, |
| 347 | } |
| 348 | } |
| 349 | |
| 350 | // --------------------------------------------------------------------------- |
| 351 | // Private — tensor data conversion and writing |
| 352 | // --------------------------------------------------------------------------- |
| 353 | |
| 354 | fn write_tensor_f32(w: &mut impl Write, data: &[f32]) -> std::io::Result<()> { |
| 355 | for &v in data { |
| 356 | w.write_all(&v.to_le_bytes())?; |
| 357 | } |
| 358 | Ok(()) |
| 359 | } |
| 360 | |
| 361 | fn write_tensor_f16(w: &mut impl Write, data: &[f32]) -> std::io::Result<()> { |
| 362 | for &v in data { |
| 363 | w.write_all(&f16::from_f32(v).to_le_bytes())?; |
| 364 | } |
| 365 | Ok(()) |
| 366 | } |
| 367 | |
| 368 | fn write_tensor_q8_0(w: &mut impl Write, data: &[f32]) -> std::io::Result<()> { |
| 369 | let block_size = 32; |
| 370 | let num_blocks = data.len().div_ceil(block_size); |
| 371 | |
| 372 | for block_idx in 0..num_blocks { |
| 373 | let start = block_idx * block_size; |
| 374 | let end = (start + block_size).min(data.len()); |
| 375 | let block = &data[start..end]; |
| 376 | |
| 377 | // Find max absolute value in the block |
| 378 | let mut amax: f32 = 0.0; |
| 379 | for &v in block { |
| 380 | let abs = v.abs(); |
| 381 | if abs > amax { |
| 382 | amax = abs; |
| 383 | } |
| 384 | } |
| 385 | |
| 386 | let scale = if amax == 0.0 { 0.0 } else { amax / 127.0 }; |
| 387 | let inv_scale = if scale == 0.0 { 0.0 } else { 1.0 / scale }; |
| 388 | |
| 389 | // Write scale as f16 (2 bytes) |
| 390 | w.write_all(&f16::from_f32(scale).to_le_bytes())?; |
| 391 | |
| 392 | // Quantize and write 32 i8 values |
| 393 | for i in 0..block_size { |
| 394 | let val = if i < block.len() { block[i] } else { 0.0 }; |
| 395 | let q = (val * inv_scale).round().clamp(-128.0, 127.0) as i8; |
| 396 | w.write_all(&[q as u8])?; |
| 397 | } |
| 398 | } |
| 399 | Ok(()) |
| 400 | } |
| 401 | |
| 402 | // --------------------------------------------------------------------------- |
| 403 | // Private — compute byte sizes of metadata |
| 404 | // --------------------------------------------------------------------------- |
| 405 | |
| 406 | /// Compute the byte size of a GGUF string (length prefix + bytes). |
| 407 | fn gguf_string_size(s: &str) -> u64 { |
| 408 | 8 + s.len() as u64 |
| 409 | } |
| 410 | |
| 411 | /// Compute the byte size of a KV string entry. |
| 412 | fn kv_string_size(key: &str, val: &str) -> u64 { |
| 413 | gguf_string_size(key) + 4 + gguf_string_size(val) |
| 414 | } |
| 415 | |
| 416 | /// Compute the byte size of a KV u32 entry. |
| 417 | fn kv_u32_size(key: &str) -> u64 { |
| 418 | gguf_string_size(key) + 4 + 4 |
| 419 | } |
| 420 | |
| 421 | /// Compute the byte size of a KV f32 entry. |
| 422 | fn kv_f32_size(key: &str) -> u64 { |
| 423 | gguf_string_size(key) + 4 + 4 |
| 424 | } |
| 425 | |
| 426 | /// Compute the byte size of a KV string-array entry. |
| 427 | fn kv_string_array_size(key: &str, vals: &[String]) -> u64 { |
| 428 | let mut size = gguf_string_size(key) + 4; // key + value_type |
| 429 | size += 4 + 8; // array element_type + count |
| 430 | for s in vals { |
| 431 | size += gguf_string_size(s); |
| 432 | } |
| 433 | size |
| 434 | } |
| 435 | |
| 436 | /// Compute the byte size of a KV i32-array entry (same layout as u32). |
| 437 | fn kv_i32_array_size(key: &str, vals: &[i32]) -> u64 { |
| 438 | gguf_string_size(key) + 4 + 4 + 8 + (vals.len() as u64 * 4) |
| 439 | } |
| 440 | |
| 441 | /// Compute the byte size of a single tensor info entry. |
| 442 | fn tensor_info_size(name: &str, ndim: usize) -> u64 { |
| 443 | // name_len(8) + name_bytes + n_dims(4) + dims(8*ndim) + type(4) + offset(8) |
| 444 | 8 + name.len() as u64 + 4 + (ndim as u64 * 8) + 4 + 8 |
| 445 | } |
| 446 | |
| 447 | // --------------------------------------------------------------------------- |
| 448 | // Private — core export logic |
| 449 | // --------------------------------------------------------------------------- |
| 450 | |
| 451 | /// Holds the info we need for each tensor in the GGUF file. |
| 452 | struct TensorEntry { |
| 453 | /// GGUF tensor name. |
| 454 | gguf_name: String, |
| 455 | /// Original HuggingFace tensor name (for loading from safetensors). |
| 456 | hf_name: String, |
| 457 | /// Dimensions in GGUF order (row-major, i.e. same order as safetensors). |
| 458 | dims: Vec<u64>, |
| 459 | /// Number of dimensions. |
| 460 | ndim: usize, |
| 461 | /// Total number of elements. |
| 462 | num_elements: usize, |
| 463 | /// GGUF tensor type ID. |
| 464 | gguf_type: u32, |
| 465 | /// Size in bytes of the tensor data. |
| 466 | data_size: usize, |
| 467 | /// Offset from start of data section. |
| 468 | offset: u64, |
| 469 | } |
| 470 | |
| 471 | fn run_gguf_export( |
| 472 | config: &GgufConfig, |
| 473 | tx: &tokio::sync::mpsc::UnboundedSender<GgufProgress>, |
| 474 | ) -> Result<()> { |
| 475 | let _ = tx.send(GgufProgress::ReadingModel); |
| 476 | |
| 477 | // ----------------------------------------------------------------------- |
| 478 | // 1. Load model config |
| 479 | // ----------------------------------------------------------------------- |
| 480 | let config_path = config.model_dir.join("config.json"); |
| 481 | let config_text = std::fs::read_to_string(&config_path) |
| 482 | .with_context(|| format!("failed to read {}", config_path.display()))?; |
| 483 | let model_cfg: ModelConfig = |
| 484 | serde_json::from_str(&config_text).context("failed to parse config.json")?; |
| 485 | |
| 486 | // ----------------------------------------------------------------------- |
| 487 | // 2. Load and parse tokenizer.json |
| 488 | // ----------------------------------------------------------------------- |
| 489 | let tok_path = config.model_dir.join("tokenizer.json"); |
| 490 | let tok_text = std::fs::read_to_string(&tok_path) |
| 491 | .with_context(|| format!("failed to read {}", tok_path.display()))?; |
| 492 | let tok: TokenizerJson = |
| 493 | serde_json::from_str(&tok_text).context("failed to parse tokenizer.json")?; |
| 494 | |
| 495 | // Read the chat template from tokenizer_config.json. |
| 496 | // This is needed by mistral.rs / llama.cpp to format multi-turn conversations. |
| 497 | // |
| 498 | // The field can be either: |
| 499 | // - a plain string (Qwen2/2.5) |
| 500 | // - an array of {"name": "...", "template": "..."} objects (Qwen3) |
| 501 | // |
| 502 | // When it's an array, pick the "default" entry, or fall back to the first. |
| 503 | // If the file is missing or the field is absent, use the standard Qwen |
| 504 | // ChatML template so the GGUF is always loadable for chat. |
| 505 | let chat_template: String = { |
| 506 | let tok_config_path = config.model_dir.join("tokenizer_config.json"); |
| 507 | let from_file = std::fs::read_to_string(&tok_config_path) |
| 508 | .ok() |
| 509 | .and_then(|text| { |
| 510 | let v: serde_json::Value = serde_json::from_str(&text).ok()?; |
| 511 | let ct = v.get("chat_template")?; |
| 512 | |
| 513 | // Plain string |
| 514 | if let Some(s) = ct.as_str() { |
| 515 | return Some(s.to_string()); |
| 516 | } |
| 517 | |
| 518 | // Array of objects — pick "default", else first entry |
| 519 | if let Some(arr) = ct.as_array() { |
| 520 | let default = arr.iter().find(|entry| { |
| 521 | entry.get("name").and_then(|n| n.as_str()) == Some("default") |
| 522 | }); |
| 523 | let chosen = default.or_else(|| arr.first()); |
| 524 | return chosen |
| 525 | .and_then(|entry| entry.get("template")) |
| 526 | .and_then(|t| t.as_str()) |
| 527 | .map(String::from); |
| 528 | } |
| 529 | |
| 530 | None |
| 531 | }); |
| 532 | |
| 533 | from_file.unwrap_or_else(|| QWEN_CHATML_TEMPLATE.to_string()) |
| 534 | }; |
| 535 | |
| 536 | // Build ordered token list sorted by ID |
| 537 | let vocab_size = model_cfg.vocab_size; |
| 538 | let mut id_to_token: Vec<(u32, String)> = |
| 539 | tok.model.vocab.into_iter().map(|(k, v)| (v, k)).collect(); |
| 540 | // Also include added_tokens that might not be in the base vocab |
| 541 | for at in &tok.added_tokens { |
| 542 | if !id_to_token.iter().any(|(id, _)| *id == at.id) { |
| 543 | id_to_token.push((at.id, at.content.clone())); |
| 544 | } |
| 545 | } |
| 546 | id_to_token.sort_by_key(|(id, _)| *id); |
| 547 | |
| 548 | // Pad to vocab_size if needed |
| 549 | let mut tokens: Vec<String> = Vec::with_capacity(vocab_size); |
| 550 | let mut next_expected: u32 = 0; |
| 551 | for (id, content) in &id_to_token { |
| 552 | while next_expected < *id && (next_expected as usize) < vocab_size { |
| 553 | tokens.push(format!("<unused{next_expected}>")); |
| 554 | next_expected += 1; |
| 555 | } |
| 556 | if (next_expected as usize) < vocab_size { |
| 557 | tokens.push(content.clone()); |
| 558 | next_expected += 1; |
| 559 | } |
| 560 | } |
| 561 | while tokens.len() < vocab_size { |
| 562 | let idx = tokens.len(); |
| 563 | tokens.push(format!("<unused{idx}>")); |
| 564 | } |
| 565 | tokens.truncate(vocab_size); |
| 566 | |
| 567 | // Build special token lookup |
| 568 | let special_set: HashMap<u32, bool> = tok |
| 569 | .added_tokens |
| 570 | .iter() |
| 571 | .filter(|at| at.special) |
| 572 | .map(|at| (at.id, true)) |
| 573 | .collect(); |
| 574 | |
| 575 | // Token types: 1 = normal, 3 = control (special). Stored as INT32 in GGUF. |
| 576 | let token_types: Vec<i32> = (0..vocab_size as u32) |
| 577 | .map(|id| if special_set.contains_key(&id) { 3 } else { 1 }) |
| 578 | .collect(); |
| 579 | |
| 580 | // BOS / EOS token IDs |
| 581 | let added_token_id = |name: &str| -> Option<u32> { |
| 582 | tok.added_tokens |
| 583 | .iter() |
| 584 | .find(|at| at.content == name) |
| 585 | .map(|at| at.id) |
| 586 | }; |
| 587 | let bos_id = added_token_id("<|endoftext|>").unwrap_or(0); |
| 588 | let eos_id = added_token_id("<|im_end|>") |
| 589 | .or_else(|| added_token_id("<|endoftext|>")) |
| 590 | .unwrap_or(0); |
| 591 | |
| 592 | let merges = tok.model.merges; |
| 593 | |
| 594 | // ----------------------------------------------------------------------- |
| 595 | // 3. Load safetensors |
| 596 | // ----------------------------------------------------------------------- |
| 597 | let st_path = config.model_dir.join("model.safetensors"); |
| 598 | // SAFETY: the file is memory-mapped read-only. |
| 599 | let safetensors = unsafe { MmapedSafetensors::new(&st_path) } |
| 600 | .with_context(|| format!("failed to mmap {}", st_path.display()))?; |
| 601 | let tensor_list = safetensors.tensors(); |
| 602 | |
| 603 | // ----------------------------------------------------------------------- |
| 604 | // 4. Build tensor entries with name mapping |
| 605 | // ----------------------------------------------------------------------- |
| 606 | let mut entries: Vec<TensorEntry> = Vec::new(); |
| 607 | for (hf_name, tensor_view) in &tensor_list { |
| 608 | let gguf_name = match map_tensor_name(hf_name) { |
| 609 | Some(n) => n, |
| 610 | None => continue, |
| 611 | }; |
| 612 | |
| 613 | let shape = tensor_view.shape(); |
| 614 | let ndim = shape.len(); |
| 615 | let num_elements: usize = shape.iter().product(); |
| 616 | // GGUF stores dimensions innermost-first (llama.cpp convention); the |
| 617 | // reader reverses them back to logical order. safetensors shapes are |
| 618 | // outermost-first, so write them reversed. For a Linear weight stored |
| 619 | // as [out, in] this yields ne = [in, out], matching what mistral.rs / |
| 620 | // candle expect after their `dimensions.reverse()`. |
| 621 | let dims: Vec<u64> = shape.iter().rev().map(|&d| d as u64).collect(); |
| 622 | let gguf_type = choose_gguf_type(&gguf_name, ndim, &config.dtype); |
| 623 | let data_size = tensor_data_size(num_elements, gguf_type); |
| 624 | |
| 625 | entries.push(TensorEntry { |
| 626 | gguf_name, |
| 627 | hf_name: hf_name.clone(), |
| 628 | dims, |
| 629 | ndim, |
| 630 | num_elements, |
| 631 | gguf_type, |
| 632 | data_size, |
| 633 | offset: 0, // computed below |
| 634 | }); |
| 635 | } |
| 636 | |
| 637 | // Handle tie_word_embeddings: if lm_head.weight is missing, duplicate |
| 638 | // embed_tokens as output.weight |
| 639 | let has_output = entries.iter().any(|e| e.gguf_name == "output.weight"); |
| 640 | if !has_output |
| 641 | && model_cfg.tie_word_embeddings |
| 642 | && let Some(embd) = entries.iter().find(|e| e.gguf_name == "token_embd.weight") |
| 643 | { |
| 644 | let gguf_type = choose_gguf_type("output.weight", embd.ndim, &config.dtype); |
| 645 | let data_size = tensor_data_size(embd.num_elements, gguf_type); |
| 646 | entries.push(TensorEntry { |
| 647 | gguf_name: "output.weight".to_string(), |
| 648 | hf_name: "model.embed_tokens.weight".to_string(), |
| 649 | dims: embd.dims.clone(), |
| 650 | ndim: embd.ndim, |
| 651 | num_elements: embd.num_elements, |
| 652 | gguf_type, |
| 653 | data_size, |
| 654 | offset: 0, |
| 655 | }); |
| 656 | } |
| 657 | |
| 658 | // Synthesize zero-filled Q/K/V bias tensors for layers that lack them. |
| 659 | // |
| 660 | // mistralrs expects `blk.{i}.attn_q.bias` etc. for the qwen2 architecture. |
| 661 | // Qwen2/2.5 have these in the safetensors; Qwen3 does not. Bartowski's |
| 662 | // llama.cpp-converted Qwen3 GGUFs include zero-filled bias tensors. |
| 663 | // We do the same here so the GGUF loads correctly for both model families. |
| 664 | let num_layers = model_cfg.num_hidden_layers; |
| 665 | let num_heads = model_cfg.num_attention_heads; |
| 666 | let num_kv_heads_usize = model_cfg.num_key_value_heads.unwrap_or(num_heads); |
| 667 | // Qwen3 decouples head_dim from hidden_size / num_heads; honour the explicit |
| 668 | // config value when present so synthesized bias tensors and the GGUF |
| 669 | // key_length/value_length metadata match the real attention dimension. |
| 670 | let head_dim = model_cfg |
| 671 | .head_dim |
| 672 | .unwrap_or(model_cfg.hidden_size / num_heads); |
| 673 | |
| 674 | for i in 0..num_layers { |
| 675 | let bias_specs: &[(&str, usize)] = &[ |
| 676 | (&format!("blk.{i}.attn_q.bias"), num_heads * head_dim), |
| 677 | ( |
| 678 | &format!("blk.{i}.attn_k.bias"), |
| 679 | num_kv_heads_usize * head_dim, |
| 680 | ), |
| 681 | ( |
| 682 | &format!("blk.{i}.attn_v.bias"), |
| 683 | num_kv_heads_usize * head_dim, |
| 684 | ), |
| 685 | ]; |
| 686 | for (gguf_name, num_elements) in bias_specs { |
| 687 | if !entries.iter().any(|e| e.gguf_name == *gguf_name) { |
| 688 | // 1-D bias → always F32 |
| 689 | let data_size = tensor_data_size(*num_elements, GGUF_TYPE_F32); |
| 690 | entries.push(TensorEntry { |
| 691 | gguf_name: gguf_name.to_string(), |
| 692 | hf_name: String::new(), // sentinel: no safetensors source |
| 693 | dims: vec![*num_elements as u64], |
| 694 | ndim: 1, |
| 695 | num_elements: *num_elements, |
| 696 | gguf_type: GGUF_TYPE_F32, |
| 697 | data_size, |
| 698 | offset: 0, |
| 699 | }); |
| 700 | } |
| 701 | } |
| 702 | } |
| 703 | |
| 704 | // Compute offsets (relative to start of data section, 32-byte aligned) |
| 705 | let mut data_offset: u64 = 0; |
| 706 | for entry in &mut entries { |
| 707 | entry.offset = data_offset; |
| 708 | let size = entry.data_size as u64; |
| 709 | data_offset += size; |
| 710 | // Pad to 32-byte alignment |
| 711 | let rem = data_offset % ALIGNMENT; |
| 712 | if rem != 0 { |
| 713 | data_offset += ALIGNMENT - rem; |
| 714 | } |
| 715 | } |
| 716 | |
| 717 | // ----------------------------------------------------------------------- |
| 718 | // 5. Compute header size to know where data section starts |
| 719 | // ----------------------------------------------------------------------- |
| 720 | let num_tensors = entries.len() as u64; |
| 721 | let num_kv_heads = model_cfg |
| 722 | .num_key_value_heads |
| 723 | .unwrap_or(model_cfg.num_attention_heads) as u32; |
| 724 | let file_type_val: u32 = match config.dtype { |
| 725 | GgufDtype::F16 => 1, |
| 726 | GgufDtype::Q8_0 => 7, |
| 727 | }; |
| 728 | |
| 729 | // GGUF architecture string. |
| 730 | // |
| 731 | // Qwen3 has the same tensor layout as Qwen2 plus mandatory QK-norm |
| 732 | // (`attn_q_norm` / `attn_k_norm`). mistral.rs ships dedicated loaders for |
| 733 | // each — `quantized_qwen` for `qwen2`, `quantized_qwen3` for `qwen3` — so |
| 734 | // declare the architecture that matches the source `model_type`. Routing a |
| 735 | // Qwen3 model through the qwen2 loader loads but produces garbage output. |
| 736 | let arch = if model_cfg.model_type.starts_with("qwen3") { |
| 737 | "qwen3" |
| 738 | } else { |
| 739 | "qwen2" |
| 740 | }; |
| 741 | |
| 742 | // Count metadata KV entries |
| 743 | let num_metadata: u64 = 20; |
| 744 | |
| 745 | // Pre-compute header + metadata + tensor info byte size |
| 746 | let mut header_size: u64 = 0; |
| 747 | // Magic + version + tensor_count + metadata_kv_count |
| 748 | header_size += 4 + 4 + 8 + 8; |
| 749 | |
| 750 | // Metadata sizes |
| 751 | header_size += kv_string_size("general.architecture", arch); |
| 752 | header_size += kv_string_size("general.name", "onde-finetuned"); |
| 753 | header_size += kv_u32_size("general.file_type"); |
| 754 | header_size += kv_u32_size(&format!("{arch}.block_count")); |
| 755 | header_size += kv_u32_size(&format!("{arch}.embedding_length")); |
| 756 | header_size += kv_u32_size(&format!("{arch}.feed_forward_length")); |
| 757 | header_size += kv_u32_size(&format!("{arch}.attention.head_count")); |
| 758 | header_size += kv_u32_size(&format!("{arch}.attention.head_count_kv")); |
| 759 | header_size += kv_u32_size(&format!("{arch}.attention.key_length")); |
| 760 | header_size += kv_u32_size(&format!("{arch}.attention.value_length")); |
| 761 | header_size += kv_f32_size(&format!("{arch}.attention.layer_norm_rms_epsilon")); |
| 762 | header_size += kv_f32_size(&format!("{arch}.rope.freq_base")); |
| 763 | header_size += kv_u32_size(&format!("{arch}.context_length")); |
| 764 | header_size += kv_string_size("tokenizer.ggml.model", "gpt2"); |
| 765 | header_size += kv_string_array_size("tokenizer.ggml.tokens", &tokens); |
| 766 | header_size += kv_i32_array_size("tokenizer.ggml.token_type", &token_types); |
| 767 | header_size += kv_string_array_size("tokenizer.ggml.merges", &merges); |
| 768 | header_size += kv_u32_size("tokenizer.ggml.bos_token_id"); |
| 769 | header_size += kv_u32_size("tokenizer.ggml.eos_token_id"); |
| 770 | header_size += kv_string_size("tokenizer.chat_template", &chat_template); |
| 771 | |
| 772 | // Tensor info sizes |
| 773 | for entry in &entries { |
| 774 | header_size += tensor_info_size(&entry.gguf_name, entry.ndim); |
| 775 | } |
| 776 | |
| 777 | // Pad header to alignment |
| 778 | let header_padded = { |
| 779 | let rem = header_size % ALIGNMENT; |
| 780 | if rem != 0 { |
| 781 | header_size + (ALIGNMENT - rem) |
| 782 | } else { |
| 783 | header_size |
| 784 | } |
| 785 | }; |
| 786 | |
| 787 | // ----------------------------------------------------------------------- |
| 788 | // 6. Write the GGUF file |
| 789 | // ----------------------------------------------------------------------- |
| 790 | let file = std::fs::File::create(&config.output_path) |
| 791 | .with_context(|| format!("failed to create {}", config.output_path.display()))?; |
| 792 | let mut w = BufWriter::new(file); |
| 793 | |
| 794 | // Header |
| 795 | write_u32(&mut w, GGUF_MAGIC)?; |
| 796 | write_u32(&mut w, GGUF_VERSION)?; |
| 797 | write_u64(&mut w, num_tensors)?; |
| 798 | write_u64(&mut w, num_metadata)?; |
| 799 | |
| 800 | // Metadata KV pairs |
| 801 | write_kv_string(&mut w, "general.architecture", arch)?; |
| 802 | write_kv_string(&mut w, "general.name", "onde-finetuned")?; |
| 803 | write_kv_u32(&mut w, "general.file_type", file_type_val)?; |
| 804 | write_kv_u32( |
| 805 | &mut w, |
| 806 | &format!("{arch}.block_count"), |
| 807 | model_cfg.num_hidden_layers as u32, |
| 808 | )?; |
| 809 | write_kv_u32( |
| 810 | &mut w, |
| 811 | &format!("{arch}.embedding_length"), |
| 812 | model_cfg.hidden_size as u32, |
| 813 | )?; |
| 814 | write_kv_u32( |
| 815 | &mut w, |
| 816 | &format!("{arch}.feed_forward_length"), |
| 817 | model_cfg.intermediate_size as u32, |
| 818 | )?; |
| 819 | write_kv_u32( |
| 820 | &mut w, |
| 821 | &format!("{arch}.attention.head_count"), |
| 822 | model_cfg.num_attention_heads as u32, |
| 823 | )?; |
| 824 | write_kv_u32( |
| 825 | &mut w, |
| 826 | &format!("{arch}.attention.head_count_kv"), |
| 827 | num_kv_heads, |
| 828 | )?; |
| 829 | // Per-head dimension. mistral.rs / llama.cpp set head_dim = key_length and |
| 830 | // fall back to embedding_length / head_count when these keys are absent — |
| 831 | // which is wrong for Qwen3 (head_dim is decoupled), so always emit them. |
| 832 | write_kv_u32( |
| 833 | &mut w, |
| 834 | &format!("{arch}.attention.key_length"), |
| 835 | head_dim as u32, |
| 836 | )?; |
| 837 | write_kv_u32( |
| 838 | &mut w, |
| 839 | &format!("{arch}.attention.value_length"), |
| 840 | head_dim as u32, |
| 841 | )?; |
| 842 | write_kv_f32( |
| 843 | &mut w, |
| 844 | &format!("{arch}.attention.layer_norm_rms_epsilon"), |
| 845 | model_cfg.rms_norm_eps as f32, |
| 846 | )?; |
| 847 | write_kv_f32( |
| 848 | &mut w, |
| 849 | &format!("{arch}.rope.freq_base"), |
| 850 | model_cfg.rope_theta as f32, |
| 851 | )?; |
| 852 | write_kv_u32( |
| 853 | &mut w, |
| 854 | &format!("{arch}.context_length"), |
| 855 | model_cfg.max_position_embeddings as u32, |
| 856 | )?; |
| 857 | write_kv_string(&mut w, "tokenizer.ggml.model", "gpt2")?; |
| 858 | write_kv_string_array(&mut w, "tokenizer.ggml.tokens", &tokens)?; |
| 859 | write_kv_i32_array(&mut w, "tokenizer.ggml.token_type", &token_types)?; |
| 860 | write_kv_string_array(&mut w, "tokenizer.ggml.merges", &merges)?; |
| 861 | write_kv_u32(&mut w, "tokenizer.ggml.bos_token_id", bos_id)?; |
| 862 | write_kv_u32(&mut w, "tokenizer.ggml.eos_token_id", eos_id)?; |
| 863 | write_kv_string(&mut w, "tokenizer.chat_template", &chat_template)?; |
| 864 | |
| 865 | // Tensor info entries |
| 866 | for entry in &entries { |
| 867 | write_gguf_string(&mut w, &entry.gguf_name)?; |
| 868 | write_u32(&mut w, entry.ndim as u32)?; |
| 869 | for &d in &entry.dims { |
| 870 | write_u64(&mut w, d)?; |
| 871 | } |
| 872 | write_u32(&mut w, entry.gguf_type)?; |
| 873 | write_u64(&mut w, entry.offset)?; |
| 874 | } |
| 875 | |
| 876 | // Pad to alignment before data section |
| 877 | let mut pos = header_size; |
| 878 | pos = pad_to_alignment(&mut w, pos, ALIGNMENT)?; |
| 879 | let _ = header_padded; // should match pos |
| 880 | |
| 881 | // ----------------------------------------------------------------------- |
| 882 | // 7. Write tensor data |
| 883 | // ----------------------------------------------------------------------- |
| 884 | let total = entries.len(); |
| 885 | for (idx, entry) in entries.iter().enumerate() { |
| 886 | let _ = tx.send(GgufProgress::WritingTensor { |
| 887 | index: idx + 1, |
| 888 | total, |
| 889 | name: entry.gguf_name.clone(), |
| 890 | }); |
| 891 | |
| 892 | // Load tensor data. Synthetic entries (empty hf_name) get zeros. |
| 893 | let data: Vec<f32> = if entry.hf_name.is_empty() { |
| 894 | vec![0.0f32; entry.num_elements] |
| 895 | } else { |
| 896 | let tensor = safetensors |
| 897 | .load(&entry.hf_name, &Device::Cpu) |
| 898 | .with_context(|| format!("failed to load tensor {}", entry.hf_name))?; |
| 899 | let tensor_f32 = tensor |
| 900 | .to_dtype(candle_core::DType::F32) |
| 901 | .with_context(|| format!("failed to convert {} to f32", entry.hf_name))?; |
| 902 | let flat = tensor_f32 |
| 903 | .flatten_all() |
| 904 | .context("failed to flatten tensor")?; |
| 905 | flat.to_vec1::<f32>() |
| 906 | .context("failed to extract f32 data from tensor")? |
| 907 | }; |
| 908 | |
| 909 | match entry.gguf_type { |
| 910 | GGUF_TYPE_F32 => write_tensor_f32(&mut w, &data)?, |
| 911 | GGUF_TYPE_F16 => write_tensor_f16(&mut w, &data)?, |
| 912 | GGUF_TYPE_Q8_0 => write_tensor_q8_0(&mut w, &data)?, |
| 913 | _ => write_tensor_f32(&mut w, &data)?, |
| 914 | } |
| 915 | |
| 916 | pos += entry.data_size as u64; |
| 917 | |
| 918 | // Pad to alignment (except possibly the last tensor) |
| 919 | if idx + 1 < total { |
| 920 | pos = pad_to_alignment(&mut w, pos, ALIGNMENT)?; |
| 921 | } |
| 922 | } |
| 923 | |
| 924 | w.flush().context("failed to flush GGUF output")?; |
| 925 | drop(w); |
| 926 | |
| 927 | // ----------------------------------------------------------------------- |
| 928 | // 8. Report completion |
| 929 | // ----------------------------------------------------------------------- |
| 930 | let size_bytes = std::fs::metadata(&config.output_path) |
| 931 | .with_context(|| format!("failed to stat {}", config.output_path.display()))? |
| 932 | .len(); |
| 933 | |
| 934 | let _ = tx.send(GgufProgress::Done { |
| 935 | output_path: config.output_path.clone(), |
| 936 | size_bytes, |
| 937 | }); |
| 938 | |
| 939 | Ok(()) |
| 940 | } |
| 941 | |
| 942 | #[cfg(test)] |
| 943 | mod tests { |
| 944 | use super::*; |
| 945 | |
| 946 | /// Resolve the locally cached merged Qwen3-0.6B model directory, if present. |
| 947 | /// Returns `None` so the heavy integration test can be skipped on machines |
| 948 | /// that have not run a fine-tune. |
| 949 | fn merged_qwen3_dir() -> Option<PathBuf> { |
| 950 | let dir = dirs::home_dir()? |
| 951 | .join("Library/Group Containers/group.com.ondeinference.apps") |
| 952 | .join("models/hub/models--Qwen--Qwen3-0.6B/snapshots/merged-model"); |
| 953 | if dir.join("config.json").exists() && dir.join("model.safetensors").exists() { |
| 954 | Some(dir) |
| 955 | } else { |
| 956 | None |
| 957 | } |
| 958 | } |
| 959 | |
| 960 | /// Resolve the cached base Qwen3-0.6B snapshot (the hex commit dir), if present. |
| 961 | fn base_qwen3_dir() -> Option<PathBuf> { |
| 962 | let snapshots = dirs::home_dir()? |
| 963 | .join("Library/Group Containers/group.com.ondeinference.apps") |
| 964 | .join("models/hub/models--Qwen--Qwen3-0.6B/snapshots"); |
| 965 | let entry = std::fs::read_dir(&snapshots).ok()?.flatten().find(|e| { |
| 966 | let n = e.file_name(); |
| 967 | let n = n.to_string_lossy(); |
| 968 | n.len() >= 7 && n.chars().all(|c| c.is_ascii_hexdigit()) |
| 969 | })?; |
| 970 | let dir = entry.path(); |
| 971 | (dir.join("config.json").exists() && dir.join("model.safetensors").exists()).then_some(dir) |
| 972 | } |
| 973 | |
| 974 | /// Load a GGUF and generate a reply via the same `onde::mistralrs` engine the |
| 975 | /// Onde SDK uses. Returns the assistant text. Panics on model/inference error. |
| 976 | fn run_gguf_inference(gguf_path: &std::path::Path, user: &str) -> String { |
| 977 | use onde::mistralrs::{ |
| 978 | GgufModelBuilder, RequestBuilder, Response, TextMessageRole, TokenSource, |
| 979 | }; |
| 980 | let rt = tokio::runtime::Runtime::new().unwrap(); |
| 981 | let path = gguf_path.to_path_buf(); |
| 982 | let user = user.to_string(); |
| 983 | rt.block_on(async move { |
| 984 | let file_name = path.file_name().unwrap().to_string_lossy().to_string(); |
| 985 | let model_root = path.parent().unwrap().to_string_lossy().to_string(); |
| 986 | let model = GgufModelBuilder::new(&model_root, vec![file_name]) |
| 987 | .with_token_source(TokenSource::None) |
| 988 | .build() |
| 989 | .await |
| 990 | .expect("exported GGUF must load in mistral.rs"); |
| 991 | let mut req = RequestBuilder::new().set_sampler_max_len(512); |
| 992 | if let Ok(system) = std::env::var("ONDE_TEST_SYSTEM_PROMPT") { |
| 993 | req = req.add_message(TextMessageRole::System, &system); |
| 994 | } |
| 995 | let req = req.add_message(TextMessageRole::User, &user); |
| 996 | let mut stream = model |
| 997 | .stream_chat_request(req) |
| 998 | .await |
| 999 | .expect("inference request must start"); |
| 1000 | let mut out = String::new(); |
| 1001 | while let Some(resp) = stream.next().await { |
| 1002 | match resp { |
| 1003 | Response::Chunk(chunk) => { |
| 1004 | if let Some(choice) = chunk.choices.first() |
| 1005 | && let Some(delta) = &choice.delta.content |
| 1006 | { |
| 1007 | out.push_str(delta); |
| 1008 | } |
| 1009 | } |
| 1010 | Response::Done(_) => break, |
| 1011 | Response::ModelError(msg, _) => panic!("model error: {msg}"), |
| 1012 | Response::InternalError(e) => panic!("internal error: {e}"), |
| 1013 | Response::ValidationError(e) => panic!("validation error: {e}"), |
| 1014 | _ => {} |
| 1015 | } |
| 1016 | } |
| 1017 | out |
| 1018 | }) |
| 1019 | } |
| 1020 | |
| 1021 | /// Full v1 pipeline regression guard: fine-tune Qwen3-0.6B → merge the LoRA |
| 1022 | /// adapter → export GGUF → load and run in mistral.rs. Proves the whole |
| 1023 | /// fine-tune → quantize → run path produces a valid, runnable model. |
| 1024 | /// |
| 1025 | /// Ignored by default (trains + loads ~1 GB on Metal/CPU); run with |
| 1026 | /// `cargo test gguf::tests::finetune_merge_export_run -- --ignored --nocapture`. |
| 1027 | #[test] |
| 1028 | #[ignore = "heavy: fine-tunes, merges, exports ~1GB GGUF and runs it"] |
| 1029 | fn finetune_merge_export_run() { |
| 1030 | let Some(base_dir) = base_qwen3_dir() else { |
| 1031 | eprintln!("base Qwen3-0.6B not found locally; skipping"); |
| 1032 | return; |
| 1033 | }; |
| 1034 | let data_path = match std::env::var("ONDE_TEST_DATA_PATH") { |
| 1035 | Ok(p) => PathBuf::from(p), |
| 1036 | Err(_) => dirs::home_dir() |
| 1037 | .unwrap() |
| 1038 | .join(".onde/datasets/Qwen/Qwen3-0.6B/train.jsonl"), |
| 1039 | }; |
| 1040 | if !data_path.exists() { |
| 1041 | eprintln!("training data {} not found; skipping", data_path.display()); |
| 1042 | return; |
| 1043 | } |
| 1044 | |
| 1045 | let work = std::env::temp_dir().join("onde-e2e-test"); |
| 1046 | let _ = std::fs::remove_dir_all(&work); |
| 1047 | std::fs::create_dir_all(&work).unwrap(); |
| 1048 | |
| 1049 | // 1. Fine-tune. |
| 1050 | let lora_dir = work.join("lora"); |
| 1051 | let (ft_tx, mut ft_rx) = tokio::sync::mpsc::unbounded_channel(); |
| 1052 | crate::finetune::start_finetune( |
| 1053 | crate::finetune::FineTuneConfig { |
| 1054 | model_dir: base_dir.clone(), |
| 1055 | data_path, |
| 1056 | output_dir: lora_dir.clone(), |
| 1057 | lora_rank: 8, |
| 1058 | lora_alpha: 16.0, |
| 1059 | learning_rate: std::env::var("ONDE_TEST_LR") |
| 1060 | .ok() |
| 1061 | .and_then(|v| v.parse().ok()) |
| 1062 | .unwrap_or(1e-4), |
| 1063 | epochs: std::env::var("ONDE_TEST_EPOCHS") |
| 1064 | .ok() |
| 1065 | .and_then(|v| v.parse().ok()) |
| 1066 | .unwrap_or(2), |
| 1067 | max_seq_len: std::env::var("ONDE_TEST_MAX_SEQ_LEN") |
| 1068 | .ok() |
| 1069 | .and_then(|v| v.parse().ok()) |
| 1070 | .unwrap_or(512), |
| 1071 | }, |
| 1072 | ft_tx, |
| 1073 | ); |
| 1074 | let mut adapter_path = None; |
| 1075 | let mut last_loss = f32::NAN; |
| 1076 | let rt = tokio::runtime::Runtime::new().unwrap(); |
| 1077 | rt.block_on(async { |
| 1078 | while let Some(p) = ft_rx.recv().await { |
| 1079 | match p { |
| 1080 | crate::finetune::FineTuneProgress::Training { |
| 1081 | loss, |
| 1082 | step, |
| 1083 | total_steps, |
| 1084 | epoch, |
| 1085 | .. |
| 1086 | } => { |
| 1087 | eprintln!("[e2e] epoch {epoch} step {step}/{total_steps} loss {loss:.4}"); |
| 1088 | last_loss = loss; |
| 1089 | } |
| 1090 | crate::finetune::FineTuneProgress::Done { adapter_path: a } => { |
| 1091 | adapter_path = Some(a); |
| 1092 | break; |
| 1093 | } |
| 1094 | crate::finetune::FineTuneProgress::Failed(e) => panic!("fine-tune failed: {e}"), |
| 1095 | _ => {} |
| 1096 | } |
| 1097 | } |
| 1098 | }); |
| 1099 | let adapter_path = adapter_path.expect("fine-tune must produce an adapter"); |
| 1100 | eprintln!("[e2e] fine-tune done, last loss {last_loss}"); |
| 1101 | assert!( |
| 1102 | last_loss.is_finite(), |
| 1103 | "training loss is NaN/Inf — training diverged" |
| 1104 | ); |
| 1105 | |
| 1106 | // 2. Merge. |
| 1107 | let merged_dir = work.join("merged"); |
| 1108 | let (mg_tx, mut mg_rx) = tokio::sync::mpsc::unbounded_channel(); |
| 1109 | crate::merge::start_merge( |
| 1110 | crate::merge::MergeConfig { |
| 1111 | base_dir, |
| 1112 | adapter_path, |
| 1113 | output_dir: merged_dir.clone(), |
| 1114 | }, |
| 1115 | mg_tx, |
| 1116 | ); |
| 1117 | let mut merged_ok = false; |
| 1118 | rt.block_on(async { |
| 1119 | while let Some(p) = mg_rx.recv().await { |
| 1120 | match p { |
| 1121 | crate::merge::MergeProgress::Done { .. } => { |
| 1122 | merged_ok = true; |
| 1123 | break; |
| 1124 | } |
| 1125 | crate::merge::MergeProgress::Failed(e) => panic!("merge failed: {e}"), |
| 1126 | _ => {} |
| 1127 | } |
| 1128 | } |
| 1129 | }); |
| 1130 | assert!(merged_ok, "merge did not complete"); |
| 1131 | |
| 1132 | // 3. Export GGUF (the dtype the CLI ships). |
| 1133 | let gguf_path = work.join("finetuned-q8_0.gguf"); |
| 1134 | let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); |
| 1135 | run_gguf_export( |
| 1136 | &GgufConfig { |
| 1137 | model_dir: merged_dir, |
| 1138 | output_path: gguf_path.clone(), |
| 1139 | dtype: GgufDtype::Q8_0, |
| 1140 | }, |
| 1141 | &tx, |
| 1142 | ) |
| 1143 | .expect("GGUF export should succeed"); |
| 1144 | |
| 1145 | // 4. Load + run in mistral.rs. Query in the training distribution so a |
| 1146 | // lightly fine-tuned model has something to say. `ONDE_TEST_EVAL_PROMPT` |
| 1147 | // overrides the default (multiple prompts separated by `||`). |
| 1148 | let eval_prompts = std::env::var("ONDE_TEST_EVAL_PROMPT") |
| 1149 | .unwrap_or_else(|_| "Apa bahasa Jawa dari \"terima kasih\"?".to_string()); |
| 1150 | for prompt in eval_prompts.split("||") { |
| 1151 | let reply = run_gguf_inference(&gguf_path, prompt.trim()); |
| 1152 | eprintln!("[e2e] eval prompt: {prompt:?}"); |
| 1153 | eprintln!("[e2e] fine-tuned GGUF reply: {reply:?}"); |
| 1154 | assert!( |
| 1155 | !reply.trim().is_empty(), |
| 1156 | "fine-tuned GGUF produced an empty reply — not runnable" |
| 1157 | ); |
| 1158 | } |
| 1159 | } |
| 1160 | |
| 1161 | /// End-to-end check: export a Qwen3 safetensors model to GGUF, then load and |
| 1162 | /// run it through the same `onde::mistralrs` engine the Onde SDK uses. |
| 1163 | /// |
| 1164 | /// This is the regression guard for the Qwen3 `head_dim` fix: with the wrong |
| 1165 | /// per-head dimension the GGUF loads with a reshape mismatch and inference |
| 1166 | /// fails. Ignored by default (downloads/loads ~1 GB, needs Metal/CPU); |
| 1167 | /// run with `cargo test gguf -- --ignored --nocapture`. |
| 1168 | #[test] |
| 1169 | #[ignore = "heavy: exports ~1GB GGUF and loads it in mistral.rs"] |
| 1170 | fn exported_qwen3_gguf_is_runnable() { |
| 1171 | // Allow overriding the source model dir and dtype to compare |
| 1172 | // base-vs-merged and F16-vs-Q8_0 without changing code. |
| 1173 | let model_dir = match std::env::var("ONDE_TEST_MODEL_DIR") { |
| 1174 | Ok(d) => PathBuf::from(d), |
| 1175 | // Default to the known-good base snapshot so this isolates the |
| 1176 | // exporter; `finetune_merge_export_run` covers the fine-tuned path. |
| 1177 | Err(_) => match base_qwen3_dir().or_else(merged_qwen3_dir) { |
| 1178 | Some(d) => d, |
| 1179 | None => { |
| 1180 | eprintln!("Qwen3-0.6B model not found locally; skipping"); |
| 1181 | return; |
| 1182 | } |
| 1183 | }, |
| 1184 | }; |
| 1185 | let dtype = match std::env::var("ONDE_TEST_DTYPE").as_deref() { |
| 1186 | Ok("f16") => GgufDtype::F16, |
| 1187 | _ => GgufDtype::Q8_0, |
| 1188 | }; |
| 1189 | eprintln!( |
| 1190 | "[test] model_dir={} dtype={}", |
| 1191 | model_dir.display(), |
| 1192 | if matches!(dtype, GgufDtype::F16) { |
| 1193 | "f16" |
| 1194 | } else { |
| 1195 | "q8_0" |
| 1196 | } |
| 1197 | ); |
| 1198 | |
| 1199 | let out_dir = std::env::temp_dir().join("onde-gguf-test"); |
| 1200 | std::fs::create_dir_all(&out_dir).unwrap(); |
| 1201 | let output_path = out_dir.join("qwen3-0.6b-test.gguf"); |
| 1202 | let _ = std::fs::remove_file(&output_path); |
| 1203 | |
| 1204 | // 1. Export GGUF. |
| 1205 | let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); |
| 1206 | run_gguf_export( |
| 1207 | &GgufConfig { |
| 1208 | model_dir, |
| 1209 | output_path: output_path.clone(), |
| 1210 | dtype, |
| 1211 | }, |
| 1212 | &tx, |
| 1213 | ) |
| 1214 | .expect("GGUF export should succeed"); |
| 1215 | assert!(output_path.exists(), "GGUF file was not written"); |
| 1216 | |
| 1217 | // 2. Load + run it through mistral.rs (the Onde SDK's engine). |
| 1218 | let reply = run_gguf_inference(&output_path, "Say hello in one short sentence."); |
| 1219 | eprintln!("exported GGUF reply: {reply:?}"); |
| 1220 | assert!( |
| 1221 | !reply.trim().is_empty(), |
| 1222 | "exported GGUF produced an empty reply — model is not runnable" |
| 1223 | ); |
| 1224 | } |
| 1225 | } |