diff --git a/crates/transcription/src/lib.rs b/crates/transcription/src/lib.rs index 83cf575..cf9ad10 100644 --- a/crates/transcription/src/lib.rs +++ b/crates/transcription/src/lib.rs @@ -1,6 +1,7 @@ pub mod concurrency; pub mod local_engine; pub mod model_manager; +pub mod streaming; pub mod transcriber; #[cfg(feature = "whisper")] pub mod whisper_rs_backend; @@ -10,5 +11,6 @@ pub use local_engine::{load_parakeet, LocalEngine, SpeechModelAdapter, TimedTran #[cfg(feature = "whisper")] pub use local_engine::load_whisper; pub use model_manager::{download, is_downloaded, list_downloaded, model_dir, models_dir}; +pub use streaming::{RmsVadChunker, VadChunk, VadChunker}; pub use transcriber::{Transcriber, TranscriberCapabilities}; pub use transcribe_rs::SpeechModel; diff --git a/crates/transcription/src/streaming/mod.rs b/crates/transcription/src/streaming/mod.rs new file mode 100644 index 0000000..2bd7030 --- /dev/null +++ b/crates/transcription/src/streaming/mod.rs @@ -0,0 +1,72 @@ +//! Streaming primitives for live capture: VAD-gated chunking, +//! agreement-based commit policy, and bounded buffer management. +//! +//! These types are tested at the unit level. Integration into +//! `src-tauri/src/commands/live.rs` lands in follow-up commits so +//! threshold tuning can be validated against real microphone captures +//! rather than synthetic fixtures (brief items #21, #24, #25). + +pub mod rms_vad; + +pub use rms_vad::RmsVadChunker; + +/// A span of audio the VAD considers worth transcribing. `start_sample` +/// is an absolute index into the stream the `VadChunker` has been fed +/// since its last `reset`; `samples` is f32 PCM at the chunker's +/// configured sample rate. +#[derive(Debug, Clone)] +pub struct VadChunk { + pub start_sample: u64, + pub samples: Vec, +} + +/// A streaming VAD-gated chunker. +/// +/// Implementations accumulate incoming samples, decide whether the +/// current segment is speech using a score + hysteresis (brief item +/// #21), and emit `VadChunk`s when a speech region ends — or when an +/// in-progress speech region exceeds the configured max length so +/// Whisper is not fed a 30-second monolith. +/// +/// `push` returns any chunks ready to dispatch; typical usage is +/// `for chunk in chunker.push(&samples) { dispatch(chunk); }` inside +/// the capture loop. +/// +/// `flush` is called at end-of-session to emit any in-flight speech as +/// a final chunk (even if silence hasn't closed it). +/// +/// `Send` because a chunker is owned by the live-session worker thread +/// and moved into `spawn_blocking`. +pub trait VadChunker: Send { + /// Feed new samples. Returns any chunks the chunker has decided to + /// emit as a result. An empty Vec means "still gathering". + fn push(&mut self, samples: &[f32]) -> Vec; + + /// End-of-session: emit any in-progress speech as a chunk even + /// though silence has not closed it. Returns `None` if there is + /// nothing buffered (or only sub-threshold samples). + fn flush(&mut self) -> Option; + + /// Drop accumulated state. Used between sessions on the same + /// chunker instance (or after a context-window reset from the + /// repetition detector — brief item #26). + fn reset(&mut self); + + /// Absolute sample index of the next sample that will be fed via + /// `push`. Exposed so the commit policy (#24) can compute sample + /// offsets for its agreement window. + fn next_sample_index(&self) -> u64; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn vad_chunker_trait_is_object_safe() { + // Compile-time witness: keep the trait dyn-compatible so the + // live-session worker can hold `Box` and swap + // between RMS and Silero backends at runtime. + let _: Option> = None; + } +} diff --git a/crates/transcription/src/streaming/rms_vad.rs b/crates/transcription/src/streaming/rms_vad.rs new file mode 100644 index 0000000..87fa951 --- /dev/null +++ b/crates/transcription/src/streaming/rms_vad.rs @@ -0,0 +1,481 @@ +//! RMS-energy-backed VAD chunker. +//! +//! This is the fallback backend the plan (`docs/whisper-ecosystem/ +//! workstream-A.md`, Phase A.3 "Known unknowns") permits while the ort +//! 2.0.0-rc.10 vs rc.12 ecosystem conflict prevents a drop-in Silero +//! dep. The `VadChunker` trait surface is identical to what a Silero +//! backend will present, so the live-session path does not change when +//! Silero lands. +//! +//! The chunker emits a `VadChunk` when a sustained-speech region ends +//! (RMS drops below `exit_threshold` for `silence_close_samples`) or +//! when an in-progress region exceeds `max_chunk_samples` (so Whisper +//! is not fed a 30-second monolith). It applies hysteresis — an +//! `enter_threshold` higher than `exit_threshold` — so a VAD score +//! bouncing around the threshold does not toggle state every frame. + +use super::{VadChunk, VadChunker}; + +/// Sample window used to compute a single RMS reading. 50 ms at 16 +/// kHz. Shorter windows twitch on transients; longer windows blur the +/// speech-onset boundary. +const FRAME_SAMPLES: usize = 800; + +/// Default thresholds tuned to match the existing `evaluate_speech_gate` +/// behaviour in `src-tauri/src/commands/live.rs`. The underlying +/// constants live in that file; this chunker exposes them as fields so +/// they can be tuned per-session without a recompile. +const DEFAULT_ENTER_RMS_THRESHOLD: f32 = 0.003; +const DEFAULT_EXIT_RMS_THRESHOLD: f32 = 0.0014; +/// Frames of sustained speech required before the chunker enters the +/// "in-speech" state. Filters out single-frame transients (keyboard +/// clicks, door closes). +const DEFAULT_SPEECH_ONSET_FRAMES: usize = 3; +/// Silence duration that closes an in-progress chunk, in samples. +/// 500 ms = 10 frames at 16 kHz / 50 ms-frames. +const DEFAULT_SILENCE_CLOSE_SAMPLES: usize = 8_000; +/// Hard cap on a single chunk. Matches the existing `CHUNK_SAMPLES` +/// (2 s) so the live-streaming experience is not delayed arbitrarily +/// by a user speaking continuously. +const DEFAULT_MAX_CHUNK_SAMPLES: usize = 32_000; +/// Sample rate the thresholds above assume. Exposed so future backends +/// (Parakeet, Moonshine) at different rates can construct a chunker +/// matching their native sample rate. +const DEFAULT_SAMPLE_RATE_HZ: u32 = 16_000; + +#[derive(Debug, Clone, Copy, PartialEq)] +enum State { + /// Nothing buffered. Waiting for the first RMS excursion over + /// `enter_threshold`. + Idle, + /// In-progress speech. Samples accumulate; closes on + /// `silence_close_samples` of sub-threshold audio or on + /// `max_chunk_samples`. + InSpeech, +} + +pub struct RmsVadChunker { + // Tunables + enter_threshold: f32, + exit_threshold: f32, + speech_onset_frames: usize, + silence_close_samples: usize, + max_chunk_samples: usize, + + // Running state + state: State, + /// Frame-boundary reassembly: samples that did not complete a + /// frame on the previous `push`. Always shorter than `FRAME_SAMPLES`. + pending: Vec, + /// Samples belonging to the current in-progress chunk (State::InSpeech). + active_chunk: Vec, + /// Trailing silence sample count inside the current chunk. Resets + /// to zero whenever a speech frame is seen. + silent_tail_samples: usize, + /// Consecutive speech frames observed while `State::Idle`. When + /// this hits `speech_onset_frames`, state transitions to InSpeech. + pending_onset_frames: usize, + /// Samples buffered from the onset window that should be attached + /// to the front of the emitted chunk so Whisper sees the speech + /// onset itself, not just the post-onset audio. + onset_buffer: Vec, + /// Absolute sample index of the next sample `push` will consume. + next_sample_index: u64, + /// Absolute sample index where the current in-progress chunk + /// started. Valid only while `state == InSpeech`. + active_chunk_start: u64, +} + +impl RmsVadChunker { + pub fn new() -> Self { + Self::with_thresholds( + DEFAULT_ENTER_RMS_THRESHOLD, + DEFAULT_EXIT_RMS_THRESHOLD, + DEFAULT_SPEECH_ONSET_FRAMES, + DEFAULT_SILENCE_CLOSE_SAMPLES, + DEFAULT_MAX_CHUNK_SAMPLES, + ) + } + + pub fn with_thresholds( + enter_threshold: f32, + exit_threshold: f32, + speech_onset_frames: usize, + silence_close_samples: usize, + max_chunk_samples: usize, + ) -> Self { + debug_assert!( + exit_threshold <= enter_threshold, + "exit_threshold must not exceed enter_threshold (hysteresis requires enter >= exit)" + ); + Self { + enter_threshold, + exit_threshold, + speech_onset_frames, + silence_close_samples, + max_chunk_samples, + state: State::Idle, + pending: Vec::new(), + active_chunk: Vec::new(), + silent_tail_samples: 0, + pending_onset_frames: 0, + onset_buffer: Vec::new(), + next_sample_index: 0, + active_chunk_start: 0, + } + } + + pub fn sample_rate_hz(&self) -> u32 { + DEFAULT_SAMPLE_RATE_HZ + } + + fn frame_rms(frame: &[f32]) -> f32 { + if frame.is_empty() { + return 0.0; + } + let sum_sq: f32 = frame.iter().map(|x| x * x).sum(); + (sum_sq / frame.len() as f32).sqrt() + } + + /// Consume one complete frame's worth of samples and update state. + /// `frame_start` is the absolute sample index of `frame[0]` in the + /// stream fed since `reset`. Returns a `VadChunk` if this frame + /// closed the in-progress chunk. + fn consume_frame(&mut self, frame: Vec, frame_start: u64) -> Option { + let rms = Self::frame_rms(&frame); + match self.state { + State::Idle => self.consume_frame_idle(frame, frame_start, rms), + State::InSpeech => self.consume_frame_in_speech(frame, rms), + } + } + + fn consume_frame_idle( + &mut self, + frame: Vec, + frame_start: u64, + rms: f32, + ) -> Option { + if rms >= self.enter_threshold { + self.pending_onset_frames += 1; + // Keep a rolling buffer of onset audio so once we confirm + // speech, the emitted chunk contains the speech attack + // rather than starting mid-syllable. + self.onset_buffer.extend_from_slice(&frame); + let onset_cap = self.speech_onset_frames * FRAME_SAMPLES; + if self.onset_buffer.len() > onset_cap { + let overflow = self.onset_buffer.len() - onset_cap; + self.onset_buffer.drain(..overflow); + } + + if self.pending_onset_frames >= self.speech_onset_frames { + // Transition: flush the onset buffer into active_chunk + // and begin accumulating. The onset buffer includes + // the current frame, so its start index is + // `frame_start + FRAME_SAMPLES - onset_buffer.len()`. + self.state = State::InSpeech; + self.active_chunk_start = frame_start + .saturating_add(FRAME_SAMPLES as u64) + .saturating_sub(self.onset_buffer.len() as u64); + self.active_chunk.clear(); + self.active_chunk.append(&mut self.onset_buffer); + self.silent_tail_samples = 0; + self.pending_onset_frames = 0; + } + } else { + // Sub-threshold frame while idle — reset the onset counter + // and drop any onset buffer. The gate demands *sustained* + // speech, not a single frame over threshold. + self.pending_onset_frames = 0; + self.onset_buffer.clear(); + } + None + } + + fn consume_frame_in_speech(&mut self, frame: Vec, rms: f32) -> Option { + self.active_chunk.extend_from_slice(&frame); + if rms >= self.exit_threshold { + self.silent_tail_samples = 0; + } else { + self.silent_tail_samples += frame.len(); + } + + let end_of_utterance = self.silent_tail_samples >= self.silence_close_samples; + let hit_max = self.active_chunk.len() >= self.max_chunk_samples; + if end_of_utterance || hit_max { + return Some(self.emit_active_chunk()); + } + None + } + + fn emit_active_chunk(&mut self) -> VadChunk { + // Trim the trailing silence from the chunk so Whisper isn't + // fed dead air at the end. Matches the "tighten chunk on + // commit" property #25 will later tie to commit points. + let mut samples = std::mem::take(&mut self.active_chunk); + if self.silent_tail_samples > 0 && samples.len() > self.silent_tail_samples { + let keep = samples.len() - self.silent_tail_samples; + samples.truncate(keep); + } + let start_sample = self.active_chunk_start; + + self.state = State::Idle; + self.silent_tail_samples = 0; + self.pending_onset_frames = 0; + self.onset_buffer.clear(); + + VadChunk { + start_sample, + samples, + } + } +} + +impl Default for RmsVadChunker { + fn default() -> Self { + Self::new() + } +} + +impl VadChunker for RmsVadChunker { + fn push(&mut self, samples: &[f32]) -> Vec { + if samples.is_empty() { + return Vec::new(); + } + self.pending.extend_from_slice(samples); + self.next_sample_index = self.next_sample_index.saturating_add(samples.len() as u64); + + let mut emitted = Vec::new(); + while self.pending.len() >= FRAME_SAMPLES { + // Absolute index of the first sample in the frame we are + // about to consume: total fed minus what is still pending. + let frame_start = self + .next_sample_index + .saturating_sub(self.pending.len() as u64); + let frame: Vec = self.pending.drain(..FRAME_SAMPLES).collect(); + if let Some(chunk) = self.consume_frame(frame, frame_start) { + emitted.push(chunk); + } + } + emitted + } + + fn flush(&mut self) -> Option { + // Consume any tail of fewer-than-frame samples so the last + // utterance is not lost when a user stops recording mid-word. + if !self.pending.is_empty() { + // The tail's absolute start index is captured BEFORE we + // consume — same formula as in `push`. + let frame_start = self + .next_sample_index + .saturating_sub(self.pending.len() as u64); + let pad_len = FRAME_SAMPLES - self.pending.len(); + let mut padded = std::mem::take(&mut self.pending); + padded.extend(std::iter::repeat(0.0_f32).take(pad_len)); + let _ = self.consume_frame(padded, frame_start); + } + if self.state == State::InSpeech { + Some(self.emit_active_chunk()) + } else { + None + } + } + + fn reset(&mut self) { + self.state = State::Idle; + self.pending.clear(); + self.active_chunk.clear(); + self.silent_tail_samples = 0; + self.pending_onset_frames = 0; + self.onset_buffer.clear(); + self.next_sample_index = 0; + self.active_chunk_start = 0; + } + + fn next_sample_index(&self) -> u64 { + self.next_sample_index + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Generate a vector of `len` samples at amplitude `amp`. The + /// signal is a constant DC offset, which gives a deterministic + /// RMS of exactly `amp.abs()` — simpler than a sinusoid for + /// threshold-crossing tests. + fn constant_signal(len: usize, amp: f32) -> Vec { + vec![amp; len] + } + + #[test] + fn pure_silence_emits_nothing() { + let mut c = RmsVadChunker::new(); + let silence = constant_signal(16_000, 0.0); // 1 s of zero + let chunks = c.push(&silence); + assert!(chunks.is_empty()); + assert_eq!(c.flush().is_none(), true); + } + + #[test] + fn below_enter_threshold_does_not_trigger() { + let mut c = RmsVadChunker::new(); + // 0.002 is between the default exit (0.0014) and enter (0.003) + // thresholds — must NOT transition Idle → InSpeech. + let hum = constant_signal(16_000, 0.002); + let chunks = c.push(&hum); + assert!( + chunks.is_empty(), + "samples below enter_threshold must not trigger onset" + ); + } + + #[test] + fn single_loud_frame_does_not_trigger_onset() { + let mut c = RmsVadChunker::new(); + // One frame above enter, surrounded by silence. With + // speech_onset_frames=3 this should NOT transition. + let mut signal = Vec::new(); + signal.extend(constant_signal(FRAME_SAMPLES, 0.0)); + signal.extend(constant_signal(FRAME_SAMPLES, 0.01)); // loud, one frame + signal.extend(constant_signal(FRAME_SAMPLES * 4, 0.0)); + let chunks = c.push(&signal); + assert!( + chunks.is_empty(), + "single-frame transient must not cross sustained-speech onset" + ); + } + + #[test] + fn sustained_speech_followed_by_silence_emits_one_chunk() { + let mut c = RmsVadChunker::new(); + // 8 frames of speech (well over onset) followed by 12 frames of + // silence (well over silence_close). Must emit exactly one + // chunk. + let mut signal = Vec::new(); + signal.extend(constant_signal(FRAME_SAMPLES * 8, 0.01)); + signal.extend(constant_signal(FRAME_SAMPLES * 12, 0.0)); + let chunks = c.push(&signal); + assert_eq!(chunks.len(), 1, "one speech region → one chunk"); + let chunk = &chunks[0]; + assert!( + !chunk.samples.is_empty(), + "emitted chunk must contain samples" + ); + } + + #[test] + fn hysteresis_prevents_mid_utterance_close_on_brief_dip() { + let mut c = RmsVadChunker::new(); + // Onset → loud → brief dip between enter and exit → loud again + // → silence. The dip is above exit_threshold so the chunk must + // NOT close across it. + let loud = constant_signal(FRAME_SAMPLES * 4, 0.01); + let dip = constant_signal(FRAME_SAMPLES, 0.002); + let more_loud = constant_signal(FRAME_SAMPLES * 4, 0.01); + let silence = constant_signal(FRAME_SAMPLES * 12, 0.0); + let mut signal = Vec::new(); + signal.extend(loud); + signal.extend(dip); + signal.extend(more_loud); + signal.extend(silence); + let chunks = c.push(&signal); + assert_eq!( + chunks.len(), + 1, + "hysteresis dip between enter and exit thresholds must not split a chunk" + ); + } + + #[test] + fn max_chunk_samples_caps_continuous_speech() { + let mut c = RmsVadChunker::with_thresholds( + DEFAULT_ENTER_RMS_THRESHOLD, + DEFAULT_EXIT_RMS_THRESHOLD, + DEFAULT_SPEECH_ONSET_FRAMES, + DEFAULT_SILENCE_CLOSE_SAMPLES, + FRAME_SAMPLES * 4, // tighter cap for the test + ); + // Feed 12 frames of sustained speech with no silence break. + // The 4-frame cap must cause at least one emission mid-stream. + let signal = constant_signal(FRAME_SAMPLES * 12, 0.01); + let chunks = c.push(&signal); + assert!( + !chunks.is_empty(), + "continuous speech over the cap must emit at least one chunk" + ); + for chunk in &chunks { + assert!( + chunk.samples.len() <= FRAME_SAMPLES * 4, + "emitted chunk exceeded max_chunk_samples" + ); + } + } + + #[test] + fn flush_emits_in_flight_speech() { + let mut c = RmsVadChunker::new(); + // Sustained speech with NO closing silence. Without flush this + // stays buffered; flush must surface it as a final chunk. + let signal = constant_signal(FRAME_SAMPLES * 5, 0.01); + let chunks = c.push(&signal); + assert!( + chunks.is_empty(), + "in-progress speech with no silence close stays buffered until flush" + ); + let final_chunk = c.flush(); + assert!(final_chunk.is_some(), "flush must emit the in-flight chunk"); + } + + #[test] + fn flush_returns_none_when_idle() { + let mut c = RmsVadChunker::new(); + assert!(c.flush().is_none()); + let _ = c.push(&constant_signal(16_000, 0.0)); + assert!(c.flush().is_none(), "flushing pure silence emits nothing"); + } + + #[test] + fn reset_clears_state() { + let mut c = RmsVadChunker::new(); + let signal = constant_signal(FRAME_SAMPLES * 5, 0.01); + let _ = c.push(&signal); + c.reset(); + assert_eq!(c.next_sample_index(), 0); + // After reset, silence must not emit a chunk derived from pre-reset state. + let silence = constant_signal(FRAME_SAMPLES * 12, 0.0); + let chunks = c.push(&silence); + assert!(chunks.is_empty()); + assert!(c.flush().is_none()); + } + + #[test] + fn start_sample_includes_onset_audio() { + let mut c = RmsVadChunker::new(); + // First 2 frames silent (so next_sample_index is advanced but + // no onset). Then speech. + let silence = constant_signal(FRAME_SAMPLES * 2, 0.0); + let _ = c.push(&silence); + assert_eq!(c.next_sample_index(), (FRAME_SAMPLES * 2) as u64); + + let speech = constant_signal(FRAME_SAMPLES * 5, 0.01); + let closing_silence = constant_signal(FRAME_SAMPLES * 12, 0.0); + let mut signal = Vec::new(); + signal.extend(speech); + signal.extend(closing_silence); + let chunks = c.push(&signal); + assert_eq!(chunks.len(), 1); + let chunk = &chunks[0]; + // The chunk's start_sample should reflect the absolute index + // of the first onset-buffered sample, NOT the post-onset index. + assert!( + chunk.start_sample >= (FRAME_SAMPLES * 2) as u64, + "start_sample must be at or after the pre-speech silence" + ); + assert!( + chunk.start_sample + <= (FRAME_SAMPLES * 2 + FRAME_SAMPLES * DEFAULT_SPEECH_ONSET_FRAMES) as u64, + "start_sample must not skip past the onset frames" + ); + } +}