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

//! Real-model bring-up for TFLite files.
//!
//! ```bash
//! cargo run --release --example load_tflite -- path/to/model.tflite
//! ```

use jouleclaw_loader_gguf::tflite::{read_tflite_file, TfliteBuffer, TfliteDtype};
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_tflite <path.tflite>");
        std::process::exit(1);
    }
    let path = PathBuf::from(&args[1]);
    println!("model:  {}", path.display());
    let t0 = Instant::now();
    let m = read_tflite_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!("  flatbuffer_schema_version: {}", m.version);
    if !m.description.is_empty() {
        println!("  description: {}", m.description);
    }
    println!("  subgraph_count: {}", m.subgraphs.len());
    println!("  buffer_count:   {}", m.buffers.len());
    println!("  mmap_len: {} ({})", m.mmap.len(), humanize_bytes(m.mmap.len() as u64));

    // Buffer-class accounting.
    let mut inline_bytes: u64 = 0;
    let mut external_bytes: u64 = 0;
    let mut empty = 0u32;
    for b in &m.buffers {
        match b {
            TfliteBuffer::Inline { len, .. } => inline_bytes += *len as u64,
            TfliteBuffer::External { size, .. } => external_bytes += *size,
            TfliteBuffer::Empty => empty += 1,
        }
    }
    println!();
    println!("=== buffer storage ===");
    println!("  inline_bytes:   {} ({})", inline_bytes, humanize_bytes(inline_bytes));
    println!("  external_bytes: {} ({})", external_bytes, humanize_bytes(external_bytes));
    println!("  empty_buffers:  {}", empty);

    // Per-subgraph summary.
    println!();
    println!("=== subgraphs ===");
    for (i, sg) in m.subgraphs.iter().enumerate() {
        let name = if sg.name.is_empty() { "(unnamed)" } else { &sg.name };
        println!(
            "  [{:>2}] {} — tensors={}  inputs={}  outputs={}",
            i, name, sg.tensors.len(), sg.inputs.len(), sg.outputs.len()
        );
    }

    // Dtype distribution across ALL tensors.
    let mut counts: BTreeMap<String, (u32, u64)> = BTreeMap::new();
    for sg in &m.subgraphs {
        for t in &sg.tensors {
            let key = format!("{:?}", t.dtype);
            let bytes = match m.buffers.get(t.buffer) {
                Some(TfliteBuffer::Inline { len, .. }) => *len as u64,
                Some(TfliteBuffer::External { size, .. }) => *size,
                _ => 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 10 tensors by buffer size.
    let mut all_tensors: Vec<_> = m.subgraphs
        .iter()
        .flat_map(|sg| sg.tensors.iter())
        .collect();
    all_tensors.sort_by_key(|t| {
        std::cmp::Reverse(match m.buffers.get(t.buffer) {
            Some(TfliteBuffer::Inline { len, .. }) => *len as u64,
            Some(TfliteBuffer::External { size, .. }) => *size,
            _ => 0,
        })
    });
    println!();
    println!("=== top 10 tensors by buffer size ===");
    for t in all_tensors.iter().take(10) {
        let bytes = match m.buffers.get(t.buffer) {
            Some(TfliteBuffer::Inline { len, .. }) => *len as u64,
            Some(TfliteBuffer::External { size, .. }) => *size,
            _ => 0,
        };
        println!(
            "  {:?} {:<60} {:>14} ({})",
            t.dtype,
            truncate(&t.name, 60),
            bytes,
            humanize_bytes(bytes),
        );
        let _ = TfliteDtype::Float32; // silence unused-import
    }
    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)
    }
}