use std::path::Path; use std::sync::Mutex; use std::time::Instant; use transcribe_rs::{SpeechModel, TranscribeOptions, TranscriptionResult}; use magnotia_core::error::{MagnotiaError, Result}; use magnotia_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. pub struct TimedTranscript { pub transcript: Transcript, pub inference_ms: u64, } /// 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: magnotia_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| MagnotiaError::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()) } } /// 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_name: EngineName, loaded_model_id: Mutex>, } impl LocalEngine { pub fn new(engine_name: EngineName) -> Self { Self { engine: Mutex::new(None), engine_name, loaded_model_id: Mutex::new(None), } } 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 .loaded_model_id .lock() .unwrap_or_else(|e| e.into_inner()); *id_guard = Some(model_id); } /// Drop the loaded model and free its backing resources (GPU VRAM, /// CPU memory, mmap'd GGML tensors). Used by the sequential-GPU /// guard (brief item A.1 #28) so loading the LLM on a tight-VRAM /// system first frees the transcription engine, and vice versa. /// /// No-op when nothing is loaded. Thread-safe — the internal Mutex /// serialises against concurrent transcribe_sync calls. pub fn unload(&self) { let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner()); *guard = None; let mut id_guard = self .loaded_model_id .lock() .unwrap_or_else(|e| e.into_inner()); *id_guard = None; } pub fn name(&self) -> &EngineName { &self.engine_name } pub fn loaded_model_id(&self) -> Option { let guard = self .loaded_model_id .lock() .unwrap_or_else(|e| e.into_inner()); guard.clone() } pub fn is_loaded(&self) -> bool { let guard = self.engine.lock().unwrap_or_else(|e| e.into_inner()); 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( &self, audio: &AudioSamples, options: &TranscriptionOptions, ) -> Result { let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner()); let backend = guard.as_mut().ok_or(MagnotiaError::EngineNotLoaded)?; let start = Instant::now(); let segments = backend.transcribe_sync(audio.samples(), options)?; let inference_ms = start.elapsed().as_millis() as u64; Ok(TimedTranscript { transcript: Transcript::new( segments, options.language.clone().unwrap_or_else(|| "en".to_string()), audio.duration_secs(), ), inference_ms, }) } } /// Thin wrapper over `ParakeetModel` that overrides `transcribe_raw` to /// request word-granularity segments. `transcribe-rs` 0.3's trait impl for /// `ParakeetModel::transcribe_raw` ignores `TranscribeOptions` and uses /// `TimestampGranularity::Token` (per-subword) — which surfaces in Magnotia as /// "T Est Ing . One , Two , Three" output. The concrete-type method /// `ParakeetModel::transcribe_with` accepts `ParakeetParams` with an /// explicit granularity; this wrapper exposes that to the trait object. struct ParakeetWordGranularity(transcribe_rs::onnx::parakeet::ParakeetModel); impl transcribe_rs::SpeechModel for ParakeetWordGranularity { fn capabilities(&self) -> transcribe_rs::ModelCapabilities { self.0.capabilities() } fn default_leading_silence_ms(&self) -> u32 { self.0.default_leading_silence_ms() } fn default_trailing_silence_ms(&self) -> u32 { self.0.default_trailing_silence_ms() } fn transcribe_raw( &mut self, samples: &[f32], options: &TranscribeOptions, ) -> std::result::Result { use transcribe_rs::onnx::parakeet::{ParakeetParams, TimestampGranularity}; let params = ParakeetParams { language: options.language.clone(), timestamp_granularity: Some(TimestampGranularity::Word), }; self.0.transcribe_with(samples, ¶ms) } } /// Load a Parakeet model from a directory path. 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| MagnotiaError::TranscriptionFailed(format!("Failed to load Parakeet: {e}")))?; Ok(Box::new(SpeechModelAdapter(Box::new( ParakeetWordGranularity(model), )))) } /// Load a Whisper model from a GGML file path via whisper-rs. #[cfg(feature = "whisper")] pub fn load_whisper(model_path: &Path) -> Result> { let backend = WhisperRsBackend::load(model_path) .map_err(|e| MagnotiaError::TranscriptionFailed(format!("Failed to load Whisper: {e}")))?; Ok(Box::new(backend)) } #[cfg(test)] mod tests { use super::*; #[test] fn engine_reports_not_available_before_loading() { let engine = LocalEngine::new(EngineName::new("test")); assert!(!engine.is_loaded()); assert!(engine.loaded_model_id().is_none()); assert!(engine.capabilities().is_none()); } }