apr_format/v2/
header_impl.rs

1//! v2 header impl + JSON metadata + tensor-index entry types (issue #2231).
2//!
3//! Formerly `include!`d into `v2/mod.rs`; now a real module. Reaches the
4//! parent-scope header/flag/const definitions via `super::` and the sibling
5//! [`super::TensorDType`] via the `v2` namespace re-export.
6
7use super::{
8    AprV2Flags, AprV2Header, TensorDType, V2FormatError, HEADER_SIZE_V2, MAGIC_V2, VERSION_V2,
9};
10use crate::crc32::crc32;
11use serde::{Deserialize, Serialize};
12use std::collections::HashMap;
13
14impl AprV2Header {
15    /// Create new v2 header with defaults
16    #[must_use]
17    pub fn new() -> Self {
18        Self {
19            magic: MAGIC_V2,
20            version: VERSION_V2,
21            flags: AprV2Flags::new(),
22            tensor_count: 0,
23            metadata_offset: HEADER_SIZE_V2 as u64,
24            metadata_size: 0,
25            tensor_index_offset: 0,
26            data_offset: 0,
27            checksum: 0,
28            reserved: [0u8; 20],
29        }
30    }
31
32    /// Check if header has valid magic number
33    #[must_use]
34    pub fn is_valid(&self) -> bool {
35        self.magic == MAGIC_V2
36    }
37
38    /// Serialize header to bytes
39    #[must_use]
40    pub fn to_bytes(&self) -> [u8; HEADER_SIZE_V2] {
41        let mut buf = [0u8; HEADER_SIZE_V2];
42
43        buf[0..4].copy_from_slice(&self.magic);
44        buf[4] = self.version.0;
45        buf[5] = self.version.1;
46        buf[6..8].copy_from_slice(&self.flags.bits().to_le_bytes());
47        buf[8..12].copy_from_slice(&self.tensor_count.to_le_bytes());
48        buf[12..20].copy_from_slice(&self.metadata_offset.to_le_bytes());
49        buf[20..24].copy_from_slice(&self.metadata_size.to_le_bytes());
50        buf[24..32].copy_from_slice(&self.tensor_index_offset.to_le_bytes());
51        buf[32..40].copy_from_slice(&self.data_offset.to_le_bytes());
52        buf[40..44].copy_from_slice(&self.checksum.to_le_bytes());
53        buf[44..64].copy_from_slice(&self.reserved);
54
55        buf
56    }
57
58    /// Deserialize header from bytes
59    ///
60    /// # Errors
61    /// Returns error if buffer is too small or magic is invalid.
62    pub fn from_bytes(buf: &[u8]) -> Result<Self, V2FormatError> {
63        if buf.len() < HEADER_SIZE_V2 {
64            return Err(V2FormatError::InvalidHeader("buffer too small".to_string()));
65        }
66
67        let magic: [u8; 4] = buf[0..4]
68            .try_into()
69            .map_err(|_| V2FormatError::InvalidHeader("failed to read magic".to_string()))?;
70
71        // Check for v2 magic only
72        if magic != MAGIC_V2 {
73            return Err(V2FormatError::InvalidMagic(magic));
74        }
75
76        let version = (buf[4], buf[5]);
77        let flags = AprV2Flags::from_bits(u16::from_le_bytes([buf[6], buf[7]]));
78        let tensor_count = u32::from_le_bytes([buf[8], buf[9], buf[10], buf[11]]);
79        let metadata_offset = u64::from_le_bytes(buf[12..20].try_into().unwrap_or([0; 8]));
80        let metadata_size = u32::from_le_bytes([buf[20], buf[21], buf[22], buf[23]]);
81        let tensor_index_offset = u64::from_le_bytes(buf[24..32].try_into().unwrap_or([0; 8]));
82        let data_offset = u64::from_le_bytes(buf[32..40].try_into().unwrap_or([0; 8]));
83        let checksum = u32::from_le_bytes([buf[40], buf[41], buf[42], buf[43]]);
84
85        let mut reserved = [0u8; 20];
86        reserved.copy_from_slice(buf.get(44..64).unwrap_or(&[0u8; 20]));
87
88        Ok(Self {
89            magic,
90            version,
91            flags,
92            tensor_count,
93            metadata_offset,
94            metadata_size,
95            tensor_index_offset,
96            data_offset,
97            checksum,
98            reserved,
99        })
100    }
101
102    /// Compute header checksum (CRC32 of header bytes excluding checksum field)
103    #[must_use]
104    pub fn compute_checksum(&self) -> u32 {
105        let bytes = self.to_bytes();
106        // Exclude checksum field (bytes 40-43) from calculation
107        // Concatenate the two regions and compute CRC32
108        let mut data = Vec::with_capacity(60);
109        data.extend_from_slice(bytes.get(0..40).unwrap_or(&[]));
110        data.extend_from_slice(bytes.get(44..64).unwrap_or(&[]));
111        crc32(&data)
112    }
113
114    /// Update checksum field
115    pub fn update_checksum(&mut self) {
116        self.checksum = self.compute_checksum();
117    }
118
119    /// Verify header checksum
120    #[must_use]
121    pub fn verify_checksum(&self) -> bool {
122        self.checksum == self.compute_checksum()
123    }
124}
125
126// ============================================================================
127// Metadata
128// ============================================================================
129
130/// APR v2 JSON metadata section
131#[derive(Debug, Clone, Default, Serialize, Deserialize)]
132pub struct AprV2Metadata {
133    /// Model type identifier
134    #[serde(default)]
135    pub model_type: String,
136
137    /// Model name
138    #[serde(default, skip_serializing_if = "Option::is_none")]
139    pub name: Option<String>,
140
141    /// Model description
142    #[serde(default, skip_serializing_if = "Option::is_none")]
143    pub description: Option<String>,
144
145    /// Model author/organization
146    #[serde(default, skip_serializing_if = "Option::is_none")]
147    pub author: Option<String>,
148
149    /// Model license (SPDX identifier; governed by C-APR-PROVENANCE).
150    /// NO skip_serializing_if here: FALSIFY-SHIP-022 requires provenance
151    /// keys to serialize as explicit `null` (FM-APR-PROV-SILENT-SKIP) —
152    /// and `license` is not a realizar alias-group member, so the null
153    /// cannot trigger the duplicate-field poison (C-APR-MERGE-RUNNABLE).
154    #[serde(default)]
155    pub license: Option<String>,
156
157    /// Training-data source (dataset identifier or "teacher-only";
158    /// governed by C-APR-PROVENANCE / AC-SHIP2-012 / FALSIFY-SHIP-022).
159    /// Explicit-null serialization required — see `license`.
160    #[serde(default)]
161    pub data_source: Option<String>,
162
163    /// SPDX license for `data_source` (governed by C-APR-PROVENANCE /
164    /// AC-SHIP2-012 / FALSIFY-SHIP-022).
165    /// Explicit-null serialization required — see `license`.
166    #[serde(default)]
167    pub data_license: Option<String>,
168
169    /// Model version string
170    #[serde(default, skip_serializing_if = "Option::is_none")]
171    pub version: Option<String>,
172
173    /// Source/provenance URI (DD6: Model provenance tracking)
174    /// Examples: "<hf://openai/whisper-tiny>", "<local://path/to/model.safetensors>"
175    #[serde(default, skip_serializing_if = "Option::is_none")]
176    pub source: Option<String>,
177
178    /// Original format before conversion
179    /// Examples: "safetensors", "gguf", "pytorch"
180    #[serde(default, skip_serializing_if = "Option::is_none")]
181    pub original_format: Option<String>,
182
183    /// Creation timestamp (ISO 8601)
184    #[serde(default, skip_serializing_if = "Option::is_none")]
185    pub created_at: Option<String>,
186
187    /// Total model size in bytes
188    #[serde(default)]
189    pub total_size: u64,
190
191    /// Parameter count
192    #[serde(default)]
193    pub param_count: u64,
194
195    /// Quantization info
196    #[serde(default, skip_serializing_if = "Option::is_none")]
197    pub quantization: Option<QuantizationMetadata>,
198
199    /// Shard info (for multi-file models)
200    #[serde(default, skip_serializing_if = "Option::is_none")]
201    pub sharding: Option<ShardingMetadata>,
202
203    /// Chat template (Jinja2 format, from tokenizer_config.json)
204    /// Per spec: chat-template-improvement-spec.md CTA-01
205    #[serde(default, skip_serializing_if = "Option::is_none")]
206    pub chat_template: Option<String>,
207
208    /// Detected chat template format
209    /// Per spec: chat-template-improvement-spec.md CTA-03
210    /// Values: "chatml", "llama2", "mistral", "phi", "alpaca", "custom", "raw"
211    #[serde(default, skip_serializing_if = "Option::is_none")]
212    pub chat_format: Option<String>,
213
214    /// Special tokens for chat templates
215    /// Per spec: chat-template-improvement-spec.md CTA-04
216    #[serde(default, skip_serializing_if = "Option::is_none")]
217    pub special_tokens: Option<ChatSpecialTokens>,
218
219    // ========================================================================
220    // Transformer Config (CRITICAL for inference - realizar::apr::AprMetadata)
221    // ========================================================================
222    /// Model architecture family (e.g., "llama", "qwen2", "phi")
223    #[serde(default, skip_serializing_if = "Option::is_none")]
224    pub architecture: Option<String>,
225
226    /// HuggingFace class name from `config.json::architectures[0]`
227    /// (e.g., "Qwen2ForCausalLM", "LlamaForCausalLM"). Distinct from
228    /// `architecture` (family) and `model_type`. PMAT-690 P0-K stamps
229    /// this so downstream `apr pretrain --init` can propagate it into
230    /// the trained checkpoint's metadata.
231    #[serde(default, skip_serializing_if = "Option::is_none")]
232    pub hf_architecture: Option<String>,
233
234    /// HuggingFace `config.json::model_type` (e.g., "qwen2", "llama").
235    /// PMAT-690 P0-K stamps this alongside `hf_architecture` so the
236    /// import→pretrain→export chain has a single source of truth.
237    #[serde(default, skip_serializing_if = "Option::is_none")]
238    pub hf_model_type: Option<String>,
239
240    /// Hidden dimension size
241    #[serde(default, skip_serializing_if = "Option::is_none")]
242    pub hidden_size: Option<usize>,
243
244    /// Number of transformer layers
245    #[serde(default, skip_serializing_if = "Option::is_none")]
246    pub num_layers: Option<usize>,
247
248    /// Number of attention heads
249    #[serde(default, skip_serializing_if = "Option::is_none")]
250    pub num_heads: Option<usize>,
251
252    /// Number of key-value heads (for GQA, defaults to num_heads)
253    #[serde(default, skip_serializing_if = "Option::is_none")]
254    pub num_kv_heads: Option<usize>,
255
256    /// Vocabulary size
257    #[serde(default, skip_serializing_if = "Option::is_none")]
258    pub vocab_size: Option<usize>,
259
260    /// FFN intermediate dimension
261    #[serde(default, skip_serializing_if = "Option::is_none")]
262    pub intermediate_size: Option<usize>,
263
264    /// Maximum context/sequence length
265    #[serde(default, skip_serializing_if = "Option::is_none")]
266    pub max_position_embeddings: Option<usize>,
267
268    /// RoPE theta for position encoding
269    #[serde(default, skip_serializing_if = "Option::is_none")]
270    pub rope_theta: Option<f32>,
271
272    /// RoPE type: 0=NORM (adjacent pairs), 2=NEOX (split halves)
273    /// CORRECTNESS-011: Qwen2.5 models require rope_type=2 (NEOX style)
274    #[serde(default, skip_serializing_if = "Option::is_none")]
275    pub rope_type: Option<u32>,
276
277    /// Layer norm epsilon
278    #[serde(default, skip_serializing_if = "Option::is_none")]
279    pub rms_norm_eps: Option<f32>,
280
281    /// Explicit head dimension (overrides hidden_size / num_heads for Qwen3+)
282    #[serde(default, skip_serializing_if = "Option::is_none")]
283    pub head_dim: Option<usize>,
284
285    /// Number of MoE experts
286    #[serde(default, skip_serializing_if = "Option::is_none")]
287    pub num_experts: Option<usize>,
288
289    /// Number of experts selected per token
290    #[serde(default, skip_serializing_if = "Option::is_none")]
291    pub num_experts_per_tok: Option<usize>,
292
293    /// MoE expert intermediate/FFN dimension
294    #[serde(default, skip_serializing_if = "Option::is_none")]
295    pub moe_intermediate_size: Option<usize>,
296
297    /// Custom key-value pairs
298    #[serde(default, flatten)]
299    pub custom: HashMap<String, serde_json::Value>,
300}
301
302/// Special tokens for chat templates (CTA-04)
303#[derive(Debug, Clone, Default, Serialize, Deserialize)]
304pub struct ChatSpecialTokens {
305    /// Beginning of sequence token
306    #[serde(default)]
307    pub bos_token: Option<String>,
308
309    /// End of sequence token
310    #[serde(default)]
311    pub eos_token: Option<String>,
312
313    /// Unknown token
314    #[serde(default)]
315    pub unk_token: Option<String>,
316
317    /// Padding token
318    #[serde(default)]
319    pub pad_token: Option<String>,
320
321    /// ChatML start token (<|im_start|>)
322    #[serde(default)]
323    pub im_start_token: Option<String>,
324
325    /// ChatML end token (<|im_end|>)
326    #[serde(default)]
327    pub im_end_token: Option<String>,
328}
329
330impl AprV2Metadata {
331    /// Create new empty metadata
332    #[must_use]
333    pub fn new(model_type: impl Into<String>) -> Self {
334        Self {
335            model_type: model_type.into(),
336            ..Default::default()
337        }
338    }
339
340    /// Serialize to JSON bytes
341    ///
342    /// # Errors
343    /// Returns error if serialization fails.
344    pub fn to_json(&self) -> Result<Vec<u8>, V2FormatError> {
345        serde_json::to_vec(self).map_err(|e| V2FormatError::MetadataError(e.to_string()))
346    }
347
348    /// Serialize to pretty JSON string
349    ///
350    /// # Errors
351    /// Returns error if serialization fails.
352    pub fn to_json_pretty(&self) -> Result<String, V2FormatError> {
353        serde_json::to_string_pretty(self).map_err(|e| V2FormatError::MetadataError(e.to_string()))
354    }
355
356    /// Canonicalize HF/GGUF-style alias keys into typed struct fields
357    /// (C-APR-MERGE-RUNNABLE / FALSIFY-APR-MERGE-RUNNABLE-001).
358    ///
359    /// Import-produced APR files carry HuggingFace-style dimension keys
360    /// (`num_hidden_layers`, `num_attention_heads`, `num_key_value_heads`, …)
361    /// which land in `custom` because this struct has no serde aliases.
362    /// Realizar's `AprMetadata` deserializer DOES alias them — so a file
363    /// containing BOTH the canonical field (even as an explicit `null`)
364    /// AND an alias key makes serde fail with "duplicate field", which
365    /// realizar's mmap loader swallows via `unwrap_or_default()` —
366    /// silently dropping ALL metadata (architecture, dims, embedded
367    /// tokenizer) and producing C-01 / "no tokenizer in APR metadata"
368    /// failures on a file that physically contains everything.
369    ///
370    /// This method promotes alias values into the typed fields (when the
371    /// field is unset) and REMOVES the alias keys from `custom`, so a
372    /// re-serialized container has exactly one spelling per dimension.
373    pub fn canonicalize_hf_aliases(&mut self) {
374        fn take_usize(
375            custom: &mut HashMap<String, serde_json::Value>,
376            keys: &[&str],
377        ) -> Option<usize> {
378            let mut found = None;
379            for k in keys {
380                if let Some(v) = custom.remove(*k) {
381                    if found.is_none() {
382                        found = v.as_u64().and_then(|n| usize::try_from(n).ok());
383                    }
384                }
385            }
386            found
387        }
388        fn take_f32(custom: &mut HashMap<String, serde_json::Value>, keys: &[&str]) -> Option<f32> {
389            let mut found = None;
390            for k in keys {
391                if let Some(v) = custom.remove(*k) {
392                    if found.is_none() {
393                        #[allow(clippy::cast_possible_truncation)]
394                        {
395                            found = v.as_f64().map(|n| n as f32);
396                        }
397                    }
398                }
399            }
400            found
401        }
402
403        // Alias groups mirror realizar::apr::AprMetadata serde aliases (PMAT-111).
404        let v = take_usize(&mut self.custom, &["hidden_dim", "d_model", "n_embd"]);
405        self.hidden_size = self.hidden_size.or(v);
406        let v = take_usize(
407            &mut self.custom,
408            &["num_hidden_layers", "n_layers", "n_layer"],
409        );
410        self.num_layers = self.num_layers.or(v);
411        let v = take_usize(
412            &mut self.custom,
413            &["num_attention_heads", "n_heads", "n_head"],
414        );
415        self.num_heads = self.num_heads.or(v);
416        let v = take_usize(&mut self.custom, &["num_key_value_heads", "n_kv_heads"]);
417        self.num_kv_heads = self.num_kv_heads.or(v);
418        let v = take_usize(&mut self.custom, &["n_vocab"]);
419        self.vocab_size = self.vocab_size.or(v);
420        let v = take_usize(
421            &mut self.custom,
422            &["ffn_dim", "intermediate_dim", "n_inner"],
423        );
424        self.intermediate_size = self.intermediate_size.or(v);
425        let v = take_usize(
426            &mut self.custom,
427            &["max_seq_len", "context_length", "n_ctx"],
428        );
429        self.max_position_embeddings = self.max_position_embeddings.or(v);
430        let v = take_f32(&mut self.custom, &["layer_norm_eps", "norm_eps"]);
431        self.rms_norm_eps = self.rms_norm_eps.or(v);
432    }
433
434    /// Deserialize from JSON bytes
435    ///
436    /// # Errors
437    /// Returns error if deserialization fails.
438    pub fn from_json(data: &[u8]) -> Result<Self, V2FormatError> {
439        // ALB-107: Parse as Value first to handle duplicate keys in metadata.
440        // Entrenar checkpoints (v1-v9) may have duplicate fields like rms_norm_eps
441        // due to #[serde(flatten)] serializing both struct field (null) and custom
442        // map entry. Value::Object deduplicates (last value wins).
443        let value: serde_json::Value = serde_json::from_slice(data)
444            .map_err(|e| V2FormatError::MetadataError(e.to_string()))?;
445        serde_json::from_value(value).map_err(|e| V2FormatError::MetadataError(e.to_string()))
446    }
447}
448
449/// Quantization metadata
450#[derive(Debug, Clone, Default, Serialize, Deserialize)]
451pub struct QuantizationMetadata {
452    /// Quantization type (e.g., "int8", "int4", "fp16")
453    pub quant_type: String,
454    /// Bits per weight
455    pub bits: u8,
456    /// Block size for block quantization
457    pub block_size: Option<usize>,
458    /// Whether symmetric quantization
459    pub symmetric: bool,
460}
461
462/// Sharding metadata for multi-file models
463#[derive(Debug, Clone, Default, Serialize, Deserialize)]
464pub struct ShardingMetadata {
465    /// Total number of shards
466    pub shard_count: usize,
467    /// This shard's index (0-based)
468    pub shard_index: usize,
469    /// Total size across all shards
470    pub total_size: u64,
471    /// Shard file pattern (e.g., "model-{:05d}-of-{:05d}.apr")
472    pub pattern: Option<String>,
473}
474
475// ============================================================================
476// Tensor Index
477// ============================================================================
478
479/// Tensor index entry (fixed size for efficient lookup)
480#[derive(Debug, Clone)]
481pub struct TensorIndexEntry {
482    /// Tensor name (up to 256 bytes)
483    pub name: String,
484    /// Data type
485    pub dtype: TensorDType,
486    /// Shape dimensions
487    pub shape: Vec<usize>,
488    /// Offset in data section (64-byte aligned)
489    pub offset: u64,
490    /// Size in bytes
491    pub size: u64,
492}
apr_format/v2/header_impl.rs

apr_format/v2/
header_impl.rs
