open-standards/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/load_onnx.rs

//! Real-model bring-up for ONNX files.
//!
//! ```bash
//! cargo run --release --example load_onnx -- path/to/model.onnx
//! ```

use jouleclaw_loader_gguf::onnx::{
    read_onnx_file, OnnxDtype, OnnxTensorData,
};
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::time::Instant;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("usage: load_onnx <path.onnx>");
        std::process::exit(1);
    }
    let path = PathBuf::from(&args[1]);
    println!("model:  {}", path.display());
    let t0 = Instant::now();
    let m = read_onnx_file(&path).map_err(|e| format!("parse: {:?}", e))?;
    let load_secs = t0.elapsed().as_secs_f64();

    println!();
    println!("=== load ===");
    println!("  load_wall_clock_secs: {:.3}", load_secs);
    println!("  ir_version: {}", m.ir_version);
    if let Some(p) = &m.producer_name {
        println!("  producer:   {}", p);
    }
    println!("  tensor_count: {}", m.tensors.len());
    println!("  mmap_len: {} ({})", m.mmap.len(), humanize_bytes(m.mmap.len() as u64));

    // Byte accounting.
    let mut raw_bytes: u64 = 0;
    let mut typed_bytes: u64 = 0;
    let mut external_count = 0u32;
    let mut external_bytes: u64 = 0;
    for t in &m.tensors {
        match &t.data {
            OnnxTensorData::RawData { len, .. } => raw_bytes += *len as u64,
            OnnxTensorData::TypedArray(a) => {
                typed_bytes += (a.float_data.len() * 4
                    + a.int32_data.len() * 4
                    + a.int64_data.len() * 8
                    + a.uint64_data.len() * 8
                    + a.double_data.len() * 8) as u64;
            }
            OnnxTensorData::External { length, .. } => {
                external_count += 1;
                external_bytes += length.unwrap_or(0);
            }
        }
    }
    println!();
    println!("=== tensor data ===");
    println!("  raw_data_bytes:   {} ({})", raw_bytes, humanize_bytes(raw_bytes));
    println!("  typed_array_bytes:{} ({})", typed_bytes, humanize_bytes(typed_bytes));
    println!(
        "  external:         {} tensors, {} ({} on disk in sibling files)",
        external_count,
        external_bytes,
        humanize_bytes(external_bytes),
    );

    // Dtype distribution.
    let mut counts: BTreeMap<String, (u32, u64)> = BTreeMap::new();
    for t in &m.tensors {
        let key = format!("{:?}", t.dtype);
        let bytes = match (&t.data, t.dtype.element_size()) {
            (OnnxTensorData::RawData { len, .. }, _) => *len as u64,
            (OnnxTensorData::External { length, .. }, _) => length.unwrap_or(0),
            (OnnxTensorData::TypedArray(_), Some(es)) => {
                (t.n_elements() as u64) * (es as u64)
            }
            (OnnxTensorData::TypedArray(_), None) => 0,
        };
        let e = counts.entry(key).or_insert((0, 0));
        e.0 += 1;
        e.1 += bytes;
    }
    println!();
    println!("=== dtype distribution ===");
    for (k, (c, b)) in &counts {
        println!("  {:<14} count={:>5}  size={}", k, c, humanize_bytes(*b));
    }

    // Top tensors by size.
    let mut sorted: Vec<&_> = m.tensors.iter().collect();
    sorted.sort_by_key(|t| {
        std::cmp::Reverse(match &t.data {
            OnnxTensorData::RawData { len, .. } => *len as u64,
            OnnxTensorData::External { length, .. } => length.unwrap_or(0),
            OnnxTensorData::TypedArray(_) => {
                (t.n_elements() as u64) * (t.dtype.element_size().unwrap_or(0) as u64)
            }
        })
    });
    println!();
    println!("=== top 10 tensors by byte size ===");
    for t in sorted.iter().take(10) {
        let bytes = match &t.data {
            OnnxTensorData::RawData { len, .. } => *len as u64,
            OnnxTensorData::External { length, .. } => length.unwrap_or(0),
            OnnxTensorData::TypedArray(_) => {
                (t.n_elements() as u64) * (t.dtype.element_size().unwrap_or(0) as u64)
            }
        };
        let kind = match &t.data {
            OnnxTensorData::RawData { .. } => "raw",
            OnnxTensorData::TypedArray(_) => "typed",
            OnnxTensorData::External { .. } => "extern",
        };
        println!(
            "  [{:>6}] {:?} {:<60} {:>14} ({})",
            kind,
            t.dtype,
            truncate(&t.name, 60),
            bytes,
            humanize_bytes(bytes),
        );
        let _ = OnnxDtype::Float; // silence unused-import in some build configs
    }

    if !m.metadata.is_empty() {
        println!();
        println!("=== metadata ===");
        for (k, v) in &m.metadata {
            let v_display = if v.len() > 80 { format!("{}…", &v[..80]) } else { v.clone() };
            println!("  {}: {}", k, v_display);
        }
    }
    Ok(())
}

fn truncate(s: &str, max: usize) -> String {
    if s.len() <= max {
        s.to_string()
    } else {
        format!("…{}", &s[s.len() - max + 1..])
    }
}

fn humanize_bytes(n: u64) -> String {
    const KB: u64 = 1024;
    const MB: u64 = KB * 1024;
    const GB: u64 = MB * 1024;
    if n >= GB {
        format!("{:.2} GiB", n as f64 / GB as f64)
    } else if n >= MB {
        format!("{:.2} MiB", n as f64 / MB as f64)
    } else if n >= KB {
        format!("{:.2} KiB", n as f64 / KB as f64)
    } else {
        format!("{} B", n)
    }
}