WHISPER!

config.yaml

@@ -5,7 +5,7 @@ whisper:
 #    n_threads: 4
     length_ms: 5000
     keep_ms: 200
-    step_ms: 500
+    step_ms: 5000
     max_tokens: 0
     audio_ctx: 0
     vad_thold: 0.6

src/config.rs

@@ -1,5 +1,4 @@
 use std::ffi::c_int;
-use std::fs;
 use std::net::IpAddr;
 use lazy_static::lazy_static;
 use serde::{Deserialize};
@@ -87,9 +86,11 @@ pub struct Config {
     pub(crate) server: Server,
 }
 
-#[tokio::test]
-async fn load() {
-    let config_str = fs::read_to_string("config.yaml").expect("failed to read config file");
-    let params: Config = serde_yaml::from_str(config_str.as_str()).expect("failed to parse config file");
-    println!("{:?}", params);
+mod tests {
+    #[tokio::test]
+    async fn load() {
+        let config_str = tokio::fs::read_to_string("config.yaml").await.expect("failed to read config file");
+        let params: crate::config::Config = serde_yaml::from_str(config_str.as_str()).expect("failed to parse config file");
+        println!("{:?}", params);
+    }
 }

src/main.rs

@@ -115,10 +115,10 @@ async fn stream_speaker(ctx: Data<&Context>, query: Query<LessonSpeakerQuery>, w
                     match msg.as_ref() {
                         Some(Ok(Message::Binary(bin))) => {
                             let _ = whisper.send(bin.to_vec()).await; // whisper test
-                            if let Err(e) = origin_tx.send(bin.to_vec()).await {
-                                tracing::warn!("failed to send voice: {}", e);
-                                break;
-                            }
+                            // if let Err(e) = origin_tx.send(bin.to_vec()).await {
+                            //     tracing::warn!("failed to send voice: {}", e);
+                            //     break;
+                            // }
                         },
                         Some(Ok(_)) => {
                             tracing::warn!("Other: {:?}", msg);

src/whisper.rs

@@ -5,7 +5,7 @@ use std::thread::sleep;
 use std::time::Duration;
 use lazy_static::lazy_static;
 use tokio::sync::{broadcast, mpsc, oneshot};
-use whisper_rs::{convert_integer_to_float_audio, WhisperState, WhisperContext, FullParams};
+use whisper_rs::{convert_integer_to_float_audio, WhisperState, WhisperContext};
 use whisper_rs_sys::WHISPER_SAMPLE_RATE;
 use crate::config::{WhisperParams, CONFIG};
 use crate::group::GroupedWithin;
@@ -79,28 +79,20 @@ pub struct WhisperHandler {
 
 impl WhisperHandler {
     pub(crate) fn new(config: WhisperParams) -> Result<Self, Error> {
-        let n_samples_step: usize = (config.step_ms * WHISPER_SAMPLE_RATE / 1000) as usize;
-        let n_samples_len: usize = (config.length_ms * WHISPER_SAMPLE_RATE / 1000) as usize;
-        let n_samples_keep: usize = (config.keep_ms * WHISPER_SAMPLE_RATE / 1000) as usize;
-        let n_new_line: usize = 1.max(config.length_ms / config.step_ms - 1) as usize;
-        let mut n_iter: usize = 0;
-
         let (stop_handle, mut stop_signal) = oneshot::channel();
         let (pcm_tx, pcm_rx) = mpsc::channel::<Vec<u8>>(128);
-        let mut grouped = GroupedWithin::new(
-            n_samples_step * 2,
-            Duration::from_millis(config.step_ms as u64),
-            pcm_rx,
-            u16::MAX as usize
-        );
         let (transcription_tx, _) = broadcast::channel::<Vec<Segment>>(128);
         let shared_transcription_tx = transcription_tx.clone();
-        let mut state = WHISPER_CONTEXT.create_state()
+        let state = WHISPER_CONTEXT.create_state()
             .map_err(|e| Error::whisper_error("failed to create WhisperState", e))?;
         tokio::task::spawn_blocking(move || {
-            let mut prompt_tokens: Vec<c_int> = Default::default();
-            let mut pcm_f32: VecDeque<f32> = VecDeque::from(vec![0f32; 30 * WHISPER_SAMPLE_RATE as usize]);
-            let mut last_offset: usize = 0;
+            let mut detector = Detector::new(state, &CONFIG.whisper);
+            let mut grouped = GroupedWithin::new(
+                detector.n_samples_step * 2,
+                Duration::from_millis(config.step_ms as u64),
+                pcm_rx,
+                u16::MAX as usize
+            );
             while let Err(oneshot::error::TryRecvError::Empty) = stop_signal.try_recv() {
                 let new_pcm_f32 = match grouped.next() {
                     Err(mpsc::error::TryRecvError::Disconnected) => break,
@@ -113,16 +105,9 @@ impl WhisperHandler {
                     }
                 };
 
-                let params = config.to_full_params(prompt_tokens.as_slice());
-                pcm_f32.extend(new_pcm_f32);
-                if pcm_f32.len() > n_samples_len + n_samples_keep {
-                    let _ = pcm_f32.drain(0..(pcm_f32.len() - n_samples_keep - n_samples_len)).len();
-                }
-                pcm_f32.make_contiguous();
-                let (data, _) = pcm_f32.as_slices();
-                let segments = match inference(&mut state, params, data, 0) {
-                    Ok((offset, result)) => {
-                        last_offset = offset;
+                detector.feed(new_pcm_f32);
+                let segments = match detector.inference() {
+                    Ok(result) => {
                         if result.is_empty() {
                             continue
                         }
@@ -134,19 +119,18 @@ impl WhisperHandler {
                     }
                 };
 
-                n_iter = n_iter + 1;
-
-                if n_iter % n_new_line == 0 {
-                    if let Err(e) = new_line(&mut pcm_f32, n_samples_keep, &config, &mut state, segments.len(), &mut prompt_tokens) {
-                        tracing::warn!("failed to new_line: {}", e);
+                if tracing::enabled!(tracing::Level::TRACE) {
+                    for segment in segments.iter() {
+                        tracing::trace!("[{}] SEGMENT: {}", detector.n_iter, segment.text);
                     }
-                    tracing::debug!("LINE: {}", segments.first().unwrap().text);
-
-                    if let Err(e) = shared_transcription_tx.send(segments) {
-                        tracing::error!("failed to send transcription: {}", e);
-                        break
-                    };
+                } else if tracing::enabled!(tracing::Level::DEBUG) {
+                    tracing::debug!("[{}] SEGMENT: {}", detector.n_iter, segments[0].text);
                 }
+
+                if let Err(e) = shared_transcription_tx.send(segments) {
+                    tracing::error!("failed to send transcription: {}", e);
+                    break
+                };
             }
         });
         Ok(Self {
@@ -165,66 +149,110 @@ impl WhisperHandler {
     }
 }
 
-fn inference(
-    state: &mut WhisperState,
-    params: FullParams,
-    pcm_f32: &[f32],
-    mut offset: usize,
-) -> Result<(usize, Vec<Segment>), Error> {
-
-    let _ = state.full(params, pcm_f32)
-        .map_err(|e| Error::whisper_error("failed to initialize WhisperState", e))?;
-
-    let num_segments = state
-        .full_n_segments()
-        .map_err(|e| Error::whisper_error("failed to get number of segments", e))?;
-    let mut segments: Vec<Segment> = Vec::with_capacity(num_segments as usize);
-    for i in 0..num_segments {
-        let segment = state
-            .full_get_segment_text(i)
-            .map_err(|e| Error::whisper_error("failed to get segment", e))?;
-        let timestamp_offset: i64 = (offset * 1000 / WHISPER_SAMPLE_RATE as usize) as i64;
-        let start_timestamp: i64 = timestamp_offset + 10 * state
-            .full_get_segment_t0(i)
-            .map_err(|e| Error::whisper_error("failed to get start timestamp", e))?;
-        let end_timestamp: i64 = timestamp_offset + 10 * state
-            .full_get_segment_t1(i)
-            .map_err(|e| Error::whisper_error("failed to get end timestamp", e))?;
-        // tracing::trace!("{}", segment);
-        segments.push(Segment { start_timestamp, end_timestamp, text: segment });
+struct Detector {
+    state: WhisperState<'static>,
+    config: &'static WhisperParams,
+    n_samples_keep: usize,
+    n_samples_step: usize,
+    n_samples_len: usize,
+    n_new_line: usize,
+    n_iter: usize,
+    prompt_tokens: Vec<c_int>,
+    pcm_f32: VecDeque<f32>,
+    offset: usize,
+}
+
+impl Detector {
+    fn new(state: WhisperState<'static>,
+           config: &'static WhisperParams) -> Self {
+        Detector {
+            state,
+            config,
+            n_samples_keep: (config.keep_ms * WHISPER_SAMPLE_RATE / 1000) as usize,
+            n_samples_step: (config.step_ms * WHISPER_SAMPLE_RATE / 1000) as usize,
+            n_samples_len: (config.length_ms * WHISPER_SAMPLE_RATE / 1000) as usize,
+            n_new_line: 1.max(config.length_ms / config.step_ms - 1) as usize,
+            n_iter: 0,
+            prompt_tokens: Default::default(),
+            pcm_f32: VecDeque::from(vec![0f32; 30 * WHISPER_SAMPLE_RATE as usize]),
+            offset: 0,
+        }
     }
 
-    Ok((offset, segments))
-}
+    fn feed(&mut self, new_pcm_f32: Vec<f32>) {
+        self.pcm_f32.extend(new_pcm_f32);
+        if self.pcm_f32.len() > self.n_samples_len + self.n_samples_keep {
+            let _ = self.pcm_f32.drain(0..(self.pcm_f32.len() - self.n_samples_keep - self.n_samples_len)).len();
+        }
+    }
 
-fn new_line(pcm_f32: &mut VecDeque<f32>,
-            n_samples_keep: usize,
-            config: &WhisperParams,
-            state: &mut WhisperState,
-            num_segments: usize,
-            prompt_tokens: &mut Vec<c_int>) -> Result<(), Error> {
+    fn inference(&mut self) -> Result<Vec<Segment>, Error> {
+        let params = self.config.to_full_params(self.prompt_tokens.as_slice());
+        let start = std::time::Instant::now();
+        let _ = self.state.full(params, self.pcm_f32.make_contiguous())
+            .map_err(|e| Error::whisper_error("failed to initialize WhisperState", e))?;
+        let end = std::time::Instant::now();
+        if end - start > Duration::from_millis(self.config.step_ms as u64) {
+            tracing::warn!("full() took {} ms too slow", (end - start).as_millis());
+        }
+
+        let num_segments = self.state
+            .full_n_segments()
+            .map_err(|e| Error::whisper_error("failed to get number of segments", e))?;
+        let mut segments: Vec<Segment> = Vec::with_capacity(num_segments as usize);
+        for i in 0..num_segments {
+            let segment = self.state
+                .full_get_segment_text(i)
+                .map_err(|e| Error::whisper_error("failed to get segment", e))?;
+            let timestamp_offset: i64 = (self.offset * 1000 / WHISPER_SAMPLE_RATE as usize) as i64;
+            let start_timestamp: i64 = timestamp_offset + 10 * self.state
+                .full_get_segment_t0(i)
+                .map_err(|e| Error::whisper_error("failed to get start timestamp", e))?;
+            let end_timestamp: i64 = timestamp_offset + 10 * self.state
+                .full_get_segment_t1(i)
+                .map_err(|e| Error::whisper_error("failed to get end timestamp", e))?;
+            // tracing::trace!("{}", segment);
+            segments.push(Segment { start_timestamp, end_timestamp, text: segment });
+        }
 
-    // keep the last n_samples_keep samples from pcm_f32
-    if pcm_f32.len() > n_samples_keep {
-        let _ = pcm_f32.drain(0..(pcm_f32.len() - n_samples_keep)).len();
+
+        self.n_iter = self.n_iter + 1;
+
+        if self.n_iter % self.n_new_line == 0 {
+            self.next_line()?;
+            Ok(segments)
+        } else {
+            Ok(vec![])
+        }
     }
 
-    if !config.no_context {
-        prompt_tokens.clear();
-
-        for i in 0..num_segments as c_int {
-            let token_count = state
-                .full_n_tokens(i)
-                .map_err(|e| Error::whisper_error("failed to get number of tokens", e))?;
-            for j in 0..token_count {
-                let token = state
-                    .full_get_token_id(i, j)
-                    .map_err(|e| Error::whisper_error("failed to get token", e))?;
-                prompt_tokens.push(token);
+    fn next_line(&mut self) -> Result<(), Error> {
+
+        // keep the last n_samples_keep samples from pcm_f32
+        if self.pcm_f32.len() > self.n_samples_keep {
+            let _ = self.pcm_f32.drain(0..(self.pcm_f32.len() - self.n_samples_keep)).len();
+        }
+
+        if !self.config.no_context {
+            self.prompt_tokens.clear();
+
+            let num_segments = self.state
+                .full_n_segments()
+                .map_err(|e| Error::whisper_error("failed to get number of segments", e))?;
+            for i in 0..num_segments {
+                let token_count = self.state
+                    .full_n_tokens(i)
+                    .map_err(|e| Error::whisper_error("failed to get number of tokens", e))?;
+                for j in 0..token_count {
+                    let token = self.state
+                        .full_get_token_id(i, j)
+                        .map_err(|e| Error::whisper_error("failed to get token", e))?;
+                    self.prompt_tokens.push(token);
+                }
             }
         }
+        Ok(())
     }
-    Ok(())
 }
 
 impl Drop for WhisperHandler {