refactor(transcription): LocalEngine dispatches SpeechBackend enum — Whisper now on whisper-rs

This commit is contained in:
2026-04-19 20:20:03 +01:00
parent c426fa7eb2
commit 4256383a5b
3 changed files with 53 additions and 44 deletions

View File

@@ -5,7 +5,7 @@ pub mod whisper_rs_backend;
pub use concurrency::run_inference; pub use concurrency::run_inference;
pub use local_engine::{ pub use local_engine::{
load_parakeet, load_whisper, LocalEngine, TimedTranscript, load_parakeet, load_whisper, LocalEngine, SpeechBackend, TimedTranscript,
}; };
pub use transcribe_rs::SpeechModel; pub use transcribe_rs::SpeechModel;
pub use model_manager::{ pub use model_manager::{

View File

@@ -10,17 +10,30 @@ use kon_core::types::{
TranscriptionOptions, TranscriptionOptions,
}; };
use crate::whisper_rs_backend::WhisperRsBackend;
/// Result of a timed transcription: transcript + inference duration. /// Result of a timed transcription: transcript + inference duration.
pub struct TimedTranscript { pub struct TimedTranscript {
pub transcript: Transcript, pub transcript: Transcript,
pub inference_ms: u64, pub inference_ms: u64,
} }
/// Public discriminator selected by the loaders (`load_parakeet`, `load_whisper`)
/// and passed to `LocalEngine::load`. `src-tauri::commands::models` names this
/// type as the return of `load_model_from_disk`, so it must be `pub`.
pub enum SpeechBackend {
/// transcribe-rs-owned model. Used for Parakeet ONNX (wrapped in
/// ParakeetWordGranularity for word-level timestamps).
Adapter(Box<dyn SpeechModel + Send>),
/// Direct whisper-rs. The only path that actually forwards `initial_prompt`.
WhisperRs(WhisperRsBackend),
}
/// Wraps any transcribe-rs engine in Kon's SpeechToText trait. /// Wraps any transcribe-rs engine in Kon's SpeechToText trait.
/// Encapsulates threading: inference always runs on a blocking thread. /// Encapsulates threading: inference always runs on a blocking thread.
/// The rest of the app never imports transcribe-rs directly. /// The rest of the app never imports transcribe-rs directly.
pub struct LocalEngine { pub struct LocalEngine {
engine: Mutex<Option<Box<dyn SpeechModel + Send>>>, engine: Mutex<Option<SpeechBackend>>,
engine_name: EngineName, engine_name: EngineName,
loaded_model_id: Mutex<Option<ModelId>>, loaded_model_id: Mutex<Option<ModelId>>,
} }
@@ -34,10 +47,10 @@ impl LocalEngine {
} }
} }
pub fn load(&self, model: Box<dyn SpeechModel + Send>, model_id: ModelId) { pub fn load(&self, backend: SpeechBackend, model_id: ModelId) {
let mut guard = let mut guard =
self.engine.lock().unwrap_or_else(|e| e.into_inner()); self.engine.lock().unwrap_or_else(|e| e.into_inner());
*guard = Some(model); *guard = Some(backend);
let mut id_guard = self let mut id_guard = self
.loaded_model_id .loaded_model_id
.lock() .lock()
@@ -72,33 +85,37 @@ impl LocalEngine {
) -> Result<TimedTranscript> { ) -> Result<TimedTranscript> {
let mut guard = let mut guard =
self.engine.lock().unwrap_or_else(|e| e.into_inner()); self.engine.lock().unwrap_or_else(|e| e.into_inner());
let engine = let backend = guard.as_mut().ok_or(KonError::EngineNotLoaded)?;
guard.as_mut().ok_or(KonError::EngineNotLoaded)?;
let opts = TranscribeOptions {
language: options.language.clone(),
translate: false,
leading_silence_ms: None,
trailing_silence_ms: None,
};
let start = Instant::now(); let start = Instant::now();
let result: TranscriptionResult = engine let segments: Vec<Segment> = match backend {
.transcribe(audio.samples(), &opts) SpeechBackend::Adapter(model) => {
.map_err(|e| KonError::TranscriptionFailed(e.to_string()))?; let opts = TranscribeOptions {
language: options.language.clone(),
translate: false,
leading_silence_ms: None,
trailing_silence_ms: None,
};
let result: TranscriptionResult = model
.transcribe(audio.samples(), &opts)
.map_err(|e| KonError::TranscriptionFailed(e.to_string()))?;
result
.segments
.unwrap_or_default()
.into_iter()
.map(|s| Segment {
start: s.start as f64,
end: s.end as f64,
text: s.text,
})
.collect()
}
SpeechBackend::WhisperRs(w) => w
.transcribe_sync(audio.samples(), options)
.map_err(|e| KonError::TranscriptionFailed(e.to_string()))?,
};
let inference_ms = start.elapsed().as_millis() as u64; let inference_ms = start.elapsed().as_millis() as u64;
let segments = result
.segments
.unwrap_or_default()
.into_iter()
.map(|s| Segment {
start: s.start as f64,
end: s.end as f64,
text: s.text,
})
.collect();
Ok(TimedTranscript { Ok(TimedTranscript {
transcript: Transcript::new( transcript: Transcript::new(
segments, segments,
@@ -150,9 +167,7 @@ impl transcribe_rs::SpeechModel for ParakeetWordGranularity {
} }
/// Load a Parakeet model from a directory path. /// Load a Parakeet model from a directory path.
pub fn load_parakeet( pub fn load_parakeet(model_dir: &Path) -> Result<SpeechBackend> {
model_dir: &Path,
) -> Result<Box<dyn SpeechModel + Send>> {
use transcribe_rs::onnx::Quantization; use transcribe_rs::onnx::Quantization;
let model = transcribe_rs::onnx::parakeet::ParakeetModel::load( let model = transcribe_rs::onnx::parakeet::ParakeetModel::load(
model_dir, model_dir,
@@ -163,21 +178,15 @@ pub fn load_parakeet(
"Failed to load Parakeet: {e}" "Failed to load Parakeet: {e}"
)) ))
})?; })?;
Ok(Box::new(ParakeetWordGranularity(model))) Ok(SpeechBackend::Adapter(Box::new(ParakeetWordGranularity(model))))
} }
/// Load a Whisper model from a GGML file path. /// Load a Whisper model from a GGML file path via whisper-rs.
pub fn load_whisper( pub fn load_whisper(model_path: &Path) -> Result<SpeechBackend> {
model_path: &Path, let backend = WhisperRsBackend::load(model_path).map_err(|e| {
) -> Result<Box<dyn SpeechModel + Send>> { KonError::TranscriptionFailed(format!("Failed to load Whisper: {e}"))
let engine = })?;
transcribe_rs::whisper_cpp::WhisperEngine::load(model_path) Ok(SpeechBackend::WhisperRs(backend))
.map_err(|e| {
KonError::TranscriptionFailed(format!(
"Failed to load Whisper: {e}"
))
})?;
Ok(Box::new(engine))
} }
#[cfg(test)] #[cfg(test)]

View File

@@ -72,7 +72,7 @@ fn model_capability(
} }
} }
pub fn load_model_from_disk(model_id: &ModelId) -> Result<Box<dyn kon_transcription::SpeechModel + Send>, String> { pub fn load_model_from_disk(model_id: &ModelId) -> Result<kon_transcription::SpeechBackend, String> {
let entry = model_registry::find_model(model_id) let entry = model_registry::find_model(model_id)
.ok_or_else(|| format!("Unknown model: {model_id}"))?; .ok_or_else(|| format!("Unknown model: {model_id}"))?;