Handy/src-tauri/src/managers/transcription.rs

78 lines
2.4 KiB
Rust

use anyhow::Result;
use std::sync::Mutex;
use tauri::{App, Manager};
use whisper_rs::install_whisper_log_trampoline;
use whisper_rs::{
FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters, WhisperState,
};
pub struct TranscriptionManager {
state: Mutex<WhisperState>,
}
impl TranscriptionManager {
pub fn new(app: &App) -> Result<Self> {
let whisper_path = app.path().resolve(
"resources/models/ggml-small.bin",
tauri::path::BaseDirectory::Resource,
)?;
let path = whisper_path
.to_str()
.expect("Path contains invalid UTF-8 Chars");
install_whisper_log_trampoline();
// Load the model
let context = WhisperContext::new_with_params(path, WhisperContextParameters::default())
.map_err(|e| anyhow::anyhow!("Failed to load whisper model: {}", e))?;
// Create state
let state = context.create_state().expect("failed to create state");
Ok(Self {
state: Mutex::new(state),
})
}
pub fn transcribe(&self, audio: Vec<f32>) -> Result<String> {
let st = std::time::Instant::now();
let mut result = String::new();
println!("Audio vector length: {}", audio.len());
if audio.len() == 0 {
println!("Empty audio vector");
// TODO error
return Ok(result);
}
let mut state = self.state.lock().unwrap();
// Initialize parameters
let mut params = FullParams::new(SamplingStrategy::default());
params.set_print_special(false);
params.set_print_progress(false);
params.set_print_realtime(false);
params.set_print_timestamps(false);
params.set_suppress_blank(true);
params.set_suppress_non_speech_tokens(true);
state
.full(params, &audio)
.expect("failed to convert samples");
let num_segments = state
.full_n_segments()
.expect("failed to get number of segments");
for i in 0..num_segments {
let segment = state
.full_get_segment_text(i)
.expect("failed to get segment");
result.push_str(&segment);
}
let et = std::time::Instant::now();
println!("\ntook {}ms", (et - st).as_millis());
Ok(result.trim().to_string())
}
}