rustytorch/crates/models/rtx-csm/examples/emotion2vec_inspect.rs

//! Phase 13.8 — slice 1 of the emotion2vec_plus_base candle port.
//!
//! Downloads the upstream FunASR/iic checkpoint, dumps every tensor's
//! shape + dtype, and prints a summary by parameter prefix. Used to
//! design the candle module shape and verify our port maps the right
//! weights.
//!
//! emotion2vec_plus_base architecture (from upstream):
//!   - Feature extractor: a stack of Conv1d layers operating on raw
//!     16 kHz waveform (HuBERT/WavLM-style — same as our wavlm_sv port)
//!   - Transformer encoder: 12 layers, 768-dim, 12 heads (HuBERT-base)
//!   - Pooling head + linear classifier producing a 9-class logit
//!     (anger/disgust/fear/happy/neutral/other/sad/surprised/unknown)
//!
//! Where this lives in the rtx-csm dependency hygiene:
//!   - We deliberately use HF Hub safetensors rather than ort or any
//!     C++ runtime — same path as Phase 8 Moonshine and Phase 5d
//!     WavLM-SV. No protobuf collisions, no ggml regression.
//!
//! Once the inspector confirms the layout, the next slice is a candle
//! module port that slots in via `EmotionDetector` (Phase 13.3 trait),
//! replacing the prosody-rule placeholder used in
//! `audio_to_manifest --auto-emotion-tag` and
//! `converse_server --reactive-emotion`.
//!
//! Usage:
//! ```bash
//! cargo run -p rtx-csm --release --features metal --example emotion2vec_inspect
//! cargo run -p rtx-csm --release --example emotion2vec_inspect -- --filter encoder
//! cargo run -p rtx-csm --release --example emotion2vec_inspect -- --path /local/model.pth
//! ```

use anyhow::{Context, Result};
use clap::Parser;
use hf_hub::api::sync::Api;
use std::path::PathBuf;

/// FunASR/iic publishes the model under this HF repo as a single
/// PyTorch pickle (`model.pt`). The repo also ships `config.yaml`
/// with architecture hyperparameters which we surface alongside the
/// tensor list to make the layout immediately legible.
const REPO: &str = "emotion2vec/emotion2vec_plus_base";
const PYTORCH_FILE: &str = "model.pt";
const CONFIG_FILE: &str = "config.yaml";
const TOKENS_FILE: &str = "tokens.txt";

#[derive(Debug, Parser)]
#[command(
    name = "emotion2vec_inspect",
    about = "Dump emotion2vec_plus_base tensor keys + shapes"
)]
struct Cli {
    /// Local checkpoint override. When set, skip the HF Hub download.
    #[arg(long)]
    path: Option<PathBuf>,
    /// Show only keys matching this substring.
    #[arg(long)]
    filter: Option<String>,
    /// Cap on number of keys printed (0 = unlimited).
    #[arg(long, default_value_t = 0)]
    limit: usize,
    /// Optional dict key to descend into. emotion2vec's `model.pt` is a
    /// wrapper around a state_dict; common keys are `model`, `best_state`,
    /// `state_dict`. Empty = root.
    #[arg(long)]
    key: Option<String>,
    /// Print the raw pickle object tree before tensor extraction. Useful
    /// for figuring out which `--key` to pass on a wrapped checkpoint.
    #[arg(long)]
    verbose: bool,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().init();
    let cli = Cli::parse();

    let path = match cli.path {
        Some(p) => p,
        None => {
            let api = Api::new().context("hf_hub init")?;
            let repo = api.model(REPO.to_string());
            // Pull the small sidecars first — they're cheap and make
            // the inspection output more useful even before the 1.1 GB
            // model.pt finishes streaming.
            if let Ok(cfg_path) = repo.get(CONFIG_FILE) {
                println!("=== {CONFIG_FILE} ===");
                match std::fs::read_to_string(&cfg_path) {
                    Ok(body) => println!("{body}"),
                    Err(e) => eprintln!("(read failed: {e})"),
                }
                println!();
            }
            if let Ok(tok_path) = repo.get(TOKENS_FILE) {
                println!("=== {TOKENS_FILE} ===");
                match std::fs::read_to_string(&tok_path) {
                    Ok(body) => println!("{}", body.trim()),
                    Err(e) => eprintln!("(read failed: {e})"),
                }
                println!();
            }
            repo.get(PYTORCH_FILE)
                .with_context(|| format!("download {PYTORCH_FILE} from {REPO}"))?
        }
    };
    let size = std::fs::metadata(&path)?.len();
    println!("=== emotion2vec_plus_base inspector ===");
    println!("path:  {}", path.display());
    println!("size:  {:.2} MB", size as f64 / 1e6);
    println!();

    // Branch on file extension. Safetensors uses `safetensors::SafeTensors`.
    // PyTorch .pt uses `pickle::read_pth_tensor_info` so we can descend
    // into a nested wrapper dict via `--key` — emotion2vec's checkpoint
    // is `{"model": state_dict, ...}`, like AudioSeal's `xp.cfg` /
    // `best_state` patterns.
    let entries: Vec<(String, Vec<usize>, String)> =
        if path.extension().and_then(|s| s.to_str()) == Some("safetensors") {
            let bytes = std::fs::read(&path)?;
            let st = safetensors::SafeTensors::deserialize(&bytes)
                .map_err(|e| anyhow::anyhow!("safetensors deserialize: {e}"))?;
            st.tensors()
                .into_iter()
                .map(|(name, view)| {
                    (
                        name.to_string(),
                        view.shape().to_vec(),
                        format!("{:?}", view.dtype()),
                    )
                })
                .collect()
        } else {
            let infos =
                candle_core::pickle::read_pth_tensor_info(&path, cli.verbose, cli.key.as_deref())
                    .context("pickle read")?;
            infos
                .into_iter()
                .map(|info| {
                    (
                        info.name,
                        info.layout.shape().dims().to_vec(),
                        format!("{:?}", info.dtype),
                    )
                })
                .collect()
        };

    println!("found {} tensor entries", entries.len());

    // Group by 2-component prefix for an architecture-shape summary.
    let mut grouped: std::collections::BTreeMap<String, (usize, usize)> = Default::default();
    let mut total_params: usize = 0;
    for (name, shape, _) in &entries {
        let n: usize = shape.iter().product::<usize>().max(1);
        total_params += n;
        let prefix = name.split('.').take(2).collect::<Vec<_>>().join(".");
        let entry = grouped.entry(prefix).or_insert((0, 0));
        entry.0 += 1;
        entry.1 += n;
    }
    println!(
        "total params: {} ({:.2} M)",
        total_params,
        total_params as f64 / 1e6
    );
    println!();
    println!("=== prefix summary ===");
    for (prefix, (n_tensors, n_params)) in &grouped {
        println!(
            "  {:<35} {:>4} tensors, {:>10} params ({:.2} M)",
            prefix,
            n_tensors,
            n_params,
            *n_params as f64 / 1e6
        );
    }
    println!();

    // Per-tensor dump (filtered + capped).
    println!("=== tensor list ===");
    let mut printed = 0usize;
    for (name, shape, dtype) in &entries {
        if let Some(f) = cli.filter.as_ref() {
            if !name.contains(f) {
                continue;
            }
        }
        println!("  {:<70} dtype={} shape={:?}", name, dtype, shape);
        printed += 1;
        if cli.limit > 0 && printed >= cli.limit {
            println!("  ... (truncated at --limit {})", cli.limit);
            break;
        }
    }
    Ok(())
}