From 8b49d0fe9c31c7d152f9cdba03e1f28141057180 Mon Sep 17 00:00:00 2001 From: Jake Date: Wed, 22 Apr 2026 04:33:23 +0100 Subject: [PATCH] feat(A.2 #13): replace SpeechBackend enum with Transcriber trait MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 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; 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 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. --- crates/transcription/Cargo.toml | 21 +++- crates/transcription/src/lib.rs | 7 +- crates/transcription/src/local_engine.rs | 109 ++++++++++-------- crates/transcription/src/transcriber.rs | 61 ++++++++++ .../transcription/src/whisper_rs_backend.rs | 31 +++-- src-tauri/src/commands/models.rs | 4 +- 6 files changed, 174 insertions(+), 59 deletions(-) create mode 100644 crates/transcription/src/transcriber.rs diff --git a/crates/transcription/Cargo.toml b/crates/transcription/Cargo.toml index f053e45..88c7eae 100644 --- a/crates/transcription/Cargo.toml +++ b/crates/transcription/Cargo.toml @@ -5,6 +5,15 @@ edition = "2021" description = "Speech-to-text engine wrappers, model management, and inference concurrency for Kon" build = "build.rs" +[features] +# Whisper backend (direct whisper-rs, vulkan-accelerated). Default on — +# gating it exists so a future Windows non-AVX2 build, or a cloud-only +# ASR configuration, can drop whisper-rs-sys entirely per brief item +# #13. Disabling this feature also drops the WhisperRsBackend module +# and the load_whisper entry point. +default = ["whisper"] +whisper = ["dep:whisper-rs", "dep:num_cpus"] + [dependencies] kon-core = { path = "../core" } @@ -21,10 +30,16 @@ futures-util = "0.3" # Download integrity verification sha2 = "0.10" -whisper-rs = { version = "0.16", default-features = false, features = ["vulkan"] } +# Gated behind the `whisper` feature (see [features] above). +whisper-rs = { version = "0.16", default-features = false, features = ["vulkan"], optional = true } -# Direct whisper-rs backend (WhisperRsBackend): thread pool sizing + typed errors. -num_cpus = "1" +# Direct whisper-rs backend (WhisperRsBackend): thread pool sizing. +# Gated alongside whisper-rs since no other code in this crate needs it. +num_cpus = { version = "1", optional = true } + +# Typed error enum used by WhisperRsBackend + elsewhere. Kept +# unconditional because it is a derive-macro crate with negligible +# build cost. thiserror = "2" # Structured logging at backend boundaries (observability for initial_prompt flow). diff --git a/crates/transcription/src/lib.rs b/crates/transcription/src/lib.rs index 00b4b32..83cf575 100644 --- a/crates/transcription/src/lib.rs +++ b/crates/transcription/src/lib.rs @@ -1,9 +1,14 @@ pub mod concurrency; pub mod local_engine; pub mod model_manager; +pub mod transcriber; +#[cfg(feature = "whisper")] pub mod whisper_rs_backend; pub use concurrency::run_inference; -pub use local_engine::{load_parakeet, load_whisper, LocalEngine, SpeechBackend, TimedTranscript}; +pub use local_engine::{load_parakeet, LocalEngine, SpeechModelAdapter, TimedTranscript}; +#[cfg(feature = "whisper")] +pub use local_engine::load_whisper; pub use model_manager::{download, is_downloaded, list_downloaded, model_dir, models_dir}; +pub use transcriber::{Transcriber, TranscriberCapabilities}; pub use transcribe_rs::SpeechModel; diff --git a/crates/transcription/src/local_engine.rs b/crates/transcription/src/local_engine.rs index 3803ac5..fc17877 100644 --- a/crates/transcription/src/local_engine.rs +++ b/crates/transcription/src/local_engine.rs @@ -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), - /// 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); + +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> { + 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`. pub struct LocalEngine { - engine: Mutex>, + engine: Mutex>>, engine_name: EngineName, loaded_model_id: Mutex>, } @@ -46,7 +80,7 @@ impl LocalEngine { } } - pub fn load(&self, backend: SpeechBackend, model_id: ModelId) { + pub fn load(&self, backend: Box, 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 { + 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 = 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 { +pub fn load_parakeet(model_dir: &Path) -> Result> { 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 { +#[cfg(feature = "whisper")] +pub fn load_whisper(model_path: &Path) -> Result> { 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()); } } diff --git a/crates/transcription/src/transcriber.rs b/crates/transcription/src/transcriber.rs new file mode 100644 index 0000000..3caa88f --- /dev/null +++ b/crates/transcription/src/transcriber.rs @@ -0,0 +1,61 @@ +//! Engine-abstraction trait for speech-to-text backends. +//! +//! Replaces the previous `SpeechBackend` enum so new backends +//! (Moonshine, whisper-rs forks, cloud ASR shims, Windows non-AVX2 +//! fallbacks) can drop in without adding a match arm in `LocalEngine`. +//! +//! Concrete implementers today: `SpeechModelAdapter` (wraps any +//! `transcribe-rs` model, currently used for Parakeet) and — behind the +//! `whisper` feature — `WhisperRsBackend` (direct whisper-rs, the only +//! path that pipes `initial_prompt`). + +use kon_core::error::Result; +use kon_core::types::{Segment, TranscriptionOptions}; + +/// Static capabilities a `Transcriber` advertises to callers. +/// +/// `sample_rate` is load-bearing for the progressive WAV writer (#19) +/// which writes live capture samples to disk at the transcriber's +/// native rate. `supports_initial_prompt` lets the Settings surface +/// hide the initial-prompt field for backends that ignore it (Parakeet +/// today). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TranscriberCapabilities { + pub sample_rate: u32, + pub channels: u16, + pub supports_initial_prompt: bool, +} + +/// Unified interface for speech-to-text backends. +/// +/// `Send` is a supertrait so `Box` travels +/// across `spawn_blocking` boundaries without a per-site bound. All +/// inference is synchronous — async callers wrap a `tokio::spawn_blocking` +/// around `transcribe_sync`. +pub trait Transcriber: Send { + fn capabilities(&self) -> TranscriberCapabilities; + + /// Synchronously transcribe 16 kHz mono f32 PCM (or whatever the + /// backend's `capabilities().sample_rate` declares). `&mut self` so + /// backends that keep per-call scratch state (whisper-rs's + /// `WhisperState`, Parakeet's decoder buffers) can mutate them + /// without interior-mutability gymnastics. + fn transcribe_sync( + &mut self, + samples: &[f32], + options: &TranscriptionOptions, + ) -> Result>; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn transcriber_trait_is_object_safe() { + // Compile-time witness: if the trait stops being object-safe + // (e.g. someone adds a generic method or a Self-returning + // method) this declaration fails to build. No runtime work. + let _: Option> = None; + } +} diff --git a/crates/transcription/src/whisper_rs_backend.rs b/crates/transcription/src/whisper_rs_backend.rs index df80c32..b15b30a 100644 --- a/crates/transcription/src/whisper_rs_backend.rs +++ b/crates/transcription/src/whisper_rs_backend.rs @@ -10,8 +10,11 @@ use std::path::Path; use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters}; +use kon_core::error::{KonError, Result}; use kon_core::types::{Segment, TranscriptionOptions}; +use crate::transcriber::{Transcriber, TranscriberCapabilities}; + #[derive(Debug, thiserror::Error)] pub enum WhisperBackendError { #[error("whisper-rs load failed: {0}")] @@ -27,20 +30,32 @@ pub struct WhisperRsBackend { } impl WhisperRsBackend { - pub fn load(model_path: &Path) -> Result { + pub fn load(model_path: &Path) -> std::result::Result { let ctx = WhisperContext::new_with_params(model_path, WhisperContextParameters::default()) .map_err(|e| WhisperBackendError::Load(e.to_string()))?; Ok(Self { ctx }) } +} + +impl Transcriber for WhisperRsBackend { + fn capabilities(&self) -> TranscriberCapabilities { + TranscriberCapabilities { + sample_rate: kon_core::constants::WHISPER_SAMPLE_RATE, + channels: 1, + supports_initial_prompt: true, + } + } /// Synchronously transcribe 16 kHz mono f32 PCM. /// - /// `options.initial_prompt` is piped directly to whisper-rs. - pub fn transcribe_sync( - &self, + /// `options.initial_prompt` is piped directly to whisper-rs — this + /// is the only backend path that honours it; `SpeechModelAdapter` + /// discards it (Parakeet has no equivalent). + fn transcribe_sync( + &mut self, samples: &[f32], options: &TranscriptionOptions, - ) -> Result, WhisperBackendError> { + ) -> Result> { tracing::info!( language = ?options.language, has_initial_prompt = options.initial_prompt.as_deref().map(|p| !p.is_empty()).unwrap_or(false), @@ -50,7 +65,7 @@ impl WhisperRsBackend { let mut state = self .ctx .create_state() - .map_err(|e| WhisperBackendError::State(e.to_string()))?; + .map_err(|e| KonError::TranscriptionFailed(WhisperBackendError::State(e.to_string()).to_string()))?; let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 1 }); if let Some(lang) = options.language.as_deref() { @@ -70,7 +85,7 @@ impl WhisperRsBackend { state .full(params, samples) - .map_err(|e| WhisperBackendError::Transcribe(e.to_string()))?; + .map_err(|e| KonError::TranscriptionFailed(WhisperBackendError::Transcribe(e.to_string()).to_string()))?; let n = state.full_n_segments(); @@ -81,7 +96,7 @@ impl WhisperRsBackend { }; let text = seg .to_str() - .map_err(|e| WhisperBackendError::Transcribe(e.to_string()))? + .map_err(|e| KonError::TranscriptionFailed(WhisperBackendError::Transcribe(e.to_string()).to_string()))? .to_string(); // whisper-rs timestamps are centiseconds (10ms units). Convert to seconds (f64). let start = seg.start_timestamp() as f64 * 0.01; diff --git a/src-tauri/src/commands/models.rs b/src-tauri/src/commands/models.rs index 5c5eaae..698b988 100644 --- a/src-tauri/src/commands/models.rs +++ b/src-tauri/src/commands/models.rs @@ -9,7 +9,7 @@ use kon_core::hardware::{self, CpuFeatures}; use kon_core::model_registry::{self, Engine, LanguageSupport, ModelEntry}; use kon_core::types::{AudioSamples, ModelId, TranscriptionOptions}; use kon_transcription::model_manager; -use kon_transcription::{load_parakeet, load_whisper, LocalEngine}; +use kon_transcription::{load_parakeet, load_whisper, LocalEngine, Transcriber}; /// Map legacy size strings to ModelId. fn whisper_model_id(size: &str) -> ModelId { @@ -73,7 +73,7 @@ fn model_capability( pub fn load_model_from_disk( model_id: &ModelId, -) -> Result { +) -> Result, String> { let entry = model_registry::find_model(model_id).ok_or_else(|| format!("Unknown model: {model_id}"))?;