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

21
crates/transcription/Cargo.toml
View File

@@ -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).

7
crates/transcription/src/lib.rs
View File

@@ -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;

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());
}
}

61
crates/transcription/src/transcriber.rs Normal file
View File

@@ -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<dyn Transcriber + Send>` 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<Vec<Segment>>;
}
#[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<Box<dyn Transcriber + Send>> = None;
}
}

31
crates/transcription/src/whisper_rs_backend.rs
View File

@@ -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<Self, WhisperBackendError> {
pub fn load(model_path: &Path) -> std::result::Result<Self, WhisperBackendError> {
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<Vec<Segment>, WhisperBackendError> {
) -> Result<Vec<Segment>> {
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;