jake/Lumotia
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(A.2 #13): replace SpeechBackend enum with Transcriber trait

Browse Source 
New crates/transcription/src/transcriber.rs defines a Transcriber
trait (Send supertrait for spawn_blocking travel) with
TranscriberCapabilities (sample_rate, channels, supports_initial_prompt).
TranscriberCapabilities.sample_rate is load-bearing for the upcoming
progressive WAV writer (#19).

Concrete impls:
- SpeechModelAdapter wraps Box<dyn transcribe_rs::SpeechModel + Send>
  for Parakeet (and any future transcribe-rs-backed engine).
- WhisperRsBackend moves its transcribe_sync body into the impl,
  widening the signature from &self to &mut self so per-call
  WhisperState can be created cleanly through the trait object.

LocalEngine now holds Box<dyn Transcriber + Send>; dispatch in
transcribe_sync collapses from a match to a direct call. Adds
LocalEngine::capabilities() for the WAV-writer.

Cargo feature flag "whisper" (default on) makes whisper-rs, num_cpus,
and the whole whisper_rs_backend module optional. cargo check
--no-default-features -p kon-transcription now builds without pulling
whisper-rs-sys — the escape hatch brief item #6 / #13 called for on
Windows / non-AVX2 / cloud-only builds. load_whisper is cfg-gated
behind the same feature.

src-tauri/src/commands/models.rs load_model_from_disk returns
Box<dyn Transcriber + Send> instead of SpeechBackend; caller chain
(ensure_model_loaded, prewarm_default_model) is unchanged.

transcriber_trait_is_object_safe test lands alongside the trait as a
compile-time witness against future Self-returning / generic-method
additions.
This commit is contained in:
Jake
2026-04-22 04:33:23 +01:00
parent dd98cb7994
commit 8b49d0fe9c
6 changed files with 174 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

109
crates/transcription/src/local_engine.rs
View File

@@ -9,6 +9,8 @@ use kon_core::types::{
AudioSamples, EngineName, ModelId, Segment, Transcript, TranscriptionOptions,
};
use crate::transcriber::{Transcriber, TranscriberCapabilities};
#[cfg(feature = "whisper")]
use crate::whisper_rs_backend::WhisperRsBackend;
/// Result of a timed transcription: transcript + inference duration.
@@ -17,22 +19,54 @@ pub struct TimedTranscript {
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),
/// Adapts any `transcribe-rs` `SpeechModel` into the `Transcriber`
/// trait. Today this is only used for Parakeet (ONNX), but the adapter
/// is the path any future transcribe-rs-backed engine plugs through —
/// Moonshine, fine-tuned Parakeet variants, etc.
pub struct SpeechModelAdapter(pub Box<dyn SpeechModel + Send>);
impl Transcriber for SpeechModelAdapter {
fn capabilities(&self) -> TranscriberCapabilities {
TranscriberCapabilities {
sample_rate: kon_core::constants::WHISPER_SAMPLE_RATE,
channels: 1,
supports_initial_prompt: false,
}
}
fn transcribe_sync(
&mut self,
samples: &[f32],
options: &TranscriptionOptions,
) -> Result<Vec<Segment>> {
let opts = TranscribeOptions {
language: options.language.clone(),
translate: false,
leading_silence_ms: None,
trailing_silence_ms: None,
};
let result: TranscriptionResult = self
.0
.transcribe(samples, &opts)
.map_err(|e| KonError::TranscriptionFailed(e.to_string()))?;
Ok(result
.segments
.unwrap_or_default()
.into_iter()
.map(|s| Segment {
start: s.start as f64,
end: s.end as f64,
text: s.text,
})
.collect())
}
}
/// Wraps any transcribe-rs engine in Kon's SpeechToText trait.
/// Encapsulates threading: inference always runs on a blocking thread.
/// The rest of the app never imports transcribe-rs directly.
/// Owns the currently-loaded speech backend and serialises inference
/// against model-swap operations via a `Mutex`. All transcription goes
/// through this struct; no caller ever holds a raw `Box<dyn Transcriber>`.
pub struct LocalEngine {
engine: Mutex<Option<SpeechBackend>>,
engine: Mutex<Option<Box<dyn Transcriber + Send>>>,
engine_name: EngineName,
loaded_model_id: Mutex<Option<ModelId>>,
}
@@ -46,7 +80,7 @@ impl LocalEngine {
}
}
pub fn load(&self, backend: SpeechBackend, model_id: ModelId) {
pub fn load(&self, backend: Box<dyn Transcriber + Send>, model_id: ModelId) {
let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner());
*guard = Some(backend);
let mut id_guard = self
@@ -90,6 +124,14 @@ impl LocalEngine {
guard.is_some()
}
/// Capabilities of the currently-loaded backend. Returns `None`
/// when nothing is loaded. Callers (live capture WAV writer, #19)
/// read sample_rate from here.
pub fn capabilities(&self) -> Option<TranscriberCapabilities> {
let guard = self.engine.lock().unwrap_or_else(|e| e.into_inner());
guard.as_ref().map(|b| b.capabilities())
}
/// Run transcription synchronously with timing.
/// Called from within spawn_blocking.
pub fn transcribe_sync(
@@ -101,32 +143,7 @@ impl LocalEngine {
let backend = guard.as_mut().ok_or(KonError::EngineNotLoaded)?;
let start = Instant::now();
let segments: Vec<Segment> = match backend {
SpeechBackend::Adapter(model) => {
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 segments = backend.transcribe_sync(audio.samples(), options)?;
let inference_ms = start.elapsed().as_millis() as u64;
Ok(TimedTranscript {
@@ -177,20 +194,21 @@ impl transcribe_rs::SpeechModel for ParakeetWordGranularity {
}
/// Load a Parakeet model from a directory path.
pub fn load_parakeet(model_dir: &Path) -> Result<SpeechBackend> {
pub fn load_parakeet(model_dir: &Path) -> Result<Box<dyn Transcriber + Send>> {
use transcribe_rs::onnx::Quantization;
let model = transcribe_rs::onnx::parakeet::ParakeetModel::load(model_dir, &Quantization::Int8)
.map_err(|e| KonError::TranscriptionFailed(format!("Failed to load Parakeet: {e}")))?;
Ok(SpeechBackend::Adapter(Box::new(ParakeetWordGranularity(
model,
Ok(Box::new(SpeechModelAdapter(Box::new(
ParakeetWordGranularity(model),
))))
}
/// Load a Whisper model from a GGML file path via whisper-rs.
pub fn load_whisper(model_path: &Path) -> Result<SpeechBackend> {
#[cfg(feature = "whisper")]
pub fn load_whisper(model_path: &Path) -> Result<Box<dyn Transcriber + Send>> {
let backend = WhisperRsBackend::load(model_path)
.map_err(|e| KonError::TranscriptionFailed(format!("Failed to load Whisper: {e}")))?;
Ok(SpeechBackend::WhisperRs(backend))
Ok(Box::new(backend))
}
#[cfg(test)]
@@ -202,5 +220,6 @@ mod tests {
let engine = LocalEngine::new(EngineName::new("test"));
assert!(!engine.is_loaded());
assert!(engine.loaded_model_id().is_none());
assert!(engine.capabilities().is_none());
}
}