axonml_serialize/
bundle.rs

1//! Model Bundle — `.axonml` container with architecture + hyperparameters + flat weights
2//!
3//! A `ModelBundle` is the portable on-disk format that `tools/model_converter/convert.py`
4//! consumes to reconstruct a trained model in PyTorch and emit ONNX. Unlike the
5//! generic `StateDict` (a dict of named tensors), a bundle carries three things
6//! the ONNX converter needs to rebuild the architecture:
7//!
8//! 1. `architecture` — an enum tag (`sentinel`, `lstm_autoencoder`, `gru_predictor`,
9//!    `rnn`, `phantom`, `conv1d`, `conv2d`, `res_net`, `vgg`, `bert`, `gpt2`, `vi_t`,
10//!    `nexus`) matching `Architecture` in the AxonML model zoo.
11//! 2. `hyperparameters` — `hidden_dim`, `num_layers`, `sequence_length`, etc.
12//!    These, together with `input_features`, completely determine layer shapes.
13//! 3. `weights` — flat `Vec<f32>` in the exact layer-by-layer order the Python
14//!    converter expects (see `tools/model_converter/convert.py::build_*`).
15//!
16//! # Binary Layout
17//!
18//! ```text
19//!   0  1  2  3  4  5  6  7  8  9 10  ...
20//! | A  X  O  N  M  L | V | H_LEN (u32 LE) | HEADER JSON | W_LEN (u32 LE) | WEIGHTS JSON |
21//! ```
22//!
23//! - `V = 1` is the current format version.
24//! - HEADER is a small metadata blob (architecture name, input_features, param count,
25//!   quantization flag). It's separately decodable without reading the weights blob —
26//!   useful for model registries that only want to display metadata.
27//! - WEIGHTS holds the full `ModelBundle` payload (architecture + hyperparameters +
28//!   flat weights + optional input-normalization stats).
29//!
30//! # File
31//! `crates/axonml-serialize/src/bundle.rs`
32//!
33//! # Author
34//! Andrew Jewell Sr. — AutomataNexus LLC
35//! ORCID: 0009-0005-2158-7060
36//!
37//! # Updated
38//! April 16, 2026 11:15 PM EST
39//!
40//! # Disclaimer
41//! Use at own risk. This software is provided "as is", without warranty of any
42//! kind, express or implied. The author and AutomataNexus shall not be held
43//! liable for any damages arising from the use of this software.
44
45use std::collections::HashMap;
46use std::fs;
47use std::io::{self, Write};
48use std::path::Path;
49
50use serde::{Deserialize, Serialize};
51
52// =============================================================================
53// On-disk Format Constants
54// =============================================================================
55
56/// Magic bytes at the start of every `.axonml` bundle file.
57pub const AXONML_MAGIC: &[u8; 6] = b"AXONML";
58
59/// Current bundle format version.
60pub const AXONML_BUNDLE_VERSION: u8 = 1;
61
62// =============================================================================
63// Error
64// =============================================================================
65
66/// Errors produced by bundle save/load.
67#[derive(Debug)]
68pub enum BundleError {
69    /// I/O failure.
70    Io(io::Error),
71    /// JSON (de)serialization failure.
72    Serde(serde_json::Error),
73    /// The file does not start with `AXONML` magic bytes.
74    BadMagic,
75    /// The file version byte is not `AXONML_BUNDLE_VERSION`.
76    BadVersion(u8),
77    /// The file ended before the header or weights section was complete.
78    Truncated(&'static str),
79    /// A field required for round-tripping is missing or malformed.
80    Invalid(String),
81}
82
83impl std::fmt::Display for BundleError {
84    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
85        match self {
86            Self::Io(e) => write!(f, "bundle I/O error: {e}"),
87            Self::Serde(e) => write!(f, "bundle JSON error: {e}"),
88            Self::BadMagic => write!(f, "invalid .axonml file: bad magic bytes"),
89            Self::BadVersion(v) => write!(
90                f,
91                "unsupported .axonml bundle version: {v} (expected {AXONML_BUNDLE_VERSION})"
92            ),
93            Self::Truncated(what) => write!(f, "invalid .axonml file: truncated {what}"),
94            Self::Invalid(msg) => write!(f, "invalid .axonml bundle: {msg}"),
95        }
96    }
97}
98
99impl std::error::Error for BundleError {}
100
101impl From<io::Error> for BundleError {
102    fn from(e: io::Error) -> Self {
103        Self::Io(e)
104    }
105}
106
107impl From<serde_json::Error> for BundleError {
108    fn from(e: serde_json::Error) -> Self {
109        Self::Serde(e)
110    }
111}
112
113/// Convenience alias for bundle operation results.
114pub type BundleResult<T> = Result<T, BundleError>;
115
116// =============================================================================
117// Header (first decodable metadata block)
118// =============================================================================
119
120/// Lightweight header decoded before the weights blob. Kept in sync with
121/// the Python converter's `parse_axonml` header JSON (must round-trip verbatim).
122#[derive(Debug, Clone, Serialize, Deserialize)]
123pub struct BundleHeader {
124    /// Architecture tag (`sentinel`, `lstm_autoencoder`, `res_net`, `bert`, etc.).
125    pub architecture: String,
126    /// Number of input features (for MLPs) or input channels (for CNNs).
127    pub input_features: usize,
128    /// Total parameter count (flat length of the weights vector).
129    pub num_parameters: usize,
130    /// Whether weights are post-quantization (INT8 stored as f32 dequantized).
131    pub quantized: bool,
132    /// Quantization bit width when `quantized = true`.
133    pub quant_bits: Option<u8>,
134}
135
136// =============================================================================
137// Bundle (full payload)
138// =============================================================================
139
140/// Full bundle payload carrying everything the ONNX converter needs.
141///
142/// The `hyperparameters` map uses flexible JSON to accommodate per-architecture
143/// configs (e.g. LSTM wants `hidden_dim` + `num_layers` + `sequence_length`,
144/// ViT wants `patch_size` + `num_heads`, etc). Keys are snake_case to match
145/// the Python converter's expectations.
146#[derive(Debug, Clone, Serialize, Deserialize)]
147pub struct ModelBundle {
148    /// Architecture tag (same as `BundleHeader::architecture`).
149    pub architecture: String,
150    /// Number of input features.
151    pub input_features: usize,
152    /// Free-form hyperparameter map (hidden_dim, num_layers, sequence_length, ...).
153    #[serde(default)]
154    pub hyperparameters: HashMap<String, serde_json::Value>,
155    /// Flat weight vector in layer-by-layer order expected by the converter.
156    pub weights: Vec<f32>,
157    /// Per-feature input-normalization means (empty if no normalization was applied).
158    #[serde(default)]
159    pub norm_means: Vec<f32>,
160    /// Per-feature input-normalization std-devs.
161    #[serde(default)]
162    pub norm_stds: Vec<f32>,
163    /// Anomaly/decision threshold for binary models (optional).
164    #[serde(default)]
165    pub anomaly_threshold: Option<f32>,
166}
167
168impl ModelBundle {
169    /// Construct a new bundle with no normalization stats and no anomaly threshold.
170    pub fn new(architecture: &str, input_features: usize, weights: Vec<f32>) -> Self {
171        Self {
172            architecture: architecture.to_string(),
173            input_features,
174            hyperparameters: HashMap::new(),
175            weights,
176            norm_means: Vec::new(),
177            norm_stds: Vec::new(),
178            anomaly_threshold: None,
179        }
180    }
181
182    /// Set a single hyperparameter (overwrites any existing value).
183    pub fn with_hyperparam(mut self, key: &str, value: impl Into<serde_json::Value>) -> Self {
184        self.hyperparameters.insert(key.to_string(), value.into());
185        self
186    }
187
188    /// Attach per-feature normalization statistics.
189    pub fn with_normalization(mut self, means: Vec<f32>, stds: Vec<f32>) -> Self {
190        self.norm_means = means;
191        self.norm_stds = stds;
192        self
193    }
194
195    /// Attach an anomaly/decision threshold.
196    pub fn with_threshold(mut self, threshold: f32) -> Self {
197        self.anomaly_threshold = Some(threshold);
198        self
199    }
200
201    /// Derive the lightweight `BundleHeader` from this bundle.
202    pub fn header(&self) -> BundleHeader {
203        BundleHeader {
204            architecture: self.architecture.clone(),
205            input_features: self.input_features,
206            num_parameters: self.weights.len(),
207            quantized: false,
208            quant_bits: None,
209        }
210    }
211}
212
213// =============================================================================
214// Save
215// =============================================================================
216
217/// Write a bundle to disk in the `.axonml` container format.
218///
219/// Creates parent directories if they don't exist. Ensures the output path has
220/// a `.axonml` extension (renames if missing). Returns the final path written.
221pub fn save_bundle<P: AsRef<Path>>(
222    bundle: &ModelBundle,
223    path: P,
224) -> BundleResult<std::path::PathBuf> {
225    let raw = path.as_ref();
226    if let Some(parent) = raw.parent() {
227        if !parent.as_os_str().is_empty() {
228            fs::create_dir_all(parent)?;
229        }
230    }
231
232    let final_path = if raw.extension().is_none_or(|e| e != "axonml") {
233        raw.with_extension("axonml")
234    } else {
235        raw.to_path_buf()
236    };
237
238    let header_json = serde_json::to_vec(&bundle.header())?;
239    let weights_json = serde_json::to_vec(bundle)?;
240
241    let mut file = fs::File::create(&final_path)?;
242    file.write_all(AXONML_MAGIC)?;
243    file.write_all(&[AXONML_BUNDLE_VERSION])?;
244    file.write_all(&(header_json.len() as u32).to_le_bytes())?;
245    file.write_all(&header_json)?;
246    file.write_all(&(weights_json.len() as u32).to_le_bytes())?;
247    file.write_all(&weights_json)?;
248
249    Ok(final_path)
250}
251
252// =============================================================================
253// Load
254// =============================================================================
255
256/// Read a bundle from disk, returning both the header (decoded eagerly) and the
257/// full payload.
258///
259/// Errors if the file is too short, the magic bytes are wrong, the version is
260/// unsupported, or either JSON blob fails to deserialize.
261pub fn load_bundle<P: AsRef<Path>>(path: P) -> BundleResult<(BundleHeader, ModelBundle)> {
262    let data = fs::read(path)?;
263    load_bundle_from_bytes(&data)
264}
265
266/// In-memory variant of [`load_bundle`] — useful for HTTP handlers or tests.
267pub fn load_bundle_from_bytes(data: &[u8]) -> BundleResult<(BundleHeader, ModelBundle)> {
268    if data.len() < 11 {
269        return Err(BundleError::Truncated("file (too short for header)"));
270    }
271    if &data[0..6] != AXONML_MAGIC {
272        return Err(BundleError::BadMagic);
273    }
274    let version = data[6];
275    if version != AXONML_BUNDLE_VERSION {
276        return Err(BundleError::BadVersion(version));
277    }
278
279    let header_len = u32::from_le_bytes([data[7], data[8], data[9], data[10]]) as usize;
280    let header_start = 11usize;
281    let header_end = header_start
282        .checked_add(header_len)
283        .ok_or(BundleError::Invalid("header length overflow".into()))?;
284    if data.len() < header_end + 4 {
285        return Err(BundleError::Truncated("header"));
286    }
287    let header: BundleHeader = serde_json::from_slice(&data[header_start..header_end])?;
288
289    let weights_len_bytes = &data[header_end..header_end + 4];
290    let weights_len = u32::from_le_bytes([
291        weights_len_bytes[0],
292        weights_len_bytes[1],
293        weights_len_bytes[2],
294        weights_len_bytes[3],
295    ]) as usize;
296    let weights_start = header_end + 4;
297    let weights_end = weights_start
298        .checked_add(weights_len)
299        .ok_or(BundleError::Invalid("weights length overflow".into()))?;
300    if data.len() < weights_end {
301        return Err(BundleError::Truncated("weights"));
302    }
303    let bundle: ModelBundle = serde_json::from_slice(&data[weights_start..weights_end])?;
304
305    Ok((header, bundle))
306}
307
308/// Decode only the header of a `.axonml` file without reading the weights blob.
309///
310/// Useful for model registries / UIs that list metadata without paying for the
311/// full weight deserialization cost on every request.
312pub fn load_header<P: AsRef<Path>>(path: P) -> BundleResult<BundleHeader> {
313    let data = fs::read(path)?;
314    if data.len() < 11 {
315        return Err(BundleError::Truncated("file (too short for header)"));
316    }
317    if &data[0..6] != AXONML_MAGIC {
318        return Err(BundleError::BadMagic);
319    }
320    let version = data[6];
321    if version != AXONML_BUNDLE_VERSION {
322        return Err(BundleError::BadVersion(version));
323    }
324    let header_len = u32::from_le_bytes([data[7], data[8], data[9], data[10]]) as usize;
325    let header_end = 11usize
326        .checked_add(header_len)
327        .ok_or(BundleError::Invalid("header length overflow".into()))?;
328    if data.len() < header_end {
329        return Err(BundleError::Truncated("header"));
330    }
331    Ok(serde_json::from_slice(&data[11..header_end])?)
332}
333
334// =============================================================================
335// Tests
336// =============================================================================
337
338#[cfg(test)]
339mod tests {
340    use super::*;
341    use tempfile::NamedTempFile;
342
343    #[test]
344    fn round_trip_sentinel_bundle() {
345        let weights: Vec<f32> = (0..256).map(|i| i as f32 * 0.01).collect();
346        let bundle = ModelBundle::new("sentinel", 11, weights.clone())
347            .with_hyperparam("hidden_dim", 128)
348            .with_hyperparam("num_layers", 2)
349            .with_normalization(vec![0.0; 11], vec![1.0; 11])
350            .with_threshold(0.5);
351
352        let tmp = NamedTempFile::new().unwrap();
353        let final_path = save_bundle(&bundle, tmp.path()).unwrap();
354
355        let (header, loaded) = load_bundle(&final_path).unwrap();
356        assert_eq!(header.architecture, "sentinel");
357        assert_eq!(header.input_features, 11);
358        assert_eq!(header.num_parameters, weights.len());
359        assert!(!header.quantized);
360
361        assert_eq!(loaded.architecture, "sentinel");
362        assert_eq!(loaded.weights, weights);
363        assert_eq!(
364            loaded
365                .hyperparameters
366                .get("hidden_dim")
367                .and_then(|v| v.as_i64()),
368            Some(128)
369        );
370        assert_eq!(loaded.anomaly_threshold, Some(0.5));
371    }
372
373    #[test]
374    fn header_only_decode_skips_weights() {
375        let weights: Vec<f32> = (0..10_000).map(|i| i as f32).collect();
376        let bundle = ModelBundle::new("bert", 128, weights);
377
378        let tmp = NamedTempFile::new().unwrap();
379        let path = save_bundle(&bundle, tmp.path()).unwrap();
380
381        let header = load_header(&path).unwrap();
382        assert_eq!(header.architecture, "bert");
383        assert_eq!(header.num_parameters, 10_000);
384    }
385
386    #[test]
387    fn rejects_bad_magic() {
388        let mut bytes = vec![b'X', b'X', b'X', b'X', b'X', b'X'];
389        bytes.extend_from_slice(&[1]);
390        bytes.extend_from_slice(&0u32.to_le_bytes());
391        let err = load_bundle_from_bytes(&bytes).unwrap_err();
392        assert!(matches!(err, BundleError::BadMagic));
393    }
394
395    #[test]
396    fn rejects_bad_version() {
397        let mut bytes = AXONML_MAGIC.to_vec();
398        bytes.push(99);
399        bytes.extend_from_slice(&0u32.to_le_bytes());
400        let err = load_bundle_from_bytes(&bytes).unwrap_err();
401        assert!(matches!(err, BundleError::BadVersion(99)));
402    }
403
404    #[test]
405    fn rejects_truncated_header() {
406        let mut bytes = AXONML_MAGIC.to_vec();
407        bytes.push(AXONML_BUNDLE_VERSION);
408        bytes.extend_from_slice(&500u32.to_le_bytes()); // claims 500-byte header
409        // but only 4 more bytes of "header" provided:
410        bytes.extend_from_slice(&[0, 0, 0, 0]);
411        let err = load_bundle_from_bytes(&bytes).unwrap_err();
412        assert!(matches!(err, BundleError::Truncated(_)));
413    }
414
415    #[test]
416    fn save_adds_axonml_extension_when_missing() {
417        let bundle = ModelBundle::new("phantom", 20, vec![1.0, 2.0, 3.0]);
418        let dir = tempfile::tempdir().unwrap();
419        let no_ext = dir.path().join("my_model");
420        let final_path = save_bundle(&bundle, &no_ext).unwrap();
421        assert_eq!(final_path.extension().unwrap(), "axonml");
422        assert!(final_path.exists());
423    }
424}
axonml_serialize/bundle.rs

axonml_serialize/
bundle.rs
