Files
rustytorch/crates/models/rtx-csm/examples/moonshine_transcribe.rs
T
osobhandClaude Opus 4.7 95015c17e1 rtx-csm: Phase 8.9 — Moonshine KV cache + profile binary
KV cache for the decoder turns greedy generation from O(T^2) into O(T)
total work. Per-token decode drops modestly on short transcripts
(7.1 -> 6.0 ms/token at 49 tokens) and compounds on longer ones.

New components in src/moonshine.rs:

  RotaryCache::apply_at(x, position, t)
      Apply RoPE for a window starting at `position`. Replaces
      `apply()` for cached step (which always called positions 0..T).

  DecoderSelfAttention::forward_step(xs, cache_k, cache_v, rope, position)
      Single-token cached self-attn. Appends new K/V to per-layer cache,
      attends across full accumulated history. No causal mask needed
      (cache only contains positions <= current).

  CrossAttention::precompute_kv(enc) -> (K, V)
      One-shot encoder K/V projection for cross-attn. Reused every step.

  CrossAttention::forward_step(xs, k, v)
      Cached cross-attn. Q computed from new token; K/V from precompute.

  DecoderCache { self_k: Vec<Option<Tensor>>, self_v, cross_k, cross_v, position }

  Decoder::precompute_cross_kv(enc) -> DecoderCache
  Decoder::step(token_id, &mut cache) -> logits (1, vocab)
  Decoder::generate_cached(enc, cfg, max_tokens) -> Vec<u32>
      Greedy loop using the cached step.

Profile (5 steady-state runs on /tmp/asr_test.flac, 10.42 s LibriSpeech):

  warm-up:                344 ms
  steady-state mean: 307 ms (p50 305, range 298-319)
  realtime factor: 0.0294x

Comparison across all STT in rtx-csm:

  Backend           RTF         Notes
  Kyutai STT 1B     1.01x       hardware-bound, 3 GB
  Whisper-tiny      0.020x      breaks CSM (in-process ggml conflict)
  Moonshine-tiny    0.0294x     pure candle, NO runtime conflict

Moonshine is the only fast STT path that integrates cleanly. ~34x
faster than realtime, ~17x faster than Kyutai 1B, no protobuf or
ggml linkage issues.

New `examples/moonshine_profile` mirrors `stt_profile` and
`whisper_profile` so all three STT backends report comparable numbers.

Phase 8.10 (next): wire as a third AsrEngine variant in converse_server
for English-only deploys. Replace the energy-VAD-gated Kyutai path
when --moonshine flag is set.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
2026-04-27 11:05:19 -07:00

99 lines
3.2 KiB
Rust

//! Phase 8.8 — end-to-end Moonshine transcription. Loads real audio,
//! runs encoder + greedy decode + tokenizer, prints the transcript.
//!
//! Usage:
//! ```bash
//! cargo run -p rtx-csm --release --features metal --example moonshine_transcribe -- \
//! --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, 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,
/// Maximum decoded tokens (Moonshine config caps at 194).
#[arg(long, default_value_t = 100)]
max_tokens: usize,
}
fn main() -> Result<()> {
let cli = Cli::parse();
let device = if candle_core::utils::metal_is_available() {
Device::new_metal(0)?
} else {
Device::Cpu
};
eprintln!("device: {device:?}");
// Download both safetensors and tokenizer.json.
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")?;
eprintln!("weights: {}", weights.display());
eprintln!("tokenizer: {}", tok_path.display());
// Audio: load + resample to 16 kHz.
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)?;
eprintln!("load: {:.2}s", load_t.elapsed().as_secs_f32());
let tok = moonshine::load_tokenizer(&tok_path)
.map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
eprintln!("tokenizer loaded (vocab {})", tok.get_vocab_size(true));
// Encode audio.
let pcm_t = Tensor::from_vec(pcm.clone(), (1, 1, pcm.len()), &device)?;
let enc_t = Instant::now();
let enc = encoder.forward(&pcm_t)?;
eprintln!(
"encode: {:.0} ms (output {:?})",
enc_t.elapsed().as_millis(),
enc.shape()
);
// Greedy decode (cached path — O(T) per token instead of O(T²)).
let dec_t = Instant::now();
let token_ids = decoder.generate_cached(&enc, &cfg, cli.max_tokens)?;
let dec_ms = dec_t.elapsed().as_millis();
eprintln!(
"decode: {dec_ms} ms ({} tokens, {:.1} ms/token, KV-cached)",
token_ids.len(),
dec_ms as f64 / token_ids.len().max(1) as f64
);
// Detokenize.
let text = tok
.decode(&token_ids, true)
.map_err(|e| anyhow::anyhow!("detok: {e}"))?;
println!();
println!("=== transcript ===");
println!("{text}");
println!();
println!("token ids: {token_ids:?}");
// Realtime factor.
let rtf = (enc_t.elapsed().as_secs_f64() + dec_ms as f64 / 1000.0) / audio_secs as f64;
println!();
println!(
"realtime factor: {rtf:.3}x ({:.1}s of compute / {:.1}s of audio)",
enc_t.elapsed().as_secs_f64() + dec_ms as f64 / 1000.0,
audio_secs
);
Ok(())
}