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

//! Phase 13.9 — slice 1 of the candle-wav2vec2 port for word-level
//! forced alignment. Drives the design of the candle module shape +
//! the pickle/safetensors key remap.
//!
//! Target model: `facebook/wav2vec2-base-960h` — 95 M params, CTC-trained
//! on 960 h LibriSpeech, English-only. The `_lv60-ft` and `-large-960h`
//! variants share the same architecture (more layers / wider) so this
//! port should generalize.
//!
//! Why CTC + Viterbi for forced alignment:
//!   - Given a known transcript T and audio A, run the model on A to get
//!     per-frame CTC log-probs over the vocab (~32 chars for `_base-960h`)
//!   - Viterbi-decode the optimal alignment of T against the per-frame
//!     log-probs — output is a (T_token, frame_start, frame_end) table
//!   - This is what WhisperX uses (via their copy of `ctc-forced-aligner`)
//!     to cut long audio at word boundaries during data prep
//!
//! Compared to the Phase 13.8 emotion2vec port:
//!   - wav2vec2 ships `model.safetensors` natively → use mmap'd VarBuilder
//!     directly (no `pickle::read_pth_with_state` intermediate)
//!   - The transformer is POST-norm (vs emotion2vec's PRE-norm). Same
//!     shape (qkv, proj, MLP) but different order in the forward.
//!   - CTC head is a single Linear → vocab_size; no 9→5 fold needed.
//!
//! Usage:
//! ```bash
//! cargo run -p rtx-csm --release --features metal --example wav2vec2_inspect
//! cargo run -p rtx-csm --release --example wav2vec2_inspect -- --filter encoder.layer
//! ```

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

const REPO: &str = "facebook/wav2vec2-base-960h";
const SAFETENSORS_FILE: &str = "model.safetensors";
const CONFIG_FILE: &str = "config.json";
const VOCAB_FILE: &str = "vocab.json";

#[derive(Debug, Parser)]
#[command(
    name = "wav2vec2_inspect",
    about = "Dump wav2vec2 tensor keys + shapes from facebook/wav2vec2-base-960h"
)]
struct Cli {
    /// Local safetensors override; if 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,
}

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 so the user sees the
            // architecture summary even before the 378 MB safetensors
            // 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(vocab_path) = repo.get(VOCAB_FILE) {
                println!("=== {VOCAB_FILE} ===");
                match std::fs::read_to_string(&vocab_path) {
                    Ok(body) => println!("{}", body.trim()),
                    Err(e) => eprintln!("(read failed: {e})"),
                }
                println!();
            }
            repo.get(SAFETENSORS_FILE)
                .with_context(|| format!("download {SAFETENSORS_FILE} from {REPO}"))?
        }
    };
    let size = std::fs::metadata(&path)?.len();
    println!("=== wav2vec2 inspector ===");
    println!("path:  {}", path.display());
    println!("size:  {:.2} MB", size as f64 / 1e6);
    println!();

    let bytes = std::fs::read(&path)?;
    let st = safetensors::SafeTensors::deserialize(&bytes)
        .map_err(|e| anyhow::anyhow!("safetensors deserialize: {e}"))?;
    let entries: Vec<(String, Vec<usize>, String)> = st
        .tensors()
        .into_iter()
        .map(|(name, view)| {
            (
                name.to_string(),
                view.shape().to_vec(),
                format!("{:?}", view.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!(
            "  {:<40} {:>4} tensors, {:>10} params ({:.2} M)",
            prefix,
            n_tensors,
            n_params,
            *n_params as f64 / 1e6
        );
    }
    println!();

    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(())
}