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

//! Phase 13.9 — slice 3 smoke test for the wav2vec2 candle port.
//!
//! Loads `facebook/wav2vec2-base-960h` from HF Hub, runs forward on a
//! real audio file, and prints the greedy CTC decode (transcript). This
//! is the first real-weight integration of the port: every safetensors
//! key must map to a candle param of matching shape, and the resulting
//! transcript should be intelligible English.
//!
//! Usage:
//! ```bash
//! cargo run -p rtx-csm --release --features metal --example wav2vec2_smoke -- \
//!     --in /tmp/asr_test.flac
//! ```

use anyhow::{Context, Result};
use candle_core::{Device, Tensor};
use clap::Parser;
use hf_hub::api::sync::Api;
use rtx_csm::audio_io;
use rtx_csm::wav2vec2::{
    CTC_BLANK_ID, VOCAB_960H, Wav2Vec2, ctc_greedy_decode, frame_to_ms, group_into_words,
    transcript_to_token_ids, viterbi_align,
};
use std::path::PathBuf;

const REPO: &str = "facebook/wav2vec2-base-960h";
const SAFETENSORS_FILE: &str = "model.safetensors";

#[derive(Debug, Parser)]
struct Cli {
    #[arg(long = "in", default_value = "/tmp/asr_test.flac")]
    input: PathBuf,
    /// Optional known transcript to force-align. When set, run CTC
    /// Viterbi to produce per-word `(start_ms, end_ms)` boundaries
    /// instead of just greedy ASR.
    #[arg(long)]
    align: Option<String>,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().init();
    let cli = Cli::parse();
    let device = if candle_core::utils::metal_is_available() {
        Device::new_metal(0)?
    } else {
        Device::Cpu
    };
    eprintln!("device: {device:?}");

    let api = Api::new().context("hf_hub init")?;
    let path = api
        .model(REPO.to_string())
        .get(SAFETENSORS_FILE)
        .with_context(|| format!("download {SAFETENSORS_FILE} from {REPO}"))?;
    eprintln!("safetensors: {}", path.display());

    let load_t = std::time::Instant::now();
    let model = Wav2Vec2::load_from_safetensors(&path, &device)?;
    eprintln!("loaded model in {:.2}s", load_t.elapsed().as_secs_f32());

    // Load + resample audio to 16 kHz.
    let pcm = audio_io::load_mono_at_rate(&cli.input, 16_000).context("load audio")?;
    eprintln!(
        "audio: {} samples ({:.2}s @ 16 kHz)",
        pcm.len(),
        pcm.len() as f32 / 16_000.0
    );

    // wav2vec2 expects pre-normalized inputs (zero mean unit variance per
    // utterance, per HF's Wav2Vec2FeatureExtractor.do_normalize).
    let mean = pcm.iter().sum::<f32>() / pcm.len().max(1) as f32;
    let var = pcm.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / pcm.len().max(1) as f32;
    let std = var.sqrt().max(1e-7);
    let norm: Vec<f32> = pcm.iter().map(|x| (x - mean) / std).collect();

    let audio_t = Tensor::from_vec(norm, (1, 1, pcm.len()), &device)?;

    let fwd_t = std::time::Instant::now();
    let logits = model.forward(&audio_t)?;
    eprintln!("forward in {} ms", fwd_t.elapsed().as_millis());
    eprintln!("logits shape: {:?}", logits.shape().dims());

    let dec_t = std::time::Instant::now();
    let transcript = ctc_greedy_decode(&logits, VOCAB_960H)?;
    eprintln!("ctc decode in {} ms", dec_t.elapsed().as_millis());

    println!();
    println!("=== transcript ===");
    println!("{}", transcript.trim());
    println!();

    // Optional forced alignment.
    if let Some(target) = cli.align.as_deref() {
        let log_probs = candle_nn::ops::log_softmax(&logits, candle_core::D::Minus1)?;
        let tokens = transcript_to_token_ids(target, VOCAB_960H)?;
        let ali_t = std::time::Instant::now();
        let aligned = viterbi_align(&log_probs, &tokens, CTC_BLANK_ID, VOCAB_960H)?;
        eprintln!(
            "viterbi alignment in {} ms ({} tokens)",
            ali_t.elapsed().as_millis(),
            aligned.len()
        );
        let words = group_into_words(&aligned);
        println!("=== forced alignment (target = {target:?}) ===");
        for w in words.iter() {
            println!(
                "  {:<20} {:>7.0} ms .. {:>7.0} ms",
                w.word,
                frame_to_ms(w.frame_start),
                frame_to_ms(w.frame_end + 1)
            );
        }
        println!();
    }
    Ok(())
}