WHISPER!

config.yaml

@@ -3,10 +3,10 @@ server:
     host: ::1
 whisper:
 #    n_threads: 4
-    step_ms: 500
-    length_ms: 5000
-    keep_ms: 5000
-    max_tokens: 32
+    length_ms: 10000
+    keep_ms: 200
+    step_ms: 1000
+    max_tokens: 0
     audio_ctx: 0
     vad_thold: 0.6
     freq_thold: 100.0

src/config.rs

@@ -1,4 +1,5 @@
 use std::ffi::c_int;
+use std::fs;
 use std::net::IpAddr;
 use lazy_static::lazy_static;
 use serde::{Deserialize};
@@ -24,8 +25,8 @@ lazy_static! {
 #[derive(Debug, Deserialize, Clone)]
 pub(crate) struct WhisperParams {
     pub(crate) n_threads: Option<usize>,
-    // pub(crate) step_ms: u32,
-    // pub(crate) length_ms: u32,
+    pub(crate) step_ms: u32,
+    pub(crate) length_ms: u32,
     pub(crate) keep_ms: u32,
     pub(crate) max_tokens: u32,
     pub(crate) audio_ctx: u32,
@@ -52,7 +53,7 @@ impl WhisperParams {
         param.set_print_realtime(false);
         param.set_print_timestamps(!self.no_timestamps);
         param.set_translate(self.translate);
-        param.set_single_segment(false);
+        param.set_single_segment(true);
         param.set_max_tokens(self.max_tokens as i32);
         let lang = self.language.as_ref().map(|s| s.as_str());
         param.set_language(lang);
@@ -62,7 +63,7 @@ impl WhisperParams {
         param.set_speed_up(self.speed_up);
         // param.set_tdrz_enable(self.tinydiarize);
         if self.no_fallback {
-            param.set_temperature_inc(0.0);
+            param.set_temperature_inc(-1.0);
         }
         if self.no_context {
             param.set_tokens(&NONE);

src/group.rs

@@ -1,3 +1,4 @@
+use std::collections::VecDeque;
 use std::time::Duration;
 use tokio::{select};
 use tokio::time::sleep;
@@ -17,13 +18,13 @@ where I: 'static + Send {
                buffer: usize) -> Self {
         let (tx, outlet) = channel::<Vec<I>>(buffer);
         tokio::spawn(async move {
-            let mut window = Vec::with_capacity(group_size);
+            let mut window = VecDeque::with_capacity(group_size);
 
             loop {
                 let grouped_fut = async {
                     while let Some(c) = inlet.recv().await {
                         window.extend(c);
-                        if window.len() > group_size {
+                        if window.len() >= group_size {
                             let will_send: Vec<I> = window.drain(0..group_size).collect();
                             return Some(will_send)
                         }
@@ -31,18 +32,22 @@ where I: 'static + Send {
                     return None
                 };
 
-                let grouped = select! {
+                let grouped: Vec<I> = select! {
                     _ = sleep(window_time) => {
                         window.drain(..).collect()
                     },
                     grouped_opt = grouped_fut => {
                         match grouped_opt {
                             None => break,
-                            Some(grouped) => grouped
+                            Some(group) => group
                         }
                     }
                 };
 
+                if grouped.is_empty() {
+                    continue
+                }
+
                 if let Err(e) = tx.send(grouped).await {
                     tracing::error!("{}", e);
                 }

src/main.rs

@@ -35,14 +35,6 @@ struct Opt {
     /// The AWS Region.
     #[structopt(short, long)]
     region: Option<String>,
-    //
-    // /// The name of the audio file.
-    // #[structopt(short, long)]
-    // audio_file: String,
-    //
-    /// Whether to display additional information.
-    #[structopt(short, long)]
-    verbose: bool,
 }
 
 #[derive(Clone)]
@@ -57,22 +49,18 @@ async fn main() -> Result<(), std::io::Error> {
 
     let Opt {
         region,
-        verbose,
     } = Opt::parse();
 
     let region_provider = RegionProviderChain::first_try(region.map(Region::new))
         .or_default_provider()
         .or_else(Region::new("us-west-2"));
 
-    println!();
-
-    if verbose {
-        println!("Transcribe client version: {}", PKG_VERSION);
-        println!(
+    if tracing::enabled!(tracing::Level::DEBUG) {
+        tracing::debug!("Transcribe client version: {}", PKG_VERSION);
+        tracing::debug!(
             "Region:                    {}",
             region_provider.region().await.unwrap().as_ref()
         );
-        println!();
     }
 
     let shared_config = aws_config::from_env().region(region_provider).load().await;
@@ -118,36 +106,37 @@ async fn stream_speaker(ctx: Data<&Context>, query: Query<LessonSpeakerQuery>, w
         loop {
             select! {
                 w = whisper_transcribe_rx.recv() => {
-                    let Ok(txt) = w else {
+                    let Ok(_txt) = w else {
+                        // TODO: handle msg
                         continue
                     };
-                    println!("Whisper: {:?}", txt)
                 }
                 msg = socket.next() => {
                     match msg.as_ref() {
                         Some(Ok(Message::Binary(bin))) => {
-                            let _ = whisper.send(bin.clone()).await; // whisper test
-                            if origin_tx.send(bin.to_vec()).await.is_err() {
-                                println!("tx closed");
+                            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;
                             }
                         },
                         Some(Ok(_)) => {
-                            println!("Other: {:?}", msg);
+                            tracing::warn!("Other: {:?}", msg);
                         },
                         Some(Err(e)) => {
-                            println!("Error: {:?}", e);
+                            tracing::warn!("Error: {:?}", e);
                         },
                         None => {
-                            let _ = socket.close().await;
-                            println!("Other: {:?}", msg);
+                            if let Err(e) = socket.close().await {
+                                tracing::debug!("Message: {:?}, {}", msg, e);
+                            }
                             break;
                         }
                     }
                 },
                 output = transcribe_rx.recv() => {
                     if let Ok(transcript) = output {
-                        println!("Transcribed: {}", transcript);
+                        tracing::trace!("Transcribed: {}", transcript);
                         let evt = LiveLessonTextEvent::Transcription { text: transcript.clone() };
                         let json = serde_json::to_string(&evt).expect("failed to serialize");
                         let _ = socket.send(Message::Text(json)).await.expect("failed to send");
@@ -177,10 +166,16 @@ enum LiveLessonTextEvent {
 #[handler]
 async fn stream_listener(ctx: Data<&Context>, query: Query<LessonListenerQuery>, ws: WebSocket) -> impl IntoResponse {
     let lesson_opt = ctx.lessons_manager.get_lesson(query.id).await;
-    println!("{:?}", query);
-    let voice_id = query.voice.parse().expect("Not supported voice");
+    tracing::debug!("listener param = {:?}", query);
 
     ws.on_upgrade(|mut socket| async move {
+        let voice_id = match query.voice.parse() {
+            Ok(id) => id,
+            Err(e) => {
+                let _ = socket.send(Message::Text(format!("invalid voice id: {}", e))).await;
+                return
+            }
+        };
         let Some(lesson) = lesson_opt else {
             let _ = socket.send(Message::Text("lesson not found".to_string())).await;
             return
@@ -197,17 +192,29 @@ async fn stream_listener(ctx: Data<&Context>, query: Query<LessonListenerQuery>,
                 transcript = transcript_rx.recv() => {
                     if let Ok(transcript) = transcript {
                         let evt = LiveLessonTextEvent::Transcription { text: transcript };
-                        let json = serde_json::to_string(&evt).expect("failed to serialize");
-                        println!("Transcribed: {}", json);
-                        let _ = socket.send(Message::Text(json)).await;
+                        match serde_json::to_string(&evt) {
+                            Ok(json) => {
+                                tracing::debug!("Transcribed: {}", json);
+                                let _ = socket.send(Message::Text(json)).await;
+                            },
+                            Err(e) => {
+                                tracing::error!("failed to serialize: {}", e);
+                            }
+                        }
                     }
                 },
                 translated = translate_rx.recv() => {
                     if let Ok(translated) = translated {
                         let evt = LiveLessonTextEvent::Translation { text: translated };
-                        let json = serde_json::to_string(&evt).expect("failed to serialize");
-                        println!("Translated: {}", json);
-                        let _ = socket.send(Message::Text(json)).await;
+                        match serde_json::to_string(&evt) {
+                            Ok(json) => {
+                                tracing::debug!("Translated: {}", json);
+                                let _ = socket.send(Message::Text(json)).await;
+                            },
+                            Err(e) => {
+                                tracing::error!("failed to serialize: {}", e);
+                            }
+                        }
                     }
                 },
                 voice = voice_rx.recv() => {

src/whisper.rs

@@ -1,10 +1,11 @@
+use std::collections::VecDeque;
 use std::ffi::c_int;
 use std::fmt::{Debug, Display, Formatter};
 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, WhisperContext, WhisperState};
+use whisper_rs::{convert_integer_to_float_audio, WhisperState, WhisperContext, FullParams};
 use whisper_rs_sys::WHISPER_SAMPLE_RATE;
 use crate::config::{WhisperParams, CONFIG};
 use crate::group::GroupedWithin;
@@ -78,22 +79,28 @@ 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_keep,
-            Duration::from_secs(5),
+            n_samples_step * 2,
+            Duration::from_millis(config.step_ms as u64),
             pcm_rx,
-            1024
+            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()
             .map_err(|e| Error::whisper_error("failed to create WhisperState", e))?;
         tokio::task::spawn_blocking(move || {
-            let mut tokens: Vec<c_int> = Default::default();
-            let mut pcm_f32: Vec<f32> = Default::default();
+            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;
             while let Err(oneshot::error::TryRecvError::Empty) = stop_signal.try_recv() {
                 let new_pcm_f32 = match grouped.next() {
                     Err(mpsc::error::TryRecvError::Disconnected) => break,
@@ -106,22 +113,39 @@ impl WhisperHandler {
                     }
                 };
 
+                let params = config.to_full_params(prompt_tokens.as_slice());
                 pcm_f32.extend(new_pcm_f32);
-                match inference(&mut state, &config, n_samples_keep, &mut tokens, &mut pcm_f32) {
-                    Ok(segments) => {
-                        if segments.is_empty() {
+                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;
+                        if result.is_empty() {
                             continue
                         }
-                        if let Err(e) = shared_transcription_tx.send(segments) {
-                            tracing::error!("failed to send transcription: {}", e);
-                            break
-                        }
+                        result
                     }
                     Err(err) => {
-                        tracing::error!("failed to run whisper: {}", err);
+                        tracing::warn!("failed to inference: {}", err);
                         continue
-                        // break
                     }
+                };
+
+                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);
+                    }
+                    tracing::debug!("LINE: {}", segments.first().unwrap().text);
+
+                    if let Err(e) = shared_transcription_tx.send(segments) {
+                        tracing::error!("failed to send transcription: {}", e);
+                        break
+                    };
                 }
             }
         });
@@ -143,17 +167,13 @@ impl WhisperHandler {
 
 fn inference(
     state: &mut WhisperState,
-    config: &WhisperParams,
-    n_samples_keep: usize,
-    prompt_tokens: &mut Vec<c_int>,
-    pcm_f32: &mut Vec<f32>
-) -> Result<Vec<Segment>, Error> {
-    let params = config.to_full_params(prompt_tokens.as_slice());
-
-    let st = std::time::Instant::now();
-    let _ = state.full(params, pcm_f32.as_slice())
+    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 et = std::time::Instant::now();
 
     let num_segments = state
         .full_n_segments()
@@ -163,28 +183,36 @@ fn inference(
         let segment = state
             .full_get_segment_text(i)
             .map_err(|e| Error::whisper_error("failed to get segment", e))?;
-        let start_timestamp = state
+        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 = state
+        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::debug!("[{} - {}]: {}", start_timestamp, end_timestamp, segment);
+        // tracing::trace!("{}", segment);
         segments.push(Segment { start_timestamp, end_timestamp, text: segment });
     }
 
+    Ok((offset, segments))
+}
+
+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> {
+
+    // 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();
+    }
+
     if !config.no_context {
         prompt_tokens.clear();
 
-        // 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)).collect::<Vec<_>>();
-        }
-
-        let n_segments = state
-            .full_n_segments()
-            .map_err(|e| Error::whisper_error("failed to get number of segments", e))?;
-        for i in 0..n_segments {
+        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))?;
@@ -196,9 +224,7 @@ fn inference(
             }
         }
     }
-
-    tracing::trace!("took {}ms", (et - st).as_millis());
-    Ok(segments)
+    Ok(())
 }
 
 impl Drop for WhisperHandler {