diff --git a/src-tauri/src/transcription/engine.rs b/src-tauri/src/transcription/engine.rs new file mode 100644 index 0000000..6008ead --- /dev/null +++ b/src-tauri/src/transcription/engine.rs @@ -0,0 +1,92 @@ +use std::path::Path; + +use sherpa_rs::whisper::{WhisperConfig, WhisperRecognizer}; + +/// Result of one ASR call: recognized text plus the detected language +/// (empty string when the engine does not report one). +pub struct AsrOutput { + pub text: String, + pub language: String, +} + +/// A speech-to-text engine that transcribes 16 kHz mono f32 samples. +/// +/// The VAD -> chunk -> transcribe loop in `worker.rs` is engine-agnostic, so new +/// engines (e.g. Parakeet) only need to implement this trait. +pub trait AsrEngine { + fn transcribe(&mut self, sample_rate: u32, samples: &[f32]) -> AsrOutput; +} + +/// Whisper large-v3-turbo (int8) via sherpa-onnx. +pub struct WhisperEngine { + recognizer: WhisperRecognizer, +} + +impl WhisperEngine { + pub fn load(model_dir: &Path) -> Result { + let recognizer = WhisperRecognizer::new(WhisperConfig { + encoder: model_dir + .join("turbo-encoder.int8.onnx") + .to_string_lossy() + .to_string(), + decoder: model_dir + .join("turbo-decoder.int8.onnx") + .to_string_lossy() + .to_string(), + tokens: model_dir + .join("turbo-tokens.txt") + .to_string_lossy() + .to_string(), + language: "".to_string(), + num_threads: Some(2), + provider: Some("cpu".to_string()), + debug: false, + ..Default::default() + }) + .map_err(|err| format!("failed to initialize whisper recognizer: {err}"))?; + + Ok(Self { recognizer }) + } +} + +impl AsrEngine for WhisperEngine { + fn transcribe(&mut self, sample_rate: u32, samples: &[f32]) -> AsrOutput { + let result = self.recognizer.transcribe(sample_rate, samples); + AsrOutput { + text: result.text, + language: result.lang, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + struct FakeEngine { + text: String, + language: String, + } + + impl AsrEngine for FakeEngine { + fn transcribe(&mut self, _sample_rate: u32, _samples: &[f32]) -> AsrOutput { + AsrOutput { + text: self.text.clone(), + language: self.language.clone(), + } + } + } + + #[test] + fn asr_engine_is_object_safe_and_returns_output() { + let mut engine: Box = Box::new(FakeEngine { + text: "hello world".to_string(), + language: "en".to_string(), + }); + + let out = engine.transcribe(16_000, &[0.0; 16]); + + assert_eq!(out.text, "hello world"); + assert_eq!(out.language, "en"); + } +} diff --git a/src-tauri/src/transcription/mod.rs b/src-tauri/src/transcription/mod.rs index 7eae782..6b9f187 100644 --- a/src-tauri/src/transcription/mod.rs +++ b/src-tauri/src/transcription/mod.rs @@ -1,3 +1,4 @@ +pub mod engine; pub mod model; pub mod resampler; pub mod worker; diff --git a/src-tauri/src/transcription/worker.rs b/src-tauri/src/transcription/worker.rs index f29885a..1ab8749 100644 --- a/src-tauri/src/transcription/worker.rs +++ b/src-tauri/src/transcription/worker.rs @@ -4,8 +4,8 @@ use std::sync::{mpsc, Arc}; use std::time::Duration; use sherpa_rs::silero_vad::{SileroVad, SileroVadConfig}; -use sherpa_rs::whisper::{WhisperConfig, WhisperRecognizer}; +use super::engine::{AsrEngine, WhisperEngine}; use super::resampler::{AudioResampler, RingAccumulator}; use super::{SegmentResult, WorkerCommand}; @@ -22,7 +22,7 @@ pub struct WorkerConfig { fn process_completed_segments( vad: &mut SileroVad, - recognizer: &mut WhisperRecognizer, + engine: &mut dyn AsrEngine, result_tx: &mpsc::Sender, recording_start_ms: u64, ) { @@ -36,13 +36,13 @@ fn process_completed_segments( continue; } - let result = recognizer.transcribe(ASR_SAMPLE_RATE, samples); + let result = engine.transcribe(ASR_SAMPLE_RATE, samples); let text = result.text.trim().to_string(); if text.is_empty() { continue; } - let lang = result.lang.trim(); + let lang = result.language.trim(); let detected_language = if lang.is_empty() { None } else { @@ -112,33 +112,18 @@ pub fn run_worker( }; eprintln!("[transcription] VAD loaded OK"); - let asr_encoder = config.model_dir.join("turbo-encoder.int8.onnx"); - let asr_decoder = config.model_dir.join("turbo-decoder.int8.onnx"); - let asr_tokens = config.model_dir.join("turbo-tokens.txt"); - eprintln!( - "[transcription] loading whisper turbo — encoder={}, decoder={}, tokens={}", - asr_encoder.display(), - asr_decoder.display(), - asr_tokens.display() + "[transcription] loading ASR engine from {}", + config.model_dir.display() ); - let mut recognizer = match WhisperRecognizer::new(WhisperConfig { - encoder: asr_encoder.to_string_lossy().to_string(), - decoder: asr_decoder.to_string_lossy().to_string(), - tokens: asr_tokens.to_string_lossy().to_string(), - language: "".to_string(), - num_threads: Some(2), - provider: Some("cpu".to_string()), - debug: false, - ..Default::default() - }) { - Ok(recognizer) => { + let mut engine: Box = match WhisperEngine::load(&config.model_dir) { + Ok(engine) => { eprintln!("[transcription] whisper turbo loaded OK"); - recognizer + Box::new(engine) } Err(err) => { - eprintln!("failed to initialize whisper recognizer: {err}"); + eprintln!("{err}"); return; } }; @@ -149,7 +134,7 @@ pub fn run_worker( vad.flush(); process_completed_segments( &mut vad, - &mut recognizer, + engine.as_mut(), &config.result_tx, config.recording_start_ms, ); @@ -175,7 +160,7 @@ pub fn run_worker( vad.accept_waveform(chunk_16k); process_completed_segments( &mut vad, - &mut recognizer, + engine.as_mut(), &config.result_tx, config.recording_start_ms, ); @@ -189,7 +174,7 @@ pub fn run_worker( vad.flush(); process_completed_segments( &mut vad, - &mut recognizer, + engine.as_mut(), &config.result_tx, config.recording_start_ms, );