WHISPER!

Cargo.lock

@@ -1419,14 +1419,17 @@ dependencies = [
  "aws-sdk-translate",
  "clap",
  "futures-util",
+ "lazy_static",
  "poem",
  "serde",
  "serde_json",
  "serde_yaml",
  "tokio",
  "tokio-stream",
+ "tracing",
  "tracing-subscriber",
  "whisper-rs",
+ "whisper-rs-sys",
 ]
 
 [[package]]
@@ -2072,9 +2075,9 @@ checksum = "b6bc1c9ce2b5135ac7f93c72918fc37feb872bdc6a5533a8b85eb4b86bfdae52"
 
 [[package]]
 name = "tracing"
-version = "0.1.39"
+version = "0.1.40"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ee2ef2af84856a50c1d430afce2fdded0a4ec7eda868db86409b4543df0797f9"
+checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef"
 dependencies = [
  "log",
  "pin-project-lite",

Cargo.toml

@@ -12,12 +12,15 @@ clap = { version = "4.4.6" , features = ["derive"]}
 tokio = {  version = "1.33.0" , features = ["full"] }
 tokio-stream = "0.1.14"
 async-stream = "0.3.5"
-tracing-subscriber = "0.3.17"
 futures-util = "0.3.28"
 serde = { version = "1.0.189", features = ["derive"] }
 serde_json = { version = "1.0.107", features = [] }
 serde_yaml = "0.9.25"
 whisper-rs = {  version = "0.8.0" , features = ["coreml"] }
+whisper-rs-sys = "0.6.1"
+tracing = "0.1.40"
+tracing-subscriber = "0.3.17"
+lazy_static = "1.4.0"
 
 [dependencies.poem]
 version = "1.3.58"

config.yaml

@@ -2,11 +2,10 @@ server:
     port: 8080
     host: ::1
 whisper:
-    n_threads: 4
+#    n_threads: 4
     step_ms: 500
     length_ms: 5000
-    keep_ms: 200
-    capture_id: -1
+    keep_ms: 5000
     max_tokens: 32
     audio_ctx: 0
     vad_thold: 0.6
@@ -15,7 +14,7 @@ whisper:
     translate: false
     no_fallback: false
     print_special: false
-    no_context: true
+    no_context: false
     no_timestamps: false
     tinydiarize: false
     language: "en"

src/config.rs

@@ -1,27 +1,43 @@
 use std::ffi::c_int;
-use std::fs;
 use std::net::IpAddr;
+use lazy_static::lazy_static;
 use serde::{Deserialize};
 use whisper_rs::FullParams;
 
-#[derive(Debug, Deserialize)]
+#[derive(Debug)]
+pub enum Error {
+    IoError(std::io::Error),
+    ConfigError(serde_yaml::Error),
+}
+
+pub(crate) fn load_config() -> Result<Config, Error> {
+    let config_str = std::fs::read_to_string("config.yaml").map_err(|e| Error::IoError(e))?;
+    let config: Config = serde_yaml::from_str(config_str.as_str())
+        .map_err(|e| Error::ConfigError(e))?;
+    return Ok(config)
+}
+
+lazy_static! {
+    pub static ref CONFIG: Config = load_config().expect("failed to load config");
+}
+
+#[derive(Debug, Deserialize, Clone)]
 pub(crate) struct WhisperParams {
     pub(crate) n_threads: Option<usize>,
-    pub(crate) step_ms: i32,
-    pub(crate) length_ms: i32,
-    pub(crate) keep_ms: i32,
-    pub(crate) capture_id: i32,
-    pub(crate) max_tokens: i32,
-    pub(crate) audio_ctx: i32,
-    pub(crate) vad_thold: f32,
-    pub(crate) freq_thold: f32,
+    // 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,
+    // pub(crate) vad_thold: f32,
+    // pub(crate) freq_thold: f32,
     pub(crate) speed_up: bool,
     pub(crate) translate: bool,
     pub(crate) no_fallback: bool,
     pub(crate) print_special: bool,
     pub(crate) no_context: bool,
     pub(crate) no_timestamps: bool,
-    pub(crate) tinydiarize: bool,
+    // pub(crate) tinydiarize: bool,
     pub(crate) language: Option<String>,
     pub(crate) model: String,
 }
@@ -29,20 +45,20 @@ pub(crate) struct WhisperParams {
 const NONE: [c_int;0] = [];
 
 impl WhisperParams {
-    pub(crate) fn to_full_params<'a, 'b>(&'a self) -> FullParams<'a, 'b> {
+    pub(crate) fn to_full_params<'a, 'b>(&'a self, tokens: &'b [c_int]) -> FullParams<'a, 'b> {
         let mut param = FullParams::new(Default::default());
         param.set_print_progress(false);
         param.set_print_special(self.print_special);
         param.set_print_realtime(false);
         param.set_print_timestamps(!self.no_timestamps);
         param.set_translate(self.translate);
-        param.set_single_segment(true);
-        param.set_max_tokens(self.max_tokens);
+        param.set_single_segment(false);
+        param.set_max_tokens(self.max_tokens as i32);
         let lang = self.language.as_ref().map(|s| s.as_str());
         param.set_language(lang);
         let num_cpus = std::thread::available_parallelism().map(|c| c.get()).unwrap_or(4);
         param.set_n_threads(self.n_threads.unwrap_or(num_cpus) as c_int);
-        param.set_audio_ctx(self.audio_ctx);
+        param.set_audio_ctx(self.audio_ctx as i32);
         param.set_speed_up(self.speed_up);
         // param.set_tdrz_enable(self.tinydiarize);
         if self.no_fallback {
@@ -50,6 +66,8 @@ impl WhisperParams {
         }
         if self.no_context {
             param.set_tokens(&NONE);
+        } else {
+            param.set_tokens(&tokens);
         }
 
         param
@@ -63,7 +81,7 @@ pub(crate) struct Server {
 }
 
 #[derive(Debug, Deserialize)]
-pub(crate) struct Config {
+pub struct Config {
     pub(crate) whisper: WhisperParams,
     pub(crate) server: Server,
 }

src/group.rs

@@ -0,0 +1,59 @@
+use std::time::Duration;
+use tokio::{select};
+use tokio::time::sleep;
+use tokio::sync::mpsc::{Receiver, channel};
+use tokio::sync::mpsc::error::TryRecvError;
+
+pub struct GroupedWithin<I>
+where I: 'static + Send {
+    outlet: Receiver<Vec<I>>
+}
+
+impl<I> GroupedWithin<I>
+where I: 'static + Send {
+    pub fn new(group_size: usize,
+               window_time: Duration,
+               mut inlet: Receiver<Vec<I>>,
+               buffer: usize) -> Self {
+        let (tx, outlet) = channel::<Vec<I>>(buffer);
+        tokio::spawn(async move {
+            let mut window = Vec::with_capacity(group_size);
+
+            loop {
+                let grouped_fut = async {
+                    while let Some(c) = inlet.recv().await {
+                        window.extend(c);
+                        if window.len() > group_size {
+                            let will_send: Vec<I> = window.drain(0..group_size).collect();
+                            return Some(will_send)
+                        }
+                    }
+                    return None
+                };
+
+                let grouped = select! {
+                    _ = sleep(window_time) => {
+                        window.drain(..).collect()
+                    },
+                    grouped_opt = grouped_fut => {
+                        match grouped_opt {
+                            None => break,
+                            Some(grouped) => grouped
+                        }
+                    }
+                };
+
+                if let Err(e) = tx.send(grouped).await {
+                    tracing::error!("{}", e);
+                }
+            }
+        });
+        Self {
+            outlet
+        }
+    }
+
+    pub fn next(&mut self) -> Result<Vec<I>, TryRecvError> {
+        self.outlet.try_recv()
+    }
+}

src/main.rs

@@ -17,17 +17,17 @@ use poem::web::websocket::{Message, WebSocket};
 use futures_util::stream::StreamExt;
 use poem::web::{Data, Query};
 
-use tokio::{fs, select};
+use tokio::{select};
 use serde::{Deserialize, Serialize};
-use whisper_rs::WhisperContext;
 use lesson::{LessonsManager};
-use crate::config::Config;
+use crate::config::CONFIG;
 use crate::lesson::Viseme;
-use crate::whisper::run_whisper;
+use crate::whisper::WhisperHandler;
 
 mod lesson;
 mod config;
 mod whisper;
+mod group;
 
 
 #[derive(Debug, Parser)]
@@ -50,26 +50,11 @@ struct Context {
     lessons_manager: LessonsManager,
 }
 
-#[derive(Debug)]
-enum Error {
-    IoError(std::io::Error),
-    ConfigError(serde_yaml::Error),
-}
-
-async fn load_config() -> Result<Config, Error> {
-    let config_str = fs::read_to_string("config.yaml").await.map_err(|e| Error::IoError(e))?;
-    let config: Config = serde_yaml::from_str(config_str.as_str())
-        .map_err(|e| Error::ConfigError(e))?;
-    return Ok(config)
-}
 
 #[tokio::main]
 async fn main() -> Result<(), std::io::Error> {
     tracing_subscriber::fmt::init();
 
-    let config = load_config().await.expect("failed to load config");
-    run_whisper(&config).await;
-
     let Opt {
         region,
         verbose,
@@ -107,7 +92,8 @@ async fn main() -> Result<(), std::io::Error> {
         .at("lesson-speaker", StaticFileEndpoint::new("./static/index.html"))
         .at("lesson-listener", StaticFileEndpoint::new("./static/index.html"))
         .data(ctx);
-    let listener = TcpListener::bind("[::]:8080");
+    let addr = format!("{}:{}", CONFIG.server.host, CONFIG.server.port);
+    let listener = TcpListener::bind(addr);
     let server = Server::new(listener);
 
     server.run(app).await
@@ -127,11 +113,20 @@ async fn stream_speaker(ctx: Data<&Context>, query: Query<LessonSpeakerQuery>, w
     ws.on_upgrade(|mut socket| async move {
         let origin_tx = lesson.voice_channel();
         let mut transcribe_rx = lesson.transcript_channel();
+        let whisper = WhisperHandler::new(CONFIG.whisper.clone()).expect("failed to create whisper");
+        let mut whisper_transcribe_rx = whisper.subscribe();
         loop {
             select! {
+                w = whisper_transcribe_rx.recv() => {
+                    let Ok(txt) = w else {
+                        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");
                                 break;
@@ -217,7 +212,6 @@ async fn stream_listener(ctx: Data<&Context>, query: Query<LessonListenerQuery>,
                 },
                 voice = voice_rx.recv() => {
                     if let Ok(voice) = voice {
-                        println!("Synthesized: {:?}", voice.len());
                         let _ = socket.send(Message::Binary(voice)).await;
                     }
                 },
@@ -225,7 +219,6 @@ async fn stream_listener(ctx: Data<&Context>, query: Query<LessonListenerQuery>,
                     if let Ok(visemes) = visemes {
                         let evt = LiveLessonTextEvent::LipSync { visemes };
                         let json = serde_json::to_string(&evt).expect("failed to serialize");
-                        println!("Visemes: {:?}", json);
                         let _ = socket.send(Message::Text(json)).await;
                     }
                 },

src/whisper.rs

@@ -1,14 +1,57 @@
-use whisper_rs::WhisperContext;
-use crate::config::Config;
+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_sys::WHISPER_SAMPLE_RATE;
+use crate::config::{WhisperParams, CONFIG};
+use crate::group::GroupedWithin;
 
-pub(crate) async fn run_whisper(config: &Config) {
-    let ctx = WhisperContext::new(&*config.whisper.model).expect("failed to load whisper context");
-    let mut _state = ctx.create_state().expect("failed to create state");
-    let params = (&config.whisper).to_full_params();
-    _state.full(params, &[]).expect("TODO: panic message");
+lazy_static! {
+    static ref WHISPER_CONTEXT: WhisperContext = {
+        WhisperContext::new(&*CONFIG.whisper.model)
+            .expect("failed to create WhisperContext")
+    };
 }
 
-async fn pcm_i16_to_f32(input: &Vec<u8>) -> Vec<f32> {
+#[derive(Debug)]
+pub(crate) enum Error {
+    WhisperError {
+        description: String,
+        error: whisper_rs::WhisperError,
+    },
+}
+
+impl Error {
+    fn whisper_error(description: &str, error: whisper_rs::WhisperError) -> Self {
+        Self::WhisperError {
+            description: description.to_string(),
+            error,
+        }
+    }
+}
+
+impl Display for Error {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::WhisperError { description, error } => {
+                write!(f, "WhisperError: {}: {}", description, error)
+            }
+        }
+    }
+}
+
+impl std::error::Error for Error {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        match self {
+            Self::WhisperError { error, .. } => Some(error),
+        }
+    }
+}
+
+fn pcm_i16_to_f32(input: &Vec<u8>) -> Vec<f32> {
     let pcm_i16 = input
         .chunks_exact(2)
         .map(|chunk| {
@@ -17,13 +60,154 @@ async fn pcm_i16_to_f32(input: &Vec<u8>) -> Vec<f32> {
             i16::from_le_bytes(buf)
         })
         .collect::<Vec<i16>>();
-    let pcm_f32 = pcm_i16
-        .iter()
-        .map(|i| *i as f32 / i16::MAX as f32)
-        .collect::<Vec<f32>>();
-    pcm_f32
+    convert_integer_to_float_audio(pcm_i16.as_slice())
+}
+
+#[derive(Clone, Debug)]
+pub struct Segment {
+    pub start_timestamp: i64,
+    pub end_timestamp: i64,
+    pub text: String,
+}
+
+pub struct WhisperHandler {
+    tx: mpsc::Sender<Vec<u8>>,
+    transcription_tx: broadcast::Sender<Vec<Segment>>,
+    stop_handle: Option<oneshot::Sender<()>>,
+}
+
+impl WhisperHandler {
+    pub(crate) fn new(config: WhisperParams) -> Result<Self, Error> {
+        let n_samples_keep: usize = (config.keep_ms * WHISPER_SAMPLE_RATE / 1000) as usize;
+        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),
+            pcm_rx,
+            1024
+        );
+        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();
+            while let Err(oneshot::error::TryRecvError::Empty) = stop_signal.try_recv() {
+                let new_pcm_f32 = match grouped.next() {
+                    Err(mpsc::error::TryRecvError::Disconnected) => break,
+                    Err(mpsc::error::TryRecvError::Empty) => {
+                        sleep(Duration::from_millis(10));
+                        continue
+                    }
+                    Ok(data) => {
+                        pcm_i16_to_f32(&data)
+                    }
+                };
+
+                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() {
+                            continue
+                        }
+                        if let Err(e) = shared_transcription_tx.send(segments) {
+                            tracing::error!("failed to send transcription: {}", e);
+                            break
+                        }
+                    }
+                    Err(err) => {
+                        tracing::error!("failed to run whisper: {}", err);
+                        continue
+                        // break
+                    }
+                }
+            }
+        });
+        Ok(Self {
+            tx: pcm_tx,
+            transcription_tx,
+            stop_handle: Some(stop_handle),
+        })
+    }
+
+    pub fn subscribe(&self) -> broadcast::Receiver<Vec<Segment>> {
+        self.transcription_tx.subscribe()
+    }
+
+    pub async fn send(&self, data: Vec<u8>) -> Result<(), mpsc::error::SendError<Vec<u8>>> {
+        self.tx.send(data).await
+    }
 }
 
-struct 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())
+        .map_err(|e| Error::whisper_error("failed to initialize WhisperState", e))?;
+    let et = std::time::Instant::now();
+
+    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 start_timestamp = state
+            .full_get_segment_t0(i)
+            .map_err(|e| Error::whisper_error("failed to get start timestamp", e))?;
+        let end_timestamp = 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);
+        segments.push(Segment { start_timestamp, end_timestamp, text: segment });
+    }
 
-}
+    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 {
+            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);
+            }
+        }
+    }
+
+    tracing::trace!("took {}ms", (et - st).as_millis());
+    Ok(segments)
+}
+
+impl Drop for WhisperHandler {
+    fn drop(&mut self) {
+        let Some(stop_handle) = self.stop_handle.take() else {
+            return tracing::warn!("WhisperHandler::drop() called without stop_handle");
+        };
+        if let Err(_) = stop_handle.send(()) {
+            tracing::warn!("WhisperHandler::drop() failed to send stop signal");
+        }
+    }
+}