diff --git a/src-tauri/src/transcription/engine.rs b/src-tauri/src/transcription/engine.rs index 6008ead..a7da5e6 100644 --- a/src-tauri/src/transcription/engine.rs +++ b/src-tauri/src/transcription/engine.rs @@ -1,5 +1,6 @@ use std::path::Path; +use sherpa_rs::transducer::{TransducerConfig, TransducerRecognizer}; use sherpa_rs::whisper::{WhisperConfig, WhisperRecognizer}; /// Result of one ASR call: recognized text plus the detected language @@ -59,6 +60,63 @@ impl AsrEngine for WhisperEngine { } } +/// NVIDIA Parakeet-TDT-v3 (int8) via sherpa-onnx's offline transducer recognizer. +pub struct ParakeetEngine { + recognizer: TransducerRecognizer, +} + +impl ParakeetEngine { + pub fn load(model_dir: &Path) -> Result { + let recognizer = TransducerRecognizer::new(TransducerConfig { + encoder: model_dir + .join("encoder.int8.onnx") + .to_string_lossy() + .to_string(), + decoder: model_dir + .join("decoder.int8.onnx") + .to_string_lossy() + .to_string(), + joiner: model_dir + .join("joiner.int8.onnx") + .to_string_lossy() + .to_string(), + tokens: model_dir + .join("tokens.txt") + .to_string_lossy() + .to_string(), + model_type: "nemo_transducer".to_string(), + num_threads: 2, + sample_rate: 16_000, + feature_dim: 80, + decoding_method: "greedy_search".to_string(), + provider: Some("cpu".to_string()), + ..Default::default() + }) + .map_err(|err| format!("failed to initialize parakeet recognizer: {err}"))?; + + Ok(Self { recognizer }) + } +} + +impl AsrEngine for ParakeetEngine { + fn transcribe(&mut self, sample_rate: u32, samples: &[f32]) -> AsrOutput { + // The transducer recognizer returns text only; Parakeet reports no language. + AsrOutput { + text: self.recognizer.transcribe(sample_rate, samples), + language: String::new(), + } + } +} + +/// Build the configured ASR engine. `engine` mirrors `meetings.asr_engine` +/// ("whisper" | "parakeet"); `model_dir` must be that engine's model directory. +pub fn load_engine(engine: &str, model_dir: &Path) -> Result, String> { + match engine { + "parakeet" => Ok(Box::new(ParakeetEngine::load(model_dir)?)), + _ => Ok(Box::new(WhisperEngine::load(model_dir)?)), + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src-tauri/src/transcription/mod.rs b/src-tauri/src/transcription/mod.rs index 6b9f187..5bdb67d 100644 --- a/src-tauri/src/transcription/mod.rs +++ b/src-tauri/src/transcription/mod.rs @@ -108,6 +108,7 @@ pub fn start_transcription_worker( let config = worker::WorkerConfig { model_dir: model::whisper_turbo_model_dir(data_dir.as_path()), + asr_engine: "whisper".to_string(), vad_model: vad_model.to_string_lossy().to_string(), recording_start_ms: 0, result_tx, diff --git a/src-tauri/src/transcription/worker.rs b/src-tauri/src/transcription/worker.rs index 1ab8749..890ed9d 100644 --- a/src-tauri/src/transcription/worker.rs +++ b/src-tauri/src/transcription/worker.rs @@ -5,7 +5,7 @@ use std::time::Duration; use sherpa_rs::silero_vad::{SileroVad, SileroVadConfig}; -use super::engine::{AsrEngine, WhisperEngine}; +use super::engine::{self, AsrEngine}; use super::resampler::{AudioResampler, RingAccumulator}; use super::{SegmentResult, WorkerCommand}; @@ -15,6 +15,7 @@ const WHISPER_MAX_DECODE_SAMPLES: usize = (ASR_SAMPLE_RATE as usize) * WHISPER_M pub struct WorkerConfig { pub model_dir: PathBuf, + pub asr_engine: String, pub vad_model: String, pub recording_start_ms: u64, pub result_tx: mpsc::Sender, @@ -117,10 +118,10 @@ pub fn run_worker( config.model_dir.display() ); - let mut engine: Box = match WhisperEngine::load(&config.model_dir) { + let mut engine = match engine::load_engine(&config.asr_engine, &config.model_dir) { Ok(engine) => { - eprintln!("[transcription] whisper turbo loaded OK"); - Box::new(engine) + eprintln!("[transcription] ASR engine '{}' loaded OK", config.asr_engine); + engine } Err(err) => { eprintln!("{err}");