This commit is contained in:
Bel LaPointe
2023-12-20 09:28:47 -05:00
parent a617d2a924
commit 32d8080251
6 changed files with 1102 additions and 46 deletions

View File

@@ -7,7 +7,7 @@ use std::thread;
use std::fs::File;
use std::io::Write;
#[derive(Parser, Debug)]
#[derive(Parser, Debug, Clone)]
pub struct Flags {
#[arg(long, default_value = "../models/ggml-tiny.en.bin")]
pub model_path: Option<String>,
@@ -26,14 +26,38 @@ pub struct Flags {
#[arg(long, default_value = "0.3")]
pub stream_tail: f32,
#[arg(long, default_value = None)]
pub wav: Option<String>,
#[arg(long, default_value = "false")]
pub debug: bool,
#[arg(long, default_value = None)]
pub wav: Option<String>,
}
pub fn main<F>(flags: Flags, handler_fn: F) where F: FnMut(Result<Transcribed, String>) + Send + 'static {
main_wav(flags.clone(), handler_fn, flags.wav.unwrap());
}
pub fn main_wav<F>(flags: Flags, handler_fn: F, wav_path: String) where F: FnMut(Result<Transcribed, String>) + Send + 'static {
let w = new_service(
flags.model_path,
flags.model_buffer,
flags.threads,
flags.stream_head,
flags.stream_tail,
handler_fn,
).unwrap();
let (header, data) = wav::read(
&mut std::fs::File::open(wav_path).expect("failed to open $WAV"),
).expect("failed to decode $WAV");
assert!(header.channel_count == 1);
assert!(header.sampling_rate == 16_000);
let data16 = data.as_sixteen().expect("wav is not 32bit floats");
let audio_data = &whisper_rs::convert_integer_to_float_audio(&data16);
w.transcribe(&audio_data);
}
pub fn main_stream<F>(flags: Flags, handler_fn: F, _stream: std::sync::mpsc::Receiver<Vec<f32>>) where F: FnMut(Result<Transcribed, String>) + Send + 'static {
let w = new_service(
flags.model_path,
flags.model_buffer,
@@ -44,53 +68,38 @@ pub fn main<F>(flags: Flags, handler_fn: F) where F: FnMut(Result<Transcribed, S
).unwrap();
let stream_retain = (flags.stream_retain * 16_000.0) as usize;
let stream_step = Duration::new(flags.stream_step, 0);
match flags.wav {
Some(wav) => {
let (header, data) = wav::read(
&mut std::fs::File::open(wav).expect("failed to open $WAV"),
).expect("failed to decode $WAV");
assert!(header.channel_count == 1);
assert!(header.sampling_rate == 16_000);
let data16 = data.as_sixteen().expect("wav is not 32bit floats");
let audio_data = &whisper_rs::convert_integer_to_float_audio(&data16);
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>| {
data.iter().for_each(|x| buffer.push(*x));
if Instant::now() - last > stream_step {
w.transcribe_async(&buffer).unwrap();
w.transcribe(&audio_data);
},
None => {
match &flags.debug {
true => { File::create("/tmp/page.rawf32audio").unwrap(); },
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 => {},
};
let mut buffer = vec![];
let mut last = Instant::now();
new_listener().listen(move |data: Vec<f32>| {
data.iter().for_each(|x| buffer.push(*x));
if Instant::now() - last > stream_step {
w.transcribe_async(&buffer).unwrap();
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();
}
});
},
};
for i in 0..stream_retain {
buffer[i] = buffer[buffer.len() - stream_retain + i];
}
buffer.truncate(stream_retain);
last = Instant::now();
}
});
}
struct Service {