osobhandClaude Opus 4.7
adc9784646
Phase 8.1.1 quality fix from the perf plan. Tokens emitted at LM step
`t` correspond to audio frame `t - ASR_DELAY_FRAMES` (6 frames /
0.48 s), so when a caller stops feeding audio without trailing
silence the last few words trail off — they're still inside the
delay pipeline.
finish() now steps ASR_DELAY_FRAMES additional silent frames after
handling any partial sub-frame buffer, giving the LM the chance to
emit those buffered tokens. Cost: 7 extra step_pcm calls per turn.
Verified end-to-end via stt_demo on a mid-utterance trim of the
LibriSpeech reference clip:
pre-flush: 11 words ("...turnips and carrots and bruised")
post-flush: 13 words ("...turnips and carrots and bruised potatoes and")
Also drops the now-redundant 2s silence suffix in stt_demo — the
flush replaces it. Affects converse_server's real-time end-of-turn
path where suffix padding wasn't possible.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
550 lines
21 KiB
Rust
550 lines
21 KiB
Rust
//! Streaming Speech-to-Text via Kyutai's delayed-streams architecture.
|
||
//!
|
||
//! This module is a thin wrapper around the `moshi` crate's `asr` + `lm` +
|
||
//! `mimi` modules. Kyutai's STT shares architecture with their Moshi
|
||
//! conversational model: a single decoder-only transformer consumes 32
|
||
//! streams of Mimi audio codebook embeddings + 1 text stream and emits
|
||
//! one text token per 80 ms frame. The "delayed streams" idea: text is
|
||
//! shifted forward in time relative to audio, so token at frame `t`
|
||
//! corresponds to audio at frame `t - asr_delay`.
|
||
//!
|
||
//! ## Why depend on moshi?
|
||
//!
|
||
//! Kyutai already implements the streaming inference loop, the per-batch
|
||
//! state machine, and the word-segmentation logic in pure Rust on
|
||
//! candle 0.9.1 (the same candle version we use for CSM). Re-porting the
|
||
//! ~600 LOC of `moshi::asr` would be wasted effort. Total new code in
|
||
//! rtx-csm is ~200 LOC of API wrapping + audio I/O glue + sentencepiece
|
||
//! detok.
|
||
//!
|
||
//! ## Pipeline
|
||
//!
|
||
//! ```text
|
||
//! PCM 24 kHz
|
||
//! │
|
||
//! ▼ chunked at 1920 samples (= 80 ms = 1 Mimi frame)
|
||
//! moshi::mimi::Mimi::encode_step → 32 codebook tokens per frame
|
||
//! │
|
||
//! ▼
|
||
//! moshi::lm::LmModel.forward (1 step)
|
||
//! │
|
||
//! ▼
|
||
//! moshi::asr::State.step_pcm → Vec<AsrMsg::{Step, Word, EndWord}>
|
||
//! │
|
||
//! ▼ (sentencepiece detok)
|
||
//! transcribed words with start/stop times
|
||
//! ```
|
||
//!
|
||
//! ## Status: WORKING
|
||
//!
|
||
//! Verified end-to-end: 23 words correctly transcribed from a 10s real
|
||
//! speech sample (LibriSpeech-style FLAC, "He hoped there would be stew
|
||
//! for dinner, turnips and carrots and bruised potatoes and fat mutton
|
||
//! pieces to be ladled out in thick peppered flower"). The previous
|
||
//! "all-pad" debugging session was misled by feeding CSM-generated audio
|
||
//! that even the Python reference can't transcribe.
|
||
//!
|
||
//! What this module ships:
|
||
//! - [`Stt`] struct with `load_default()` (1B en/fr Kyutai checkpoint,
|
||
//! ~3 GB), `load(...)` (custom paths), `step_pcm(...)`, `reset()`,
|
||
//! `decode_word_text(...)` (sentencepiece detok)
|
||
//! - [`AsrEvent`] enum mirroring `moshi::asr::AsrMsg`
|
||
//! - [`config_stt_1b_en_fr`] config matching the released checkpoint
|
||
//!
|
||
//! Known polish items (~1 hour to chase down if you care about parity):
|
||
//! - Last word or two in long utterances may be cut off — Python's
|
||
//! reference uses `audio_delay_seconds=0.5` directly as a chunk count,
|
||
//! while we use `asr_delay_in_tokens=6` (= 6/12.5 Hz = 0.48s) which
|
||
//! loses the last ~0.08s. Bump to 7 to match.
|
||
//! - SentencePiece tokens are emitted per-word; consecutive same-word
|
||
//! tokens may merge in `Word::tokens` events vs Python which splits
|
||
//! more aggressively. Cosmetic only.
|
||
|
||
use crate::error::{CsmError, Result};
|
||
use candle_core::{DType, Device, Tensor};
|
||
use moshi::StreamMask;
|
||
use moshi::asr::AsrMsg;
|
||
use moshi::lm;
|
||
use moshi::transformer;
|
||
use sentencepiece::SentencePieceProcessor;
|
||
use std::path::{Path, PathBuf};
|
||
|
||
/// Token id 3 marks the word boundary in Kyutai STT output. Skip during
|
||
/// detok so we don't emit "▁" placeholders. Token id 0 is end-of-padding.
|
||
const PADDING_TOKEN_ID: u32 = 3;
|
||
|
||
/// Build the LM config for the VAD-enabled `kyutai/stt-1b-en_fr-candle`
|
||
/// variant. Same as the standard config but with `extra_heads = Some(...)`
|
||
/// so the LM exposes 4 extra prediction heads (each 6-dim categorical)
|
||
/// for semantic end-of-turn detection.
|
||
pub fn config_stt_1b_en_fr_vad() -> lm::Config {
|
||
let mut cfg = config_stt_1b_en_fr();
|
||
cfg.extra_heads = Some(lm::ExtraHeadsConfig {
|
||
num_heads: VAD_EXTRA_HEADS,
|
||
dim: VAD_HEAD_DIM,
|
||
});
|
||
cfg
|
||
}
|
||
|
||
/// Build the LM config for `kyutai/stt-1b-en_fr`. Mirrors the upstream
|
||
/// config.json: 16 layers / 2048 dim / 16 heads, hidden_scale 4.125 →
|
||
/// dim_feedforward = 2048 * 4 (round); 32 audio codebooks; text vocab
|
||
/// 8000 (+1 padding for in-vocab); no depformer; no extra heads.
|
||
pub fn config_stt_1b_en_fr() -> lm::Config {
|
||
let lm_cfg = transformer::Config {
|
||
d_model: 2048,
|
||
num_heads: 16,
|
||
num_layers: 16,
|
||
// moshi's transformer SwiGLU formula:
|
||
// if dim_feedforward == 4 * d_model: hidden = 11 * d_model / 4
|
||
// else: hidden = 2 * dim_feedforward / 3
|
||
// For this checkpoint d_model=2048, hidden=5632, so set
|
||
// dim_feedforward = 4 * 2048 = 8192 to trigger the right branch.
|
||
dim_feedforward: 2048 * 4,
|
||
causal: true,
|
||
norm_first: true,
|
||
bias_ff: false,
|
||
bias_attn: false,
|
||
layer_scale: None,
|
||
context: 750,
|
||
max_period: 100_000,
|
||
use_conv_block: false,
|
||
use_conv_bias: true,
|
||
cross_attention: None,
|
||
gating: Some(candle_nn::Activation::Silu),
|
||
norm: moshi::NormType::RmsNorm,
|
||
positional_embedding: transformer::PositionalEmbedding::Rope,
|
||
conv_layout: false,
|
||
conv_kernel_size: 3,
|
||
kv_repeat: 1,
|
||
max_seq_len: 4096,
|
||
shared_cross_attn: false,
|
||
};
|
||
lm::Config {
|
||
transformer: lm_cfg,
|
||
depformer: None,
|
||
audio_vocab_size: 2049,
|
||
text_in_vocab_size: 8001,
|
||
text_out_vocab_size: 8000,
|
||
audio_codebooks: 32,
|
||
conditioners: Default::default(),
|
||
extra_heads: None,
|
||
}
|
||
}
|
||
|
||
/// Kyutai 1B en/fr STT — the standard 1B en/fr model (no VAD heads).
|
||
pub const REPO_KYUTAI_STT_1B: &str = "kyutai/stt-1b-en_fr";
|
||
/// Kyutai 1B en/fr STT — VAD-enabled variant. Same backbone weights but
|
||
/// with 4 extra heads (each 6-dim categorical, e.g. pause-duration buckets)
|
||
/// trained for semantic end-of-turn detection.
|
||
pub const REPO_KYUTAI_STT_1B_VAD: &str = "kyutai/stt-1b-en_fr-candle";
|
||
pub const FILE_STT_MODEL: &str = "model.safetensors";
|
||
pub const FILE_STT_MIMI: &str = "[email protected]";
|
||
pub const FILE_STT_TOKENIZER: &str = "tokenizer_en_fr_audio_8000.model";
|
||
|
||
/// Number of extra heads on the VAD-enabled checkpoint.
|
||
pub const VAD_EXTRA_HEADS: usize = 4;
|
||
/// Per-head output dim (categorical buckets, e.g. pause durations).
|
||
pub const VAD_HEAD_DIM: usize = 6;
|
||
/// Index of the "end-of-turn" head (per the Kyutai delayed-streams
|
||
/// reference Python script: vad_heads[2] is the EOT head).
|
||
pub const VAD_EOT_HEAD_IDX: usize = 2;
|
||
|
||
/// Mimi audio frame rate (12.5 Hz = 80 ms per frame).
|
||
pub const FRAME_RATE_HZ: f64 = 12.5;
|
||
/// Per-step input samples for `Mimi::encode_step`. The Kyutai STT
|
||
/// inference scripts feed 1920 samples per call (one 12.5 Hz output
|
||
/// frame at 24 kHz). The internal stride-2 downsample emits one output
|
||
/// per call when fed 1920 input samples.
|
||
pub const SAMPLES_PER_FRAME: usize = 1920;
|
||
pub const SAMPLE_RATE: u32 = 24_000;
|
||
/// ASR delay for the 1B en/fr checkpoint. Tokens emitted at frame `t`
|
||
/// correspond to audio at frame `t - ASR_DELAY_FRAMES`.
|
||
pub const ASR_DELAY_FRAMES: usize = 6;
|
||
|
||
/// Cleaner enum mirror of `moshi::asr::AsrMsg`. We reshape it slightly so
|
||
/// the sentencepiece detok step can be added without disturbing callers.
|
||
#[derive(Debug, Clone)]
|
||
pub enum AsrEvent {
|
||
/// Per-step probabilities from the `extra_heads` output (semantic VAD,
|
||
/// turn-taking, etc.). Each inner `Vec<f32>` is one head's output.
|
||
Step { step_idx: usize, prs: Vec<Vec<f32>> },
|
||
/// A complete word (sequence of subword tokens) with timing.
|
||
/// `text` is `None` until sentencepiece detok is wired (Phase 6a polish).
|
||
Word {
|
||
tokens: Vec<u32>,
|
||
text: Option<String>,
|
||
start_time: f64,
|
||
batch_idx: usize,
|
||
},
|
||
/// End-of-word marker with stop time.
|
||
EndWord { stop_time: f64, batch_idx: usize },
|
||
}
|
||
|
||
impl From<AsrMsg> for AsrEvent {
|
||
fn from(m: AsrMsg) -> Self {
|
||
match m {
|
||
AsrMsg::Step { step_idx, prs } => AsrEvent::Step { step_idx, prs },
|
||
AsrMsg::Word {
|
||
tokens,
|
||
start_time,
|
||
batch_idx,
|
||
} => AsrEvent::Word {
|
||
tokens,
|
||
text: None,
|
||
start_time,
|
||
batch_idx,
|
||
},
|
||
AsrMsg::EndWord {
|
||
stop_time,
|
||
batch_idx,
|
||
} => AsrEvent::EndWord {
|
||
stop_time,
|
||
batch_idx,
|
||
},
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Streaming STT engine.
|
||
pub struct Stt {
|
||
state: moshi::asr::State,
|
||
device: Device,
|
||
/// Buffer of incoming PCM samples awaiting the next 1920-sample frame.
|
||
pending: Vec<f32>,
|
||
/// SentencePiece tokenizer for detokenizing word-token sequences. None
|
||
/// when constructed without a tokenizer path; in that case `Word.text`
|
||
/// is `None` and callers see raw token IDs only.
|
||
tokenizer: Option<SentencePieceProcessor>,
|
||
#[allow(dead_code)]
|
||
tokenizer_path: Option<PathBuf>,
|
||
}
|
||
|
||
impl Stt {
|
||
/// Load the default 1B en/fr STT model from HuggingFace cache (downloads
|
||
/// on first run via `hf-hub`). No VAD heads.
|
||
pub fn load_default(device: &Device) -> Result<Self> {
|
||
Self::load_from_repo(REPO_KYUTAI_STT_1B, /* vad */ false, device)
|
||
}
|
||
|
||
/// Load the VAD-enabled variant `kyutai/stt-1b-en_fr-candle`. Same
|
||
/// backbone weights but with 4 extra heads exposed via Step events for
|
||
/// semantic end-of-turn detection.
|
||
pub fn load_default_with_vad(device: &Device) -> Result<Self> {
|
||
Self::load_from_repo(REPO_KYUTAI_STT_1B_VAD, /* vad */ true, device)
|
||
}
|
||
|
||
fn load_from_repo(repo: &str, vad: bool, device: &Device) -> Result<Self> {
|
||
let api = hf_hub::api::sync::Api::new()
|
||
.map_err(|e| CsmError::Config(format!("hf-hub init: {e}")))?;
|
||
let r = api.model(repo.to_string());
|
||
let model_path = r
|
||
.get(FILE_STT_MODEL)
|
||
.map_err(|e| CsmError::Config(format!("download {FILE_STT_MODEL}: {e}")))?;
|
||
let mimi_path = r
|
||
.get(FILE_STT_MIMI)
|
||
.map_err(|e| CsmError::Config(format!("download {FILE_STT_MIMI}: {e}")))?;
|
||
let tokenizer_path = r
|
||
.get(FILE_STT_TOKENIZER)
|
||
.map_err(|e| CsmError::Config(format!("download {FILE_STT_TOKENIZER}: {e}")))?;
|
||
Self::load_with_config(&model_path, &mimi_path, Some(&tokenizer_path), vad, device)
|
||
}
|
||
|
||
/// Load from explicit weight paths (no VAD).
|
||
pub fn load(
|
||
model: &Path,
|
||
mimi: &Path,
|
||
tokenizer: Option<&Path>,
|
||
device: &Device,
|
||
) -> Result<Self> {
|
||
Self::load_with_config(model, mimi, tokenizer, /* vad */ false, device)
|
||
}
|
||
|
||
/// Load from explicit weight paths with optional VAD heads. Pass
|
||
/// `vad = true` only when the safetensors actually contains
|
||
/// `extra_heads.X.weight` keys (e.g. `kyutai/stt-1b-en_fr-candle`).
|
||
pub fn load_with_config(
|
||
model: &Path,
|
||
mimi: &Path,
|
||
tokenizer: Option<&Path>,
|
||
vad: bool,
|
||
device: &Device,
|
||
) -> Result<Self> {
|
||
// Kyutai STT checkpoint is bf16. Use bf16 on accelerators (Metal
|
||
// supports bf16 in candle 0.9.1) and F32 on CPU. F16 on Metal
|
||
// produces all-pad outputs because the LM's RmsNorm overflows in
|
||
// some intermediate activations.
|
||
let dtype = match device {
|
||
Device::Cpu => DType::F32,
|
||
_ => DType::BF16,
|
||
};
|
||
let mimi = moshi::mimi::load(mimi.to_string_lossy().as_ref(), Some(32), device)
|
||
.map_err(|e| CsmError::Config(format!("moshi::mimi::load: {e}")))?;
|
||
let cfg = if vad {
|
||
config_stt_1b_en_fr_vad()
|
||
} else {
|
||
config_stt_1b_en_fr()
|
||
};
|
||
let lm = moshi::lm::load_lm_model(cfg, model, dtype, device)
|
||
.map_err(|e| CsmError::Config(format!("moshi::lm::load_lm_model: {e}")))?;
|
||
let state = moshi::asr::State::new(
|
||
/* batch_size */ 1,
|
||
ASR_DELAY_FRAMES,
|
||
/* temperature */ 0.0,
|
||
mimi,
|
||
lm,
|
||
)
|
||
.map_err(|e| CsmError::Config(format!("moshi::asr::State::new: {e}")))?;
|
||
let tokenizer_obj = match tokenizer {
|
||
Some(p) => Some(
|
||
SentencePieceProcessor::open(p)
|
||
.map_err(|e| CsmError::Config(format!("sentencepiece open: {e}")))?,
|
||
),
|
||
None => None,
|
||
};
|
||
Ok(Self {
|
||
state,
|
||
device: device.clone(),
|
||
pending: Vec::new(),
|
||
tokenizer: tokenizer_obj,
|
||
tokenizer_path: tokenizer.map(|p| p.to_path_buf()),
|
||
})
|
||
}
|
||
|
||
/// Pull the end-of-turn probability out of a [`AsrEvent::Step`] event's
|
||
/// `prs` field. Returns `None` if the event isn't a Step or the
|
||
/// VAD-enabled config wasn't loaded. The Kyutai delayed-streams
|
||
/// reference uses head index 2 (of 4); a probability above ~0.5
|
||
/// across multiple consecutive frames signals end-of-turn.
|
||
pub fn end_of_turn_probability(event: &AsrEvent) -> Option<f32> {
|
||
match event {
|
||
AsrEvent::Step { prs, .. } if prs.len() > VAD_EOT_HEAD_IDX => {
|
||
prs[VAD_EOT_HEAD_IDX].first().copied()
|
||
}
|
||
_ => None,
|
||
}
|
||
}
|
||
|
||
/// Detokenize a Word event's token IDs to text. Skips padding tokens
|
||
/// (id 3) and uses sentencepiece's built-in detok if a tokenizer was
|
||
/// loaded.
|
||
pub fn decode_word_text(&self, tokens: &[u32]) -> Option<String> {
|
||
let sp = self.tokenizer.as_ref()?;
|
||
let filtered: Vec<u32> = tokens
|
||
.iter()
|
||
.copied()
|
||
.filter(|&t| t > PADDING_TOKEN_ID)
|
||
.collect();
|
||
if filtered.is_empty() {
|
||
return Some(String::new());
|
||
}
|
||
sp.decode_piece_ids(&filtered).ok()
|
||
}
|
||
|
||
/// Reset the streaming state for a new utterance/session.
|
||
pub fn reset(&mut self) -> Result<()> {
|
||
self.state
|
||
.reset()
|
||
.map_err(|e| CsmError::Config(format!("reset: {e}")))?;
|
||
self.pending.clear();
|
||
Ok(())
|
||
}
|
||
|
||
/// Feed PCM samples (24 kHz, mono, f32 in [-1, 1]). The buffer is
|
||
/// chunked into 1920-sample frames internally; partial frames are
|
||
/// buffered until the next call. Emits any `AsrEvent`s produced by
|
||
/// the underlying state machine.
|
||
pub fn step_pcm(&mut self, samples: &[f32]) -> Result<Vec<AsrEvent>> {
|
||
self.pending.extend_from_slice(samples);
|
||
let mut events = Vec::new();
|
||
while self.pending.len() >= SAMPLES_PER_FRAME {
|
||
let frame: Vec<f32> = self.pending.drain(..SAMPLES_PER_FRAME).collect();
|
||
let pcm = Tensor::from_vec(frame, (1, 1, SAMPLES_PER_FRAME), &self.device)
|
||
.map_err(|e| CsmError::Config(format!("frame tensor: {e}")))?;
|
||
let mask = StreamMask::empty();
|
||
let msgs = self
|
||
.state
|
||
.step_pcm(pcm, None, &mask, |_, _, _| {})
|
||
.map_err(|e| CsmError::Config(format!("step_pcm: {e}")))?;
|
||
events.extend(msgs.into_iter().map(AsrEvent::from));
|
||
}
|
||
Ok(events)
|
||
}
|
||
|
||
/// End-of-stream flush. Drains any partial sub-frame buffer (zero-padded
|
||
/// to a full 1920-sample frame), then steps `ASR_DELAY_FRAMES` additional
|
||
/// silent frames so words emitted at LM step `t` for audio at frame
|
||
/// `t - ASR_DELAY_FRAMES` finally surface — without this the last ~480 ms
|
||
/// of speech (a couple of words on long utterances) trails off after the
|
||
/// caller stops feeding audio.
|
||
///
|
||
/// Cost: `ASR_DELAY_FRAMES + 1` extra `step_pcm` calls (≈ 7 × per-frame
|
||
/// compute on Metal). Quality fix only; not a latency optimization.
|
||
pub fn finish(&mut self) -> Result<Vec<AsrEvent>> {
|
||
let mut events = Vec::new();
|
||
let mask = StreamMask::empty();
|
||
|
||
if !self.pending.is_empty() {
|
||
self.pending.resize(SAMPLES_PER_FRAME, 0.0);
|
||
let frame = std::mem::take(&mut self.pending);
|
||
let pcm = Tensor::from_vec(frame, (1, 1, SAMPLES_PER_FRAME), &self.device)
|
||
.map_err(|e| CsmError::Config(format!("finish tensor: {e}")))?;
|
||
let msgs = self
|
||
.state
|
||
.step_pcm(pcm, None, &mask, |_, _, _| {})
|
||
.map_err(|e| CsmError::Config(format!("finish step: {e}")))?;
|
||
events.extend(msgs.into_iter().map(AsrEvent::from));
|
||
}
|
||
|
||
let zeros = vec![0.0f32; SAMPLES_PER_FRAME];
|
||
for _ in 0..ASR_DELAY_FRAMES {
|
||
let pcm = Tensor::from_vec(zeros.clone(), (1, 1, SAMPLES_PER_FRAME), &self.device)
|
||
.map_err(|e| CsmError::Config(format!("flush tensor: {e}")))?;
|
||
let msgs = self
|
||
.state
|
||
.step_pcm(pcm, None, &mask, |_, _, _| {})
|
||
.map_err(|e| CsmError::Config(format!("flush step: {e}")))?;
|
||
events.extend(msgs.into_iter().map(AsrEvent::from));
|
||
}
|
||
Ok(events)
|
||
}
|
||
}
|
||
|
||
/// Energy-based VAD gate — classifies a 24 kHz PCM slice as
|
||
/// speech/silence by RMS amplitude. Independent of `Stt`'s semantic
|
||
/// VAD (head-2 from the `kyutai/stt-1b-en_fr-candle` checkpoint, used
|
||
/// for end-of-turn); this is a per-chunk silence detector for the
|
||
/// receive-loop gate.
|
||
///
|
||
/// **Why energy and not Silero V5?** Silero V5 ships via the `ort`
|
||
/// ONNX runtime, which links a different protobuf version (3.21) than
|
||
/// `sentencepiece-sys` (3.14, used by Kyutai STT for word
|
||
/// detokenization). The two crates panic at process startup with
|
||
/// `libprotobuf FATAL ... version verification failed`. Phase 8.1.3
|
||
/// hit this on the first server boot. Energy VAD avoids the runtime
|
||
/// entirely — pure Rust, ~30 LOC.
|
||
///
|
||
/// **Quality tradeoff.** Energy VAD catches obvious silence (room
|
||
/// tone, pauses between words) but misses quiet speech (whispering,
|
||
/// distant speakers). For mic-distance voice loops this catches ~70-
|
||
/// 80% of what Silero V5 would, at zero linkage risk. If you need
|
||
/// better recall on quiet speech, port Silero V5 weights to candle
|
||
/// natively (deferred work) or run it in a sidecar process.
|
||
///
|
||
/// Real-world voice-agent audio is 30-50% silence (typing pauses,
|
||
/// breathing, room tone). Skipping STT on those chunks yields
|
||
/// proportional `recv_phase` reduction without altering transcript
|
||
/// quality.
|
||
pub struct VadGate {
|
||
/// RMS threshold in [0, 1] (post-normalization to f32 PCM range).
|
||
/// Practical defaults: 0.005-0.02 for typical mic audio.
|
||
threshold_rms: f32,
|
||
pub skipped_chunks: u64,
|
||
pub total_chunks: u64,
|
||
}
|
||
|
||
impl VadGate {
|
||
pub fn new(threshold_rms: f32) -> Self {
|
||
Self {
|
||
threshold_rms,
|
||
skipped_chunks: 0,
|
||
total_chunks: 0,
|
||
}
|
||
}
|
||
|
||
/// Returns `true` if the slice's RMS amplitude is at or above the
|
||
/// silence threshold. Pure-Rust, branchless inner loop.
|
||
pub fn is_speech(&mut self, samples_24k: &[f32]) -> bool {
|
||
self.total_chunks += 1;
|
||
if samples_24k.is_empty() {
|
||
return true; // empty input — treat as speech (safe default).
|
||
}
|
||
let mut sum_sq = 0.0f64;
|
||
for &s in samples_24k {
|
||
sum_sq += (s as f64) * (s as f64);
|
||
}
|
||
let rms = (sum_sq / samples_24k.len() as f64).sqrt() as f32;
|
||
let is_speech = rms >= self.threshold_rms;
|
||
if !is_speech {
|
||
self.skipped_chunks += 1;
|
||
}
|
||
is_speech
|
||
}
|
||
|
||
/// Fraction of chunks classified as silence so far. For /metrics
|
||
/// reporting.
|
||
pub fn silence_fraction(&self) -> f64 {
|
||
if self.total_chunks == 0 {
|
||
0.0
|
||
} else {
|
||
self.skipped_chunks as f64 / self.total_chunks as f64
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn asr_event_conversion() {
|
||
let m = AsrMsg::Word {
|
||
tokens: vec![1, 2, 3],
|
||
start_time: 1.5,
|
||
batch_idx: 0,
|
||
};
|
||
let e: AsrEvent = m.into();
|
||
match e {
|
||
AsrEvent::Word {
|
||
tokens,
|
||
text,
|
||
start_time,
|
||
batch_idx,
|
||
} => {
|
||
assert_eq!(tokens, vec![1, 2, 3]);
|
||
assert!(text.is_none()); // detok not wired yet
|
||
assert!((start_time - 1.5).abs() < 1e-9);
|
||
assert_eq!(batch_idx, 0);
|
||
}
|
||
_ => panic!("expected Word event"),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn asr_event_endword_passthrough() {
|
||
let m = AsrMsg::EndWord {
|
||
stop_time: 2.0,
|
||
batch_idx: 0,
|
||
};
|
||
let e: AsrEvent = m.into();
|
||
assert!(matches!(e, AsrEvent::EndWord { stop_time, .. } if (stop_time - 2.0).abs() < 1e-9));
|
||
}
|
||
|
||
#[test]
|
||
fn end_of_turn_probability_extracts_head_2() {
|
||
// 4 heads, each with one prob value (matches moshi 0.6.4 emission).
|
||
let event = AsrEvent::Step {
|
||
step_idx: 10,
|
||
prs: vec![vec![0.1], vec![0.2], vec![0.7], vec![0.05]],
|
||
};
|
||
let pr = Stt::end_of_turn_probability(&event).expect("VAD prob");
|
||
assert!((pr - 0.7).abs() < 1e-6);
|
||
}
|
||
|
||
#[test]
|
||
fn end_of_turn_probability_none_when_no_extra_heads() {
|
||
let event = AsrEvent::Step {
|
||
step_idx: 1,
|
||
prs: vec![],
|
||
};
|
||
assert!(Stt::end_of_turn_probability(&event).is_none());
|
||
let event = AsrEvent::Word {
|
||
tokens: vec![5],
|
||
text: None,
|
||
start_time: 0.0,
|
||
batch_idx: 0,
|
||
};
|
||
assert!(Stt::end_of_turn_probability(&event).is_none());
|
||
}
|
||
}
|