feat: use vad

Author: quake

Cargo.lock

@@ -193,5 +193,7 @@ objc2-core-foundation = "0.3"
 objc2-foundation = "0.3"
 objc2-user-notifications = "0.3"
 
+voice_activity_detector = "0.2"
+
 [patch.crates-io]
 cpal = { git = "https://github.com/RustAudio/cpal", rev = "51c3b43" }

Cargo.toml

@@ -8,7 +8,6 @@ hound = { workspace = true }
 hypr-data = { workspace = true }
 
 [dependencies]
-hypr-vad = { workspace = true }
 kalosm-sound = { workspace = true, default-features = false }
 rodio = { workspace = true, features = ["wav"] }
 
@@ -17,3 +16,4 @@ serde = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread", "macros"] }
 tracing = { workspace = true }
+voice_activity_detector = { workspace = true }

crates/chunker/Cargo.toml

@@ -1,24 +1,17 @@
-mod error;
-mod predictor;
 mod stream;
 
-pub use error::*;
-pub use predictor::*;
 pub use stream::*;
 
 use kalosm_sound::AsyncSource;
 use std::time::Duration;
+use voice_activity_detector::VoiceActivityDetector;
 
 pub trait ChunkerExt: AsyncSource + Sized {
-    fn chunks<P: Predictor + Unpin>(
-        self,
-        predictor: P,
-        chunk_duration: Duration,
-    ) -> ChunkStream<Self, P>
+    fn chunks(self, vad: VoiceActivityDetector, chunk_duration: Duration) -> ChunkStream<Self>
     where
         Self: Unpin,
     {
-        ChunkStream::new(self, predictor, chunk_duration)
+        ChunkStream::new(self, vad, chunk_duration)
     }
 }
 
@@ -28,6 +21,7 @@ impl<T: AsyncSource> ChunkerExt for T {}
 mod tests {
     use super::*;
     use futures_util::StreamExt;
+    use voice_activity_detector::VoiceActivityDetector;
 
     #[tokio::test]
     async fn test_chunker() {
@@ -43,7 +37,12 @@ mod tests {
             sample_format: hound::SampleFormat::Float,
         };
 
-        let mut stream = audio_source.chunks(RMS::new(), Duration::from_secs(15));
+        let vad = VoiceActivityDetector::builder()
+            .sample_rate(16000)
+            .chunk_size(512usize)
+            .build()
+            .unwrap();
+        let mut stream = audio_source.chunks(vad, Duration::from_secs(15));
         let mut i = 0;
 
         std::fs::remove_dir_all("tmp/english_1").unwrap();

crates/chunker/src/error.rs

@@ -7,21 +7,20 @@ use std::{
 
 use kalosm_sound::AsyncSource;
 use rodio::buffer::SamplesBuffer;
+use voice_activity_detector::{IteratorExt, VoiceActivityDetector};
 
-use crate::Predictor;
-
-pub struct ChunkStream<S: AsyncSource + Unpin, P: Predictor + Unpin> {
+pub struct ChunkStream<S: AsyncSource + Unpin> {
     source: S,
-    predictor: P,
+    vad: VoiceActivityDetector,
     buffer: Vec<f32>,
     max_duration: Duration,
 }
 
-impl<S: AsyncSource + Unpin, P: Predictor + Unpin> ChunkStream<S, P> {
-    pub fn new(source: S, predictor: P, max_duration: Duration) -> Self {
+impl<S: AsyncSource + Unpin> ChunkStream<S> {
+    pub fn new(source: S, vad: VoiceActivityDetector, max_duration: Duration) -> Self {
         Self {
             source,
-            predictor,
+            vad,
             buffer: Vec::new(),
             max_duration,
         }
@@ -34,26 +33,9 @@ impl<S: AsyncSource + Unpin, P: Predictor + Unpin> ChunkStream<S, P> {
     fn samples_for_duration(&self, duration: Duration) -> usize {
         (self.source.sample_rate() as f64 * duration.as_secs_f64()) as usize
     }
-
-    fn trim_silence(predictor: &P, data: &mut Vec<f32>) {
-        const WINDOW_SIZE: usize = 100;
-
-        let mut trim_index = 0;
-        for start_idx in (0..data.len()).step_by(WINDOW_SIZE) {
-            let end_idx = (start_idx + WINDOW_SIZE).min(data.len());
-            let window = &data[start_idx..end_idx];
-
-            if let Ok(false) = predictor.predict(window) {
-                trim_index = start_idx;
-                break;
-            }
-        }
-
-        data.drain(0..trim_index);
-    }
 }
 
-impl<S: AsyncSource + Unpin, P: Predictor + Unpin> Stream for ChunkStream<S, P> {
+impl<S: AsyncSource + Unpin> Stream for ChunkStream<S> {
     type Item = SamplesBuffer<f32>;
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
@@ -62,8 +44,6 @@ impl<S: AsyncSource + Unpin, P: Predictor + Unpin> Stream for ChunkStream<S, P>
         let sample_rate = this.source.sample_rate();
 
         let min_buffer_samples = this.samples_for_duration(Duration::from_secs(6));
-        let silence_window_samples = this.samples_for_duration(Duration::from_millis(500));
-
         let stream = this.source.as_stream();
         let mut stream = std::pin::pin!(stream);
 
@@ -73,32 +53,41 @@ impl<S: AsyncSource + Unpin, P: Predictor + Unpin> Stream for ChunkStream<S, P>
                     this.buffer.push(sample);
 
                     if this.buffer.len() >= min_buffer_samples {
-                        let buffer_len = this.buffer.len();
-                        let silence_start = buffer_len.saturating_sub(silence_window_samples);
-                        let last_samples = &this.buffer[silence_start..buffer_len];
-
-                        if let Ok(false) = this.predictor.predict(last_samples) {
-                            let mut data = std::mem::take(&mut this.buffer);
-                            Self::trim_silence(&this.predictor, &mut data);
-
-                            return Poll::Ready(Some(SamplesBuffer::new(1, sample_rate, data)));
-                        }
+                        let data = std::mem::take(&mut this.buffer);
+                        let speech = filter_speech_chunks(&mut this.vad, data);
+                        return Poll::Ready(Some(SamplesBuffer::new(1, sample_rate, speech)));
                     }
                 }
                 Poll::Ready(None) if !this.buffer.is_empty() => {
-                    let mut data = std::mem::take(&mut this.buffer);
-                    Self::trim_silence(&this.predictor, &mut data);
-
-                    return Poll::Ready(Some(SamplesBuffer::new(1, sample_rate, data)));
+                    let data = std::mem::take(&mut this.buffer);
+                    let speech = filter_speech_chunks(&mut this.vad, data);
+                    return Poll::Ready(Some(SamplesBuffer::new(1, sample_rate, speech)));
                 }
                 Poll::Ready(None) => return Poll::Ready(None),
                 Poll::Pending => return Poll::Pending,
             }
         }
 
-        let mut chunk: Vec<_> = this.buffer.drain(0..max_samples).collect();
-        Self::trim_silence(&this.predictor, &mut chunk);
-
-        Poll::Ready(Some(SamplesBuffer::new(1, sample_rate, chunk)))
+        let data = this.buffer.drain(0..max_samples);
+        let speech = filter_speech_chunks(&mut this.vad, data);
+        Poll::Ready(Some(SamplesBuffer::new(1, sample_rate, speech)))
     }
 }
+
+// helper function to filter speech chunks
+fn filter_speech_chunks<D: IntoIterator<Item = f32>>(
+    vad: &mut VoiceActivityDetector,
+    data: D,
+) -> Vec<f32> {
+    data.into_iter()
+        .label(vad, 0.75, 3)
+        .filter_map(|label| {
+            if label.is_speech() {
+                Some(label.into_iter())
+            } else {
+                None
+            }
+        })
+        .flatten()
+        .collect()
+}

crates/chunker/src/lib.rs

@@ -50,6 +50,8 @@ futures-util = { workspace = true }
 tokio = { workspace = true, features = ["rt", "macros"] }
 tokio-util = { workspace = true }
 
+voice_activity_detector = { workspace = true }
+
 [target.'cfg(not(target_os = "macos"))'.dependencies]
 kalosm-sound = { workspace = true, default-features = false }
 

crates/chunker/src/predictor.rs

@@ -21,6 +21,8 @@ use hypr_chunker::ChunkerExt;
 use hypr_listener_interface::{ListenOutputChunk, ListenParams, Word};
 use hypr_ws_utils::WebSocketAudioSource;
 
+use voice_activity_detector::VoiceActivityDetector;
+
 use crate::manager::{ConnectionGuard, ConnectionManager};
 
 #[derive(Default)]
@@ -144,9 +146,14 @@ async fn websocket_with_model(
 async fn websocket(socket: WebSocket, model: hypr_whisper::local::Whisper, guard: ConnectionGuard) {
     let (mut ws_sender, ws_receiver) = socket.split();
     let mut stream = {
-        let audio_source = WebSocketAudioSource::new(ws_receiver, 16 * 1000);
-        let chunked =
-            audio_source.chunks(hypr_chunker::RMS::new(), std::time::Duration::from_secs(15));
+        let sample_rate = 16_000;
+        let audio_source = WebSocketAudioSource::new(ws_receiver, sample_rate);
+        let vad = VoiceActivityDetector::builder()
+            .sample_rate(sample_rate)
+            .chunk_size(512usize)
+            .build()
+            .expect("vad config is valid");
+        let chunked = audio_source.chunks(vad, std::time::Duration::from_secs(15));
         hypr_whisper::local::TranscribeChunkedAudioStreamExt::transcribe(chunked, model)
     };