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]>
99 lines
3.2 KiB
Rust
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(())
|
|
}
|