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

//! Streaming STT demo: transcribe a WAV via Kyutai's 1B en/fr model.
//!
//! First run downloads ~3 GB from `kyutai/stt-1b-en_fr` to the HF cache.
//!
//! Usage:
//! ```
//! cargo run -p rtx-csm --release --features metal --example stt_demo -- \
//!     --in /tmp/csm_24k.wav
//! ```
//!
//! Note: until sentencepiece detok is wired (Phase 6a polish), the output
//! is raw token IDs per word. The first version is intentionally minimal —
//! demonstrates that the streaming pipeline is connected end to end.

use anyhow::Result;
use clap::Parser;
use rtx_csm::{audio_io, stt::{AsrEvent, Stt, SAMPLE_RATE}};
use std::path::PathBuf;

#[derive(Debug, Parser)]
#[command(name = "stt_demo")]
struct Cli {
    /// Input WAV (any rate / channels — resampled to 24 kHz mono).
    #[arg(long = "in")]
    input: PathBuf,
    /// Force CPU device.
    #[arg(long)]
    cpu: bool,
}

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

    let device = if cli.cpu {
        candle_core::Device::Cpu
    } else if candle_core::utils::metal_is_available() {
        candle_core::Device::new_metal(0)?
    } else {
        candle_core::Device::Cpu
    };
    println!("device: {device:?}");

    let t = std::time::Instant::now();
    let mut stt = Stt::load_default(&device)?;
    println!("loaded Kyutai STT 1B en/fr in {:.2}s", t.elapsed().as_secs_f32());

    // Load WAV at 24 kHz mono (Mimi's expected input rate).
    let samples = audio_io::load_mono_at_rate(&cli.input, SAMPLE_RATE)?;
    println!(
        "loaded {}: {} samples ({:.2}s @ {} Hz)",
        cli.input.display(),
        samples.len(),
        samples.len() as f32 / SAMPLE_RATE as f32,
        SAMPLE_RATE
    );

    // The 1B en/fr STT model expects:
    //   - 0.0 seconds of silence prefix (no warmup needed)
    //   - 0.5 seconds of silence suffix (= 6.25 frames @ 12.5 Hz, round up to 7)
    //     to flush the asr_delay-buffered predictions at end of audio.
    // Without the suffix the model produces only pad tokens. See HF
    // config.json `stt_config` and the reference Python script.
    const PREFIX_SILENCE_SECS: f32 = 0.0;
    const SUFFIX_SILENCE_SECS: f32 = 2.0;
    let mut audio_with_padding =
        vec![0.0f32; (PREFIX_SILENCE_SECS * SAMPLE_RATE as f32) as usize];
    audio_with_padding.extend_from_slice(&samples);
    audio_with_padding
        .extend(std::iter::repeat(0.0f32).take((SUFFIX_SILENCE_SECS * SAMPLE_RATE as f32) as usize));
    println!(
        "padded with {:.1}s prefix + {:.1}s suffix silence -> {} samples",
        PREFIX_SILENCE_SECS,
        SUFFIX_SILENCE_SECS,
        audio_with_padding.len()
    );

    // Stream in 1-second chunks so we can observe streaming behavior.
    let chunk_size = SAMPLE_RATE as usize;
    let mut all_events: Vec<AsrEvent> = Vec::new();
    let t = std::time::Instant::now();
    for (i, chunk) in audio_with_padding.chunks(chunk_size).enumerate() {
        let evs = stt.step_pcm(chunk)?;
        let n_step = evs.iter().filter(|e| matches!(e, AsrEvent::Step { .. })).count();
        let n_word = evs.iter().filter(|e| matches!(e, AsrEvent::Word { .. })).count();
        let n_end = evs.iter().filter(|e| matches!(e, AsrEvent::EndWord { .. })).count();
        println!("[chunk {i}] events: step={n_step} word={n_word} endword={n_end}");
        all_events.extend(evs);
    }
    let evs_finish = stt.finish()?;
    all_events.extend(evs_finish);
    println!("inference: {:.2}s", t.elapsed().as_secs_f32());

    // Pair Word with the next EndWord to get full timing, then detokenize.
    let mut words = 0usize;
    let mut full_text = String::new();
    let mut pending: Option<(Vec<u32>, f64)> = None;
    for ev in &all_events {
        match ev {
            AsrEvent::Word {
                tokens,
                start_time,
                ..
            } => {
                pending = Some((tokens.clone(), *start_time));
            }
            AsrEvent::EndWord { stop_time, .. } => {
                if let Some((tokens, start)) = pending.take() {
                    let text = stt
                        .decode_word_text(&tokens)
                        .unwrap_or_default();
                    println!("  ({:.2}s - {:.2}s) {}", start, stop_time, text);
                    if !full_text.is_empty() && !text.is_empty() {
                        full_text.push(' ');
                    }
                    full_text.push_str(&text);
                    words += 1;
                }
            }
            AsrEvent::Step { .. } => {}
        }
    }
    println!("\n== transcript ({} words) ==", words);
    println!("{}", full_text.trim());

    Ok(())
}