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

//! Phase 8.9 — standalone profile of Moonshine STT, comparable to
//! `examples/stt_profile` (Kyutai) and `examples/whisper_profile`
//! (Whisper-tiny via whisper-rs). Runs N transcriptions of the same
//! audio, reports per-call latency stats + realtime factor.
//!
//! Usage:
//! ```bash
//! cargo run -p rtx-csm --release --features metal --example moonshine_profile -- \
//!     --in /tmp/asr_test.flac --repeat 5
//! ```

use anyhow::{Context, Result};
use candle_core::{Device, Tensor};
use clap::Parser;
use hf_hub::api::sync::Api;
use rtx_csm::{audio_io, moonshine};
use std::path::PathBuf;
use std::time::Instant;

#[derive(Debug, Parser)]
struct Cli {
    #[arg(long = "in", default_value = "/tmp/asr_test.flac")]
    input: PathBuf,
    #[arg(long, default_value_t = 5)]
    repeat: usize,
    #[arg(long, default_value_t = 100)]
    max_tokens: usize,
    /// Force CPU device.
    #[arg(long)]
    cpu: bool,
}

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

    let api = Api::new()?;
    let repo = api.model("UsefulSensors/moonshine-tiny".to_string());
    let weights = repo.get("model.safetensors")?;
    let tok_path = repo.get("tokenizer.json")?;

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

    let cfg = moonshine::MoonshineConfig::tiny();
    let load_t = Instant::now();
    let (encoder, decoder) = moonshine::load_full(&weights, &device, &cfg)?;
    let tok = moonshine::load_tokenizer(&tok_path)
        .map_err(|e| anyhow::anyhow!("tok: {e}"))?;
    eprintln!("load: {:.2}s", load_t.elapsed().as_secs_f32());

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

    // Warm-up: first call pays JIT + cache init.
    let warm_t = Instant::now();
    let enc = encoder.forward(&pcm_tensor)?;
    let _warm_tokens = decoder.generate_cached(&enc, &cfg, cli.max_tokens)?;
    eprintln!("warm-up: {} ms", warm_t.elapsed().as_millis());

    // Steady-state runs.
    let mut per_call_ms: Vec<f64> = Vec::with_capacity(cli.repeat);
    let mut last_text = String::new();
    let mut last_token_count = 0usize;
    for i in 0..cli.repeat {
        let t = Instant::now();
        let enc = encoder.forward(&pcm_tensor)?;
        let tokens = decoder.generate_cached(&enc, &cfg, cli.max_tokens)?;
        let ms = t.elapsed().as_secs_f64() * 1000.0;
        per_call_ms.push(ms);
        last_text = tok.decode(&tokens, true).map_err(|e| anyhow::anyhow!("detok: {e}"))?;
        last_token_count = tokens.len();
        eprintln!("  run {}: {ms:.0} ms ({} tokens)", i + 1, tokens.len());
    }

    per_call_ms.sort_by(|a, b| a.partial_cmp(b).unwrap());
    let n = per_call_ms.len() as f64;
    let mean = per_call_ms.iter().sum::<f64>() / n;
    let p50 = per_call_ms[per_call_ms.len() / 2];
    let realtime_factor = (mean / 1000.0) / audio_secs as f64;

    println!();
    println!("=== Moonshine-tiny profile ===");
    println!("input: {} ({audio_secs:.2} s of audio)", cli.input.display());
    println!("steady-state runs: {}", per_call_ms.len());
    println!();
    println!("per-call latency:");
    println!("  mean = {mean:.0} ms");
    println!("  p50  = {p50:.0} ms");
    println!("  min  = {:.0} ms", per_call_ms[0]);
    println!("  max  = {:.0} ms", per_call_ms[per_call_ms.len() - 1]);
    println!();
    println!("realtime factor: {realtime_factor:.4}x");
    println!("  (mean / audio_duration; sub-1.0 = faster than realtime)");
    println!();
    println!("transcript ({last_token_count} tokens):");
    println!("  {last_text}");

    // For direct comparison with the other STT profiles in this crate:
    println!();
    println!("=== A/B against other STT backends in rtx-csm ===");
    println!("  Kyutai STT 1B   ~1.01x realtime  (3 GB, hardware-bound)");
    println!("  Whisper-tiny    ~0.020x          (in-process ggml, breaks CSM)");
    println!("  Moonshine-tiny  {realtime_factor:.4}x          (pure candle, no conflict)");

    Ok(())
}