125 lines
4.1 KiB
Rust
125 lines
4.1 KiB
Rust
//! 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::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}")]
|
|
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) -> 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 — 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>> {
|
|
tracing::info!(
|
|
language = ?options.language,
|
|
has_initial_prompt = options.initial_prompt.as_deref().map(|p| !p.is_empty()).unwrap_or(false),
|
|
"WhisperRsBackend::transcribe_sync entering"
|
|
);
|
|
|
|
let mut state = self.ctx.create_state().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() {
|
|
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| {
|
|
KonError::TranscriptionFailed(
|
|
WhisperBackendError::Transcribe(e.to_string()).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| {
|
|
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;
|
|
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"));
|
|
}
|
|
}
|