inferencelayer 0.2.9

//! `chatterbox-dbg` — Phase-0 manifest dumper for the Chatterbox checkpoints.
//!
//! Rust-only (no torch): parses each `.safetensors` header and prints a COLLAPSED tensor manifest —
//! per-layer repetition (`...layers.0...`, `...layers.1...`) is folded to `...layers.{i}...` with a
//! count, so a 30-layer backbone reads as a handful of lines. This is how the P1 loaders learn the
//! ground-truth tensor names/shapes instead of guessing PyTorch module paths.
//!
//! Usage:
//!   cargo run --release --features "cli encoder-cpu" --bin chatterbox-dbg -- [DIR]
//!   (DIR defaults to ~/.cache/inferencelayer/chatterbox)

use anyhow::Result;
use inferencelayer::chatterbox::{TensorInfo, cfg, read_manifest};
use std::collections::BTreeMap;
use std::path::PathBuf;

fn main() -> Result<()> {
    let dir: PathBuf = std::env::args()
        .nth(1)
        .map(PathBuf::from)
        .unwrap_or_else(default_dir);
    println!("chatterbox checkpoint dir: {}\n", dir.display());

    let files = [
        ("VE  (voice encoder)", cfg::FILE_VE),
        ("T3  (multilingual)", cfg::FILE_T3_MTL),
        ("T3  (english base)", cfg::FILE_T3_EN),
        ("S3Gen (multilingual)", cfg::FILE_S3GEN_MTL),
        ("S3Gen (english base)", cfg::FILE_S3GEN_EN),
    ];
    for (label, fname) in files {
        let path = dir.join(fname);
        if !path.exists() {
            println!("── {label}: {fname}  [not present]\n");
            continue;
        }
        match read_manifest(&path) {
            Ok(m) => dump(label, fname, &m),
            Err(e) => println!("── {label}: {fname}  [ERROR: {e:#}]\n"),
        }
    }
    Ok(())
}

fn default_dir() -> PathBuf {
    let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
    PathBuf::from(home).join(".cache/inferencelayer/chatterbox")
}

/// Fold digit-runs between dots to `{i}` so per-layer tensors collapse into one templated row.
fn template(name: &str) -> String {
    name.split('.')
        .map(|seg| {
            if !seg.is_empty() && seg.bytes().all(|b| b.is_ascii_digit()) {
                "{i}"
            } else {
                seg
            }
        })
        .collect::<Vec<_>>()
        .join(".")
}

struct Group {
    count: usize,
    dtype: String,
    example_shape: Vec<usize>,
    numel: usize,
}

fn dump(label: &str, fname: &str, manifest: &[TensorInfo]) {
    let total_numel: usize = manifest.iter().map(|t| t.numel()).sum();
    let mut groups: BTreeMap<String, Group> = BTreeMap::new();
    for t in manifest {
        let key = template(&t.name);
        let g = groups.entry(key).or_insert_with(|| Group {
            count: 0,
            dtype: t.dtype.clone(),
            example_shape: t.shape.clone(),
            numel: 0,
        });
        g.count += 1;
        g.numel += t.numel();
    }
    println!("── {label}: {fname}",);
    println!(
        "   {} tensors, {} templated groups, {:.1}M params",
        manifest.len(),
        groups.len(),
        total_numel as f64 / 1e6
    );
    for (tmpl, g) in &groups {
        let mult = if g.count > 1 {
            format!("×{}", g.count)
        } else {
            String::new()
        };
        println!(
            "   {:<58} {:<5} {:<4} {:?}",
            tmpl, g.dtype, mult, g.example_shape
        );
    }
    println!();
}