Context

Cargo.lock

@@ -200,9 +200,9 @@ dependencies = [
 
 [[package]]
 name = "aws-sdk-polly"
-version = "0.33.0"
+version = "0.34.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4a4157655fb2b0d28bb4b838e0b0b111ea14a19d53d8b406d80b9528eb3c1eae"
+checksum = "6950a413936042ea34341331eb2fe7635f33d760ffff3dd11afa4ad35d5151a5"
 dependencies = [
  "aws-credential-types",
  "aws-http",
@@ -273,9 +273,9 @@ dependencies = [
 
 [[package]]
 name = "aws-sdk-transcribestreaming"
-version = "0.33.0"
+version = "0.34.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e004a8f20dfcca5e4548a5d071049bf704184c702832f78748588727f229336d"
+checksum = "bf1b3a068ad56440585dccc3210a537bbe9653b07897698339e0d04c234463d2"
 dependencies = [
  "aws-credential-types",
  "aws-http",
@@ -299,9 +299,9 @@ dependencies = [
 
 [[package]]
 name = "aws-sdk-translate"
-version = "0.33.0"
+version = "0.34.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fc4a37674add1d56a2e0dc7c5cbe41b8db4f352f117dc4ef5c892f3d58344d5b"
+checksum = "23f27dd4ee2f7aeb6736cb24bad9f0ef0bb08194bd598d8fa8e96261c6689409"
 dependencies = [
  "aws-credential-types",
  "aws-http",

Cargo.toml

@@ -5,25 +5,25 @@ edition = "2021"
 
 [dependencies]
 aws-config= { version = "0.56.1" }
-aws-sdk-transcribestreaming= { version = "0.33.0" }
-aws-sdk-translate = "0.33.0"
-aws-sdk-polly = "0.33.0"
-clap = { version = "4.4.6" , features = ["derive"]}
-tokio = {  version = "1.33.0" , features = ["full"] }
-tokio-stream = "0.1.14"
-async-stream = "0.3.5"
-futures-util = "0.3.28"
-serde = { version = "1.0.189", features = ["derive"] }
-serde_json = { version = "1.0.107" }
-serde_yaml = "0.9.25"
-whisper-rs = {  version = "0.8.0" }
-whisper-rs-sys = "0.6.1"
-tracing = "0.1.40"
-tracing-subscriber = "0.3.17"
-lazy_static = "1.4.0"
+aws-sdk-transcribestreaming= { version = "0.34" }
+aws-sdk-translate = "0.34"
+aws-sdk-polly = "0.34"
+clap = { version = "4.4" , features = ["derive"]}
+tokio = {  version = "1.33" , features = ["full"] }
+tokio-stream = "0.1"
+async-stream = "0.3"
+futures-util = "0.3"
+serde = { version = "1.0", features = ["derive"] }
+serde_json = { version = "1.0" }
+serde_yaml = "0.9"
+whisper-rs = {  version = "0.8" }
+whisper-rs-sys = "0.6"
+tracing = "0.1"
+tracing-subscriber = "0.3"
+lazy_static = "1.4"
 
 [dependencies.poem]
-version = "1.3.58"
+version = "1.3"
 features = ["websocket", "static-files"]
 
 [target.aarch64-apple-darwin.dependencies.whisper-rs]

config.yaml

@@ -2,20 +2,24 @@ server:
     port: 8080
     host: ::1
 whisper:
-#    n_threads: 4
-    length_ms: 5000
+    length_ms: 10000
     keep_ms: 200
     step_ms: 5000
-    max_tokens: 0
-    audio_ctx: 0
-    vad_thold: 0.6
-    freq_thold: 100.0
-    speed_up: false
-    translate: false
-    no_fallback: false
-    print_special: false
-    no_context: false
-    no_timestamps: false
-    tinydiarize: false
-    language: "en"
-    model: "models/ggml-base.bin"
+    model: "models/ggml-base.bin"
+    max_prompt_tokens: 128
+    params:
+        #n_threads: 4
+        max_tokens: 0
+        audio_ctx: 0
+        speed_up: false
+        single_segment: false
+        translate: false
+        no_fallback: false
+        temperature_inc: -1.0
+        print_special: false
+        print_progress: false
+        print_realtime: false
+        no_context: false
+        no_timestamps: false
+        tinydiarize: false
+        language: "en"

src/config.rs

@@ -22,24 +22,31 @@ lazy_static! {
 }
 
 #[derive(Debug, Deserialize, Clone)]
-pub(crate) struct WhisperParams {
-    pub(crate) n_threads: Option<usize>,
+pub(crate) struct WhisperConfig {
+    pub(crate) params: WhisperParams,
     pub(crate) step_ms: u32,
     pub(crate) length_ms: u32,
     pub(crate) keep_ms: u32,
+    pub(crate) model: String,
+    pub(crate) max_prompt_tokens: usize,
+}
+
+#[derive(Debug, Deserialize, Clone)]
+pub(crate) struct WhisperParams {
+    pub(crate) n_threads: Option<usize>,
     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) print_realtime: bool,
+    pub(crate) print_progress: bool,
     pub(crate) no_timestamps: bool,
+    pub(crate) temperature_inc: f32,
+    pub(crate) single_segment: bool,
     // pub(crate) tinydiarize: bool,
     pub(crate) language: Option<String>,
-    pub(crate) model: String,
 }
 
 const NONE: [c_int;0] = [];
@@ -47,12 +54,12 @@ const NONE: [c_int;0] = [];
 impl WhisperParams {
     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_progress(self.print_progress);
         param.set_print_special(self.print_special);
-        param.set_print_realtime(false);
+        param.set_print_realtime(self.print_realtime);
         param.set_print_timestamps(!self.no_timestamps);
         param.set_translate(self.translate);
-        param.set_single_segment(true);
+        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);
@@ -61,14 +68,7 @@ impl WhisperParams {
         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 {
-            param.set_temperature_inc(-1.0);
-        }
-        if self.no_context {
-            param.set_tokens(&NONE);
-        } else {
-            param.set_tokens(&tokens);
-        }
+        param.set_temperature_inc(self.temperature_inc);
 
         param
     }
@@ -82,7 +82,7 @@ pub(crate) struct Server {
 
 #[derive(Debug, Deserialize)]
 pub struct Config {
-    pub(crate) whisper: WhisperParams,
+    pub(crate) whisper: WhisperConfig,
     pub(crate) server: Server,
 }
 

src/whisper.rs

@@ -7,7 +7,7 @@ use lazy_static::lazy_static;
 use tokio::sync::{broadcast, mpsc, oneshot};
 use whisper_rs::{convert_integer_to_float_audio, WhisperState, WhisperContext};
 use whisper_rs_sys::WHISPER_SAMPLE_RATE;
-use crate::config::{WhisperParams, CONFIG};
+use crate::config::{WhisperParams, CONFIG, WhisperConfig};
 use crate::group::GroupedWithin;
 
 lazy_static! {
@@ -69,6 +69,7 @@ pub struct Segment {
     pub start_timestamp: i64,
     pub end_timestamp: i64,
     pub text: String,
+    tokens: Vec<c_int>,
 }
 
 pub struct WhisperHandler {
@@ -78,7 +79,7 @@ pub struct WhisperHandler {
 }
 
 impl WhisperHandler {
-    pub(crate) fn new(config: WhisperParams) -> Result<Self, Error> {
+    pub(crate) fn new(config: WhisperConfig) -> Result<Self, Error> {
         let (stop_handle, mut stop_signal) = oneshot::channel();
         let (pcm_tx, pcm_rx) = mpsc::channel::<Vec<u8>>(128);
         let (transcription_tx, _) = broadcast::channel::<Vec<Segment>>(128);
@@ -121,10 +122,11 @@ impl WhisperHandler {
 
                 if tracing::enabled!(tracing::Level::TRACE) {
                     for segment in segments.iter() {
-                        tracing::trace!("[{}] SEGMENT: {}", detector.n_iter, segment.text);
+                        tracing::trace!("[{}-{}]s SEGMENT: {}",
+                            segment.start_timestamp as f32 / 1000.0,
+                            segment.end_timestamp as f32 / 1000.0,
+                            segment.text);
                     }
-                } 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) {
@@ -151,107 +153,114 @@ impl WhisperHandler {
 
 struct Detector {
     state: WhisperState<'static>,
-    config: &'static WhisperParams,
+    config: &'static WhisperConfig,
     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,
+    stable_offset: usize,
 }
 
 impl Detector {
     fn new(state: WhisperState<'static>,
-           config: &'static WhisperParams) -> Self {
+           config: &'static WhisperConfig) -> 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,
+            stable_offset: 0
         }
     }
 
     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();
+        if self.pcm_f32.len() < self.n_samples_len {
+            return
         }
+        let len_to_drain = self.pcm_f32.drain(0..(self.pcm_f32.len() - self.n_samples_len)).len();
+        self.offset += len_to_drain;
     }
 
     fn inference(&mut self) -> Result<Vec<Segment>, Error> {
-        let params = self.config.to_full_params(self.prompt_tokens.as_slice());
+        let params = self.config.params.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());
+            tracing::warn!("full([{}]) took {} ms too slow", self.pcm_f32.len(), (end - start).as_millis());
         }
 
+        let timestamp_offset: i64 = (self.offset * 1000 / WHISPER_SAMPLE_RATE as usize) as i64;
+        let stable_offset: i64 = (self.stable_offset * 1000 / WHISPER_SAMPLE_RATE as usize) as i64;
         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 });
+            if end_timestamp <= stable_offset {
+                continue
+            }
+
+            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 segment = self.state
+                .full_get_segment_text(i)
+                .map_err(|e| Error::whisper_error("failed to get segment", e))?;
+            let num_tokens = self.state
+                .full_n_tokens(i)
+                .map_err(|e| Error::whisper_error("failed to get segment tokens", e))?;
+            let mut segment_tokens = Vec::with_capacity(num_tokens as usize);
+            for j in 0..num_tokens {
+                segment_tokens.push(
+                    self.state
+                        .full_get_token_id(i, j)
+                        .map_err(|e| Error::whisper_error("failed to get token", e))?
+                );
+            }
+
+            segments.push(Segment { start_timestamp, end_timestamp, text: segment.trim().to_string(), tokens: segment_tokens });
         }
 
+        let Some((_last, init)) = segments.split_last() else {
+            return Ok(Vec::default())
+        };
 
-        self.n_iter = self.n_iter + 1;
+        let Some((last_2_seg, _)) = init.split_last() else {
+            return Ok(Vec::default())
+        };
 
-        if self.n_iter % self.n_new_line == 0 {
-            self.next_line()?;
-            Ok(segments)
-        } else {
-            Ok(vec![])
-        }
+        let offset = (last_2_seg.end_timestamp - timestamp_offset) as usize / 1000 * WHISPER_SAMPLE_RATE as usize;
+        self.stable_offset = offset;
+        self.drop_stable_by_segments(init);
+        Ok(init.into())
     }
 
-    fn next_line(&mut self) -> Result<(), Error> {
+    fn drop_stable_by_segments(&mut self, stable_segments: &[Segment]) {
+        let Some(last) = stable_segments.last() else {
+            return
+        };
+        let drop_offset: usize = (last.end_timestamp as usize / 1000 * WHISPER_SAMPLE_RATE as usize - self.offset) as usize;
+        let len_to_drain = self.pcm_f32.drain(0..drop_offset).len();
+        self.offset += len_to_drain;
 
-        // 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();
+        for segment in stable_segments.into_iter() {
+            self.prompt_tokens.extend(&segment.tokens);
         }
-
-        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);
-                }
-            }
+        if self.prompt_tokens.len() > self.config.max_prompt_tokens {
+            let _ = self.prompt_tokens.drain(0..(self.prompt_tokens.len() - self.config.max_prompt_tokens)).len();
         }
-        Ok(())
     }
 }