pictor_model/convert/
common.rs

1//! Shared helpers for HuggingFace / ONNX → GGUF conversion pipelines.
2//!
3//! Both the safetensors path (`convert::convert_hf_to_gguf`) and the ONNX path
4//! (`convert::onnx::convert_onnx_to_gguf`) emit the same GGUF layout and share:
5//!
6//! * Parsing of the sibling `config.json`.
7//! * Writing Qwen3 metadata (architecture, dimensions, norm epsilon, rope base).
8//! * Padding `f32` weights to a multiple of the TQ2_0_g128 block size.
9//! * Serialising `BlockTQ2_0_g128` blocks into raw GGUF tensor bytes.
10//! * A single `ConvertStats` result struct so callers can report progress
11//!   uniformly.
12
13use std::path::Path;
14
15use anyhow::Context;
16use serde_json::Value;
17
18use pictor_core::gguf::tensor_info::keys;
19use pictor_core::gguf::writer::{GgufWriter, MetadataWriteValue};
20use pictor_core::quant_ternary::{BlockTQ2_0_g128, BLOCK_TQ2_0_G128_BYTES};
21
22/// Statistics returned after a successful conversion.
23///
24/// Shared between the safetensors (`convert_hf_to_gguf`) and ONNX
25/// (`convert_onnx_to_gguf`) pipelines so CLI callers can report a single
26/// struct regardless of source format.
27#[derive(Debug, Clone, Default)]
28pub struct ConvertStats {
29    /// Total number of tensors written to the GGUF file.
30    pub n_tensors: usize,
31    /// Number of tensors quantized to TQ2_0_g128.
32    pub n_ternary: usize,
33    /// Number of tensors stored as FP32.
34    pub n_fp32: usize,
35    /// Number of tensors stored verbatim as BF16 (e.g. MLX skip-pattern tensors).
36    pub n_bf16: usize,
37    /// Number of tensors stored verbatim as F16.
38    pub n_f16: usize,
39    /// Total size of the output GGUF file in bytes.
40    pub output_bytes: usize,
41}
42
43/// Read and parse `config.json` at the given path.
44///
45/// Callers are responsible for locating the file (the safetensors path uses
46/// `from_dir.join("config.json")`; the ONNX path may need to search the ONNX
47/// parent and grandparent directories).
48pub fn read_config_json(config_path: &Path) -> anyhow::Result<Value> {
49    let raw = std::fs::read_to_string(config_path)
50        .with_context(|| format!("reading {:?}", config_path))?;
51    let value: Value =
52        serde_json::from_str(&raw).with_context(|| format!("parsing {:?}", config_path))?;
53    Ok(value)
54}
55
56/// Write Qwen3 metadata from `config.json` into a GGUF writer.
57///
58/// The caller provides the human-readable model name; for HF this is usually
59/// the directory basename, and for ONNX the `.onnx` file stem or repository
60/// identifier.
61pub fn write_metadata(
62    writer: &mut GgufWriter,
63    config: &Value,
64    model_name: &str,
65) -> anyhow::Result<()> {
66    // Architecture constant
67    writer.add_metadata(
68        keys::GENERAL_ARCHITECTURE,
69        MetadataWriteValue::Str("qwen3".to_string()),
70    );
71
72    // Human-readable model name
73    writer.add_metadata(
74        keys::GENERAL_NAME,
75        MetadataWriteValue::Str(model_name.to_string()),
76    );
77
78    // Quantisation version string
79    writer.add_metadata(
80        "general.quantization_version",
81        MetadataWriteValue::Str("TQ2_0_G128".to_string()),
82    );
83
84    // Integer keys (u32) - required.
85    let u32_keys = [
86        (keys::LLM_BLOCK_COUNT, "num_hidden_layers"),
87        (keys::LLM_EMBEDDING_LENGTH, "hidden_size"),
88        (keys::LLM_FEED_FORWARD_LENGTH, "intermediate_size"),
89        (keys::LLM_ATTENTION_HEAD_COUNT, "num_attention_heads"),
90        (keys::LLM_ATTENTION_HEAD_COUNT_KV, "num_key_value_heads"),
91        (keys::LLM_CONTEXT_LENGTH, "max_position_embeddings"),
92        (keys::LLM_VOCAB_SIZE, "vocab_size"),
93    ];
94    for (gguf_key, json_key) in &u32_keys {
95        if let Some(val) = config.get(*json_key).and_then(Value::as_u64) {
96            writer.add_metadata(gguf_key, MetadataWriteValue::U32(val as u32));
97        } else {
98            tracing::warn!(json_key, "missing or non-u64 field in config.json");
99        }
100    }
101
102    // head_dim is optional in config.json: Qwen3 (1.7B/4B/8B/14B) sets it
103    // explicitly because head_dim is decoupled from hidden_size/num_heads
104    // (notably Qwen3-4B has hidden=2560, heads=32, head_dim=128, so the
105    // derived hidden/heads=80 is wrong). Older Qwen2 configs omit it and
106    // the reader falls back to the hidden/heads derivation.
107    if let Some(val) = config.get("head_dim").and_then(Value::as_u64) {
108        writer.add_metadata(
109            keys::LLM_ATTENTION_KEY_LENGTH,
110            MetadataWriteValue::U32(val as u32),
111        );
112    }
113
114    // rms_norm_eps → F32
115    if let Some(eps) = config.get("rms_norm_eps").and_then(Value::as_f64) {
116        writer.add_metadata(
117            keys::LLM_ATTENTION_LAYER_NORM_RMS_EPSILON,
118            MetadataWriteValue::F32(eps as f32),
119        );
120    }
121
122    // rope_theta → F32 (default 10000.0 if absent)
123    //
124    // Resolution order:
125    //   1. `config["rope_theta"]`                    (top-level, legacy Qwen2 layout)
126    //   2. `config["rope_parameters"]["rope_theta"]` (nested, Qwen3 ONNX/newer layout)
127    //   3. 10000.0 fallback (with `tracing::warn!`)
128    let rope_theta = resolve_rope_theta(config);
129    writer.add_metadata(
130        keys::LLM_ROPE_FREQ_BASE,
131        MetadataWriteValue::F32(rope_theta as f32),
132    );
133
134    // If the nested `rope_parameters` block indicates YARN scaling, note it.
135    //
136    // The existing native `Ternary-Bonsai-1.7B.gguf` only carries
137    // `llm.rope.freq_base` — no `llm.rope.scaling.*` keys — so for now we log
138    // the YARN parameters at info level and rely on architecture defaults
139    // rather than inventing new metadata keys.
140    if let Some(rp) = config.get("rope_parameters").and_then(Value::as_object) {
141        let rope_type = rp.get("rope_type").and_then(Value::as_str).unwrap_or("");
142        if rope_type.eq_ignore_ascii_case("yarn") {
143            let factor = rp.get("factor").and_then(Value::as_f64);
144            let original_max_pos = rp
145                .get("original_max_position_embeddings")
146                .and_then(Value::as_u64);
147            tracing::info!(
148                ?factor,
149                ?original_max_pos,
150                "YARN rope_parameters detected; GGUF YARN metadata not plumbed, relying on architecture defaults"
151            );
152        }
153    }
154
155    Ok(())
156}
157
158/// Resolve `rope_theta` from a HuggingFace `config.json` value.
159///
160/// Looks in this order:
161///   1. Top-level `rope_theta` (legacy Qwen2 layout).
162///   2. Nested `rope_parameters.rope_theta` (Qwen3 ONNX/newer layout).
163///   3. Fallback `10000.0` with a `tracing::warn!` describing the absence.
164fn resolve_rope_theta(config: &Value) -> f64 {
165    if let Some(v) = config.get("rope_theta").and_then(Value::as_f64) {
166        return v;
167    }
168    if let Some(v) = config
169        .get("rope_parameters")
170        .and_then(|rp| rp.get("rope_theta"))
171        .and_then(Value::as_f64)
172    {
173        return v;
174    }
175    tracing::warn!(
176        "config.json missing both `rope_theta` and `rope_parameters.rope_theta`; \
177         falling back to default 10000.0"
178    );
179    10000.0
180}
181
182/// Pad an f32 slice to a multiple of 128 elements for TQ2_0_g128 quantisation.
183///
184/// If the length is already block-aligned, the slice is copied verbatim.
185/// Otherwise the tail is zero-padded up to the next multiple of 128.
186pub fn pad_to_multiple_of_128(f32_data: &[f32]) -> Vec<f32> {
187    let len = f32_data.len();
188    let remainder = len % 128;
189    if remainder == 0 {
190        f32_data.to_vec()
191    } else {
192        let padded_len = len + (128 - remainder);
193        let mut padded = f32_data.to_vec();
194        padded.resize(padded_len, 0.0_f32);
195        padded
196    }
197}
198
199/// Serialise a slice of `BlockTQ2_0_g128` blocks into raw bytes.
200///
201/// Each block is 34 bytes: 32 bytes of packed `qs` + 2 bytes of FP16 `d`.
202///
203/// # Safety
204///
205/// `BlockTQ2_0_g128` is `#[repr(C)]` with a compile-time size assertion of
206/// exactly 34 bytes. The cast is safe because we size the output slice using
207/// `blocks.len() * BLOCK_TQ2_0_G128_BYTES`.
208pub fn blocks_to_bytes(blocks: &[BlockTQ2_0_g128]) -> Vec<u8> {
209    let total = blocks.len() * BLOCK_TQ2_0_G128_BYTES;
210    // SAFETY: repr(C) layout with compile-time size check; byte length verified.
211    let bytes: &[u8] = unsafe { std::slice::from_raw_parts(blocks.as_ptr() as *const u8, total) };
212    bytes.to_vec()
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218    use serde_json::json;
219
220    #[test]
221    fn pad_aligned_is_identity() {
222        let v = vec![1.0_f32; 128];
223        assert_eq!(pad_to_multiple_of_128(&v), v);
224    }
225
226    #[test]
227    fn pad_extends_to_next_block() {
228        let v = vec![1.0_f32; 130];
229        let padded = pad_to_multiple_of_128(&v);
230        assert_eq!(padded.len(), 256);
231        assert_eq!(&padded[..130], &v[..]);
232        assert!(padded[130..].iter().all(|&x| x == 0.0));
233    }
234
235    #[test]
236    fn empty_input_stays_empty() {
237        let v: Vec<f32> = Vec::new();
238        assert!(pad_to_multiple_of_128(&v).is_empty());
239    }
240
241    #[test]
242    fn rope_theta_top_level_wins() {
243        // Legacy Qwen2 layout: `rope_theta` at the top level.
244        let cfg = json!({
245            "rope_theta": 500000.0,
246        });
247        assert_eq!(resolve_rope_theta(&cfg), 500000.0);
248    }
249
250    #[test]
251    fn rope_theta_nested_under_rope_parameters() {
252        // Qwen3 ONNX layout: nested under `rope_parameters`.
253        let cfg = json!({
254            "rope_parameters": {
255                "factor": 4.0,
256                "original_max_position_embeddings": 8192,
257                "rope_theta": 1_000_000.0,
258                "rope_type": "yarn",
259            },
260        });
261        assert_eq!(resolve_rope_theta(&cfg), 1_000_000.0);
262    }
263
264    #[test]
265    fn rope_theta_top_level_takes_precedence_over_nested() {
266        // If both are present, the top-level value wins.
267        let cfg = json!({
268            "rope_theta": 250000.0,
269            "rope_parameters": {
270                "rope_theta": 1_000_000.0,
271                "rope_type": "yarn",
272            },
273        });
274        assert_eq!(resolve_rope_theta(&cfg), 250000.0);
275    }
276
277    #[test]
278    fn rope_theta_fallback_when_missing() {
279        // Neither key present: fall back to the default 10000.0.
280        let cfg = json!({
281            "hidden_size": 2048,
282        });
283        assert_eq!(resolve_rope_theta(&cfg), 10000.0);
284    }
285}
pictor_model/convert/common.rs

pictor_model/convert/
common.rs
