update --stream-* defaults

src/main.rs

@@ -5,6 +5,8 @@ use std::time::{Duration, Instant};
 use chrono;
 use clap::Parser;
 use std::thread;
+use std::fs::File;
+use std::io::Write;
 
 #[derive(Parser, Debug)]
 struct Flags {
@@ -14,18 +16,29 @@ struct Flags {
    #[arg(long, default_value = "8")]
    threads: i32,
 
-   #[arg(long, default_value = "1.0")]
-   stream_retain: f32,
    #[arg(long, default_value = "5")]
    stream_step: u64,
+   #[arg(long, default_value = "0.6")]
+   stream_retain: f32,
+   #[arg(long, default_value = "0.3")]
+   stream_head: f32,
+   #[arg(long, default_value = "0.3")]
+   stream_tail: f32,
 
    wav: Option<String>,
+
+   #[arg(long, default_value = "false")]
+   debug: bool,
 }
 
 fn main() {
    let flags = Flags::parse();
 
-   let w = new_whisper(flags.model, flags.threads, Handler{}).unwrap();
+   let handler = Handler{
+      head: flags.stream_head,
+      tail: flags.stream_tail,
+   };
+   let w = new_whisper(flags.model, flags.threads, handler).unwrap();
    let stream_retain = (flags.stream_retain * 16_000.0) as usize;
    let stream_step = Duration::new(flags.stream_step, 0);
    match flags.wav {
@@ -41,6 +54,10 @@ fn main() {
          w.transcribe(&audio_data);
       },
       None => {
+         match &flags.debug {
+            true => { File::create("/tmp/page.rawf32audio").unwrap(); },
+            false => {},
+         };
          let mut buffer = vec![];
          let mut last = Instant::now();
          new_listener().listen(move |data: Vec<f32>| {
@@ -51,8 +68,22 @@ fn main() {
                   Err(msg) => eprintln!("{}", msg),
                };
 
-               for i in stream_retain..buffer.len() {
-                  buffer[i - stream_retain] = buffer[i]
+               match &flags.debug {
+                  true => {
+                     let mut f = File::options().append(true).open("/tmp/page.rawf32audio").unwrap();
+                     let mut wav_data = vec![];
+                     for i in buffer.iter() {
+                        for j in i.to_le_bytes() {
+                           wav_data.push(j);
+                        }
+                     }
+                     f.write_all(wav_data.as_slice()).unwrap();
+                  },
+                  false => {},
+               };
+
+               for i in 0..stream_retain {
+                  buffer[i] = buffer[buffer.len() - stream_retain + i];
                }
                buffer.truncate(stream_retain);
                last = Instant::now();
@@ -139,7 +170,7 @@ impl WhisperEngine {
       }
    }
 
-   fn transcribe(&self, data: &Vec<f32>) -> Result<String, WhisperError> {
+   fn transcribe(&self, data: &Vec<f32>) -> Result<Whispered, WhisperError> {
       let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 0 });
       params.set_no_context(true);
       params.set_n_threads(self.threads);
@@ -154,11 +185,13 @@ impl WhisperEngine {
       let mut state = self.ctx.create_state()?;
       state.full(params, &data[..])?;
 
+      let mut result = new_whispered();
       let num_segments = state.full_n_segments()?;
-      let mut result = "".to_string();
       for i in 0..num_segments {
-         let segment = state.full_get_segment_text(i)?;
-         result = format!("{} {}", result, segment);
+         let data = state.full_get_segment_text(i)?;
+         let start = state.full_get_segment_t0(i)?;
+         let stop = state.full_get_segment_t1(i)?;
+         result.push(data, start, stop);
       }
 
       Ok(result)
@@ -170,12 +203,86 @@ struct AWhisper {
    ack: Option<std::sync::mpsc::SyncSender<bool>>,
 }
 
-struct Handler {}
+#[derive(Clone, Debug)]
+struct Whispered {
+   data: Vec<AWhispered>,
+}
+
+#[derive(Clone, Debug)]
+struct AWhispered {
+   data: String,
+   offset: i64,
+   length: i64,
+}
+
+fn new_whispered() -> Whispered {
+   Whispered{data: vec![]}
+}
+
+fn new_a_whispered(data: String, start: i64, stop: i64) -> AWhispered {
+   AWhispered{
+      data: data,
+      offset: start.clone(),
+      length: stop - start,
+   }
+}
+
+impl Whispered {
+   fn to_string(&self) -> String {
+      let mut result = "".to_string();
+      for i in 0..self.data.len() {
+         result = format!("{} {}", result, &self.data[i].data);
+      }
+      result
+   }
+
+   fn after(&self, t: &f32) -> Whispered {
+      let mut result = new_whispered();
+      self.data
+         .iter()
+         .filter(|x| x.offset as f32 >= *t)
+         .for_each(|x| result.data.push(x.clone()));
+      result
+   }
+
+   fn before(&self, t: &f32) -> Whispered {
+      let mut result = new_whispered();
+      let end = match self.data.iter().map(|x| x.offset + x.length).max() {
+         Some(x) => x,
+         None => 1,
+      };
+      let t = (end as f32) - *t;
+      self.data
+         .iter()
+         .filter(|x| ((x.offset) as f32) <= t)
+         .for_each(|x| result.data.push(x.clone()));
+      result
+   }
+
+   fn push(&mut self, data: String, start: i64, stop: i64) {
+      let words: Vec<_> = data.split_whitespace().collect();
+      let per_word = (stop - start) / (words.len() as i64);
+      for i in 0..words.len() {
+         let start = (i as i64) * per_word;
+         let stop = start.clone() + per_word;
+         self.data.push(new_a_whispered(words[i].to_string(), start, stop));
+      }
+   }
+}
+
+struct Handler {
+   head: f32,
+   tail: f32,
+}
 
 impl Handler {
-   fn on_success(&self, result: String) {
-      eprintln!("{}", chrono::Local::now());
-      println!("{}", result);
+   fn on_success(&self, result: Whispered) {
+      eprintln!("{}: {:?}", chrono::Local::now(), &result);
+      println!("{}", result
+         .after(&(self.head * 100.0))
+         .before(&(self.tail * 100.0))
+         .to_string(),
+      );
    }
    fn on_error(&self, msg: String) {
       eprintln!("error: {}", msg);