feat(kon): add transcription crate — transcribe-rs wrapper, model manager
- LocalEngine wraps transcribe-rs SpeechModel behind Kon's own abstraction - load_parakeet() and load_whisper() factory functions - Unified model manager: download with progress callback, check, list, path resolution - Atomic downloads: .part suffix with rename on completion - run_inference() encapsulates spawn_blocking threading - Zero Tauri dependency in this crate (progress via callback, not events) - transcribe-rs v0.3.2 with onnx + whisper-cpp features - 4 tests passing, clippy clean Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -6,3 +6,13 @@ description = "Speech-to-text engine wrappers, model management, and inference c
|
|||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
kon-core = { path = "../core" }
|
kon-core = { path = "../core" }
|
||||||
|
|
||||||
|
# Unified STT engine (Parakeet via ONNX, Whisper via whisper.cpp)
|
||||||
|
transcribe-rs = { version = "0.3", features = ["onnx", "whisper-cpp"] }
|
||||||
|
|
||||||
|
# Async runtime for spawn_blocking
|
||||||
|
tokio = { version = "1", features = ["rt", "sync"] }
|
||||||
|
|
||||||
|
# Model downloads
|
||||||
|
reqwest = { version = "0.12", features = ["stream"] }
|
||||||
|
futures-util = "0.3"
|
||||||
|
|||||||
20
crates/transcription/src/concurrency.rs
Normal file
20
crates/transcription/src/concurrency.rs
Normal file
@@ -0,0 +1,20 @@
|
|||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use kon_core::error::{KonError, Result};
|
||||||
|
use kon_core::types::{AudioSamples, Transcript, TranscriptionOptions};
|
||||||
|
|
||||||
|
use crate::local_engine::LocalEngine;
|
||||||
|
|
||||||
|
/// Runs inference on a blocking thread. Encapsulates all threading concerns.
|
||||||
|
/// Callers never see spawn_blocking — they call this async function.
|
||||||
|
pub async fn run_inference(
|
||||||
|
engine: Arc<LocalEngine>,
|
||||||
|
audio: AudioSamples,
|
||||||
|
options: TranscriptionOptions,
|
||||||
|
) -> Result<Transcript> {
|
||||||
|
tokio::task::spawn_blocking(move || {
|
||||||
|
engine.transcribe_sync(&audio, &options)
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.map_err(|e| KonError::TranscriptionFailed(format!("Task join error: {e}")))?
|
||||||
|
}
|
||||||
@@ -1,2 +1,7 @@
|
|||||||
// kon-transcription: transcribe-rs wrapper behind SpeechToText trait,
|
pub mod concurrency;
|
||||||
// unified model manager, and inference threading.
|
pub mod local_engine;
|
||||||
|
pub mod model_manager;
|
||||||
|
|
||||||
|
pub use concurrency::run_inference;
|
||||||
|
pub use local_engine::{load_parakeet, load_whisper, LocalEngine};
|
||||||
|
pub use model_manager::{download, is_downloaded, list_downloaded, model_dir, models_dir};
|
||||||
|
|||||||
118
crates/transcription/src/local_engine.rs
Normal file
118
crates/transcription/src/local_engine.rs
Normal file
@@ -0,0 +1,118 @@
|
|||||||
|
use std::path::Path;
|
||||||
|
use std::sync::Mutex;
|
||||||
|
|
||||||
|
use transcribe_rs::{SpeechModel, TranscribeOptions, TranscriptionResult};
|
||||||
|
|
||||||
|
use kon_core::error::{KonError, Result};
|
||||||
|
use kon_core::types::{
|
||||||
|
AudioSamples, EngineName, Segment, Transcript, TranscriptionOptions,
|
||||||
|
};
|
||||||
|
|
||||||
|
/// 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.
|
||||||
|
pub struct LocalEngine {
|
||||||
|
engine: Mutex<Option<Box<dyn SpeechModel + Send>>>,
|
||||||
|
engine_name: EngineName,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl LocalEngine {
|
||||||
|
pub fn new(engine_name: EngineName) -> Self {
|
||||||
|
Self {
|
||||||
|
engine: Mutex::new(None),
|
||||||
|
engine_name,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn load(&self, model: Box<dyn SpeechModel + Send>) {
|
||||||
|
let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
|
*guard = Some(model);
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn name(&self) -> &EngineName {
|
||||||
|
&self.engine_name
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn is_loaded(&self) -> bool {
|
||||||
|
let guard = self.engine.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
|
guard.is_some()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run transcription synchronously. Called from within spawn_blocking.
|
||||||
|
pub fn transcribe_sync(
|
||||||
|
&self,
|
||||||
|
audio: &AudioSamples,
|
||||||
|
options: &TranscriptionOptions,
|
||||||
|
) -> Result<Transcript> {
|
||||||
|
let mut guard =
|
||||||
|
self.engine.lock().unwrap_or_else(|e| e.into_inner());
|
||||||
|
let engine = guard
|
||||||
|
.as_mut()
|
||||||
|
.ok_or(KonError::EngineNotLoaded)?;
|
||||||
|
|
||||||
|
let opts = TranscribeOptions {
|
||||||
|
language: options.language.clone(),
|
||||||
|
translate: false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let result: TranscriptionResult = engine
|
||||||
|
.transcribe(audio.samples(), &opts)
|
||||||
|
.map_err(|e| KonError::TranscriptionFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
let segments = result
|
||||||
|
.segments
|
||||||
|
.unwrap_or_default()
|
||||||
|
.into_iter()
|
||||||
|
.map(|s| Segment {
|
||||||
|
start: s.start as f64,
|
||||||
|
end: s.end as f64,
|
||||||
|
text: s.text,
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
Ok(Transcript::new(
|
||||||
|
segments,
|
||||||
|
options
|
||||||
|
.language
|
||||||
|
.clone()
|
||||||
|
.unwrap_or_else(|| "en".to_string()),
|
||||||
|
audio.duration_secs(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Load a Parakeet model from a directory path.
|
||||||
|
/// The directory should contain model_int8.onnx (+ .onnx_data) and tokenizer.json.
|
||||||
|
pub fn load_parakeet(model_dir: &Path) -> Result<Box<dyn SpeechModel + 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(Box::new(model))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Load a Whisper model from a GGML file path.
|
||||||
|
pub fn load_whisper(model_path: &Path) -> Result<Box<dyn SpeechModel + Send>> {
|
||||||
|
let engine = transcribe_rs::whisper_cpp::WhisperEngine::load(model_path)
|
||||||
|
.map_err(|e| {
|
||||||
|
KonError::TranscriptionFailed(format!(
|
||||||
|
"Failed to load Whisper: {e}"
|
||||||
|
))
|
||||||
|
})?;
|
||||||
|
Ok(Box::new(engine))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[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());
|
||||||
|
}
|
||||||
|
}
|
||||||
165
crates/transcription/src/model_manager.rs
Normal file
165
crates/transcription/src/model_manager.rs
Normal file
@@ -0,0 +1,165 @@
|
|||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
|
use kon_core::error::{KonError, Result};
|
||||||
|
use kon_core::model_registry::{find_model, ModelFile};
|
||||||
|
use kon_core::types::{DownloadProgress, ModelId};
|
||||||
|
|
||||||
|
/// Resolve the models storage directory.
|
||||||
|
/// Windows: %LOCALAPPDATA%/kon/models
|
||||||
|
/// Unix: ~/.kon/models
|
||||||
|
pub fn models_dir() -> PathBuf {
|
||||||
|
if cfg!(target_os = "windows") {
|
||||||
|
let local_app_data = std::env::var("LOCALAPPDATA")
|
||||||
|
.unwrap_or_else(|_| ".".to_string());
|
||||||
|
PathBuf::from(local_app_data).join("kon").join("models")
|
||||||
|
} else {
|
||||||
|
dirs_path().join("models")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn dirs_path() -> PathBuf {
|
||||||
|
if cfg!(target_os = "windows") {
|
||||||
|
let local_app_data = std::env::var("LOCALAPPDATA")
|
||||||
|
.unwrap_or_else(|_| ".".to_string());
|
||||||
|
PathBuf::from(local_app_data).join("kon")
|
||||||
|
} else {
|
||||||
|
let home =
|
||||||
|
std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
|
||||||
|
PathBuf::from(home).join(".kon")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the directory path where a specific model's files are stored.
|
||||||
|
pub fn model_dir(id: &ModelId) -> PathBuf {
|
||||||
|
models_dir().join(id.as_str())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check whether all files for a model have been downloaded.
|
||||||
|
pub fn is_downloaded(id: &ModelId) -> bool {
|
||||||
|
let entry = match find_model(id) {
|
||||||
|
Some(e) => e,
|
||||||
|
None => return false,
|
||||||
|
};
|
||||||
|
let dir = model_dir(id);
|
||||||
|
entry.files.iter().all(|f| dir.join(f.filename).exists())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all downloaded model IDs.
|
||||||
|
pub fn list_downloaded() -> Vec<ModelId> {
|
||||||
|
kon_core::model_registry::all_models()
|
||||||
|
.iter()
|
||||||
|
.filter(|m| is_downloaded(&m.id))
|
||||||
|
.map(|m| m.id.clone())
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Download all files for a model, calling the progress callback per chunk.
|
||||||
|
/// Files are downloaded to a .part suffix and atomically renamed on completion.
|
||||||
|
pub async fn download(
|
||||||
|
id: &ModelId,
|
||||||
|
progress: impl Fn(DownloadProgress) + Send + 'static,
|
||||||
|
) -> Result<()> {
|
||||||
|
let entry = find_model(id)
|
||||||
|
.ok_or_else(|| KonError::ModelNotFound(id.clone()))?;
|
||||||
|
|
||||||
|
let dir = model_dir(id);
|
||||||
|
std::fs::create_dir_all(&dir)?;
|
||||||
|
|
||||||
|
for file in &entry.files {
|
||||||
|
let dest = dir.join(file.filename);
|
||||||
|
if dest.exists() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
download_file(file, &dest, id, &progress).await?;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn download_file(
|
||||||
|
file: &ModelFile,
|
||||||
|
dest: &Path,
|
||||||
|
model_id: &ModelId,
|
||||||
|
progress: &(impl Fn(DownloadProgress) + Send),
|
||||||
|
) -> Result<()> {
|
||||||
|
use futures_util::StreamExt;
|
||||||
|
|
||||||
|
let part_path = dest.with_extension(
|
||||||
|
dest.extension()
|
||||||
|
.map(|e| format!("{}.part", e.to_string_lossy()))
|
||||||
|
.unwrap_or_else(|| "part".to_string()),
|
||||||
|
);
|
||||||
|
|
||||||
|
let client = reqwest::Client::builder()
|
||||||
|
.timeout(std::time::Duration::from_secs(30))
|
||||||
|
.build()
|
||||||
|
.map_err(|e| KonError::DownloadFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
let response = client
|
||||||
|
.get(file.url)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|e| KonError::DownloadFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
let total_bytes = response.content_length().unwrap_or(0);
|
||||||
|
let mut stream = response.bytes_stream();
|
||||||
|
let mut downloaded: u64 = 0;
|
||||||
|
let mut last_percent: u8 = 0;
|
||||||
|
|
||||||
|
let mut out = std::fs::File::create(&part_path)?;
|
||||||
|
|
||||||
|
while let Some(chunk) = stream.next().await {
|
||||||
|
let chunk = chunk
|
||||||
|
.map_err(|e| KonError::DownloadFailed(e.to_string()))?;
|
||||||
|
std::io::Write::write_all(&mut out, &chunk)?;
|
||||||
|
downloaded += chunk.len() as u64;
|
||||||
|
|
||||||
|
let percent = if total_bytes > 0 {
|
||||||
|
((downloaded as f64 / total_bytes as f64) * 100.0) as u8
|
||||||
|
} else {
|
||||||
|
0
|
||||||
|
};
|
||||||
|
|
||||||
|
if percent != last_percent {
|
||||||
|
last_percent = percent;
|
||||||
|
progress(DownloadProgress {
|
||||||
|
model_id: model_id.clone(),
|
||||||
|
file_name: file.filename.to_string(),
|
||||||
|
bytes_downloaded: downloaded,
|
||||||
|
total_bytes,
|
||||||
|
percent,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
drop(out);
|
||||||
|
std::fs::rename(&part_path, dest)?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn model_dir_returns_correct_path() {
|
||||||
|
let id = ModelId::new("whisper-tiny-en");
|
||||||
|
let path = model_dir(&id);
|
||||||
|
assert!(path.to_string_lossy().contains("whisper-tiny-en"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn is_downloaded_returns_false_for_missing_model() {
|
||||||
|
let id = ModelId::new("nonexistent-model");
|
||||||
|
assert!(!is_downloaded(&id));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn list_downloaded_returns_empty_when_no_models() {
|
||||||
|
let list = list_downloaded();
|
||||||
|
// In test environment, no models are downloaded
|
||||||
|
// This just verifies the function doesn't panic
|
||||||
|
assert!(list.len() <= kon_core::model_registry::all_models().len());
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user