diff --git a/crates/transcription/Cargo.toml b/crates/transcription/Cargo.toml index 8069037..3587ddf 100644 --- a/crates/transcription/Cargo.toml +++ b/crates/transcription/Cargo.toml @@ -21,3 +21,7 @@ futures-util = "0.3" sha2 = "0.10" whisper-rs = { version = "0.16", default-features = false, features = ["vulkan"] } + +# Direct whisper-rs backend (WhisperRsBackend): thread pool sizing + typed errors. +num_cpus = "1" +thiserror = "2" diff --git a/crates/transcription/src/lib.rs b/crates/transcription/src/lib.rs index 8c297fd..84afc65 100644 --- a/crates/transcription/src/lib.rs +++ b/crates/transcription/src/lib.rs @@ -1,6 +1,7 @@ pub mod concurrency; pub mod local_engine; pub mod model_manager; +pub mod whisper_rs_backend; pub use concurrency::run_inference; pub use local_engine::{ diff --git a/crates/transcription/src/whisper_rs_backend.rs b/crates/transcription/src/whisper_rs_backend.rs new file mode 100644 index 0000000..673a307 --- /dev/null +++ b/crates/transcription/src/whisper_rs_backend.rs @@ -0,0 +1,101 @@ +//! Direct whisper-rs backend. Owns a WhisperContext; each call builds a +//! fresh WhisperState (state can be reused, but fresh-per-call is simpler +//! and matches the transcribe-rs call style we are replacing). +//! +//! Exists because transcribe-rs does not expose set_initial_prompt; this +//! wrapper is the only path that can pipe per-capture vocabulary context +//! into Whisper. + +use std::path::Path; + +use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters}; + +use kon_core::types::{Segment, TranscriptionOptions}; + +#[derive(Debug, thiserror::Error)] +pub enum WhisperBackendError { + #[error("whisper-rs load failed: {0}")] + Load(String), + #[error("whisper-rs state creation failed: {0}")] + State(String), + #[error("whisper-rs transcribe failed: {0}")] + Transcribe(String), +} + +pub struct WhisperRsBackend { + ctx: WhisperContext, +} + +impl WhisperRsBackend { + pub fn load(model_path: &Path) -> Result { + let ctx = WhisperContext::new_with_params( + model_path, + WhisperContextParameters::default(), + ) + .map_err(|e| WhisperBackendError::Load(e.to_string()))?; + Ok(Self { ctx }) + } + + /// Synchronously transcribe 16 kHz mono f32 PCM. + /// + /// `options.initial_prompt` is piped directly to whisper-rs. + pub fn transcribe_sync( + &self, + samples: &[f32], + options: &TranscriptionOptions, + ) -> Result, WhisperBackendError> { + let mut state = self + .ctx + .create_state() + .map_err(|e| WhisperBackendError::State(e.to_string()))?; + + let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 1 }); + if let Some(lang) = options.language.as_deref() { + if !lang.is_empty() { + params.set_language(Some(lang)); + } + } + if let Some(prompt) = options.initial_prompt.as_deref() { + if !prompt.is_empty() { + params.set_initial_prompt(prompt); + } + } + params.set_n_threads(num_cpus::get() as i32); + params.set_print_special(false); + params.set_print_progress(false); + params.set_print_realtime(false); + + state + .full(params, samples) + .map_err(|e| WhisperBackendError::Transcribe(e.to_string()))?; + + let n = state.full_n_segments(); + + let mut out = Vec::with_capacity(n.max(0) as usize); + for i in 0..n { + let Some(seg) = state.get_segment(i) else { + continue; + }; + let text = seg + .to_str() + .map_err(|e| WhisperBackendError::Transcribe(e.to_string()))? + .to_string(); + // whisper-rs timestamps are centiseconds (10ms units). Convert to seconds (f64). + let start = seg.start_timestamp() as f64 * 0.01; + let end = seg.end_timestamp() as f64 * 0.01; + out.push(Segment { start, end, text }); + } + Ok(out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn backend_error_displays() { + let e = WhisperBackendError::Load("oops".into()); + assert!(e.to_string().contains("oops")); + } +}