From bef9481ee0e1e549c7c44ec9a4c412988e9e0872 Mon Sep 17 00:00:00 2001 From: Grzegorz <19194188+farce1@users.noreply.github.com> Date: Sun, 14 Jun 2026 13:58:46 +0200 Subject: [PATCH] feat(transcription): add Parakeet-v3 model download (sherpa-onnx tarball) Track-B step 3. Adds model::parakeet_model_dir + check_parakeet_assets_ready, PARAKEET_V3_ARCHIVE (k2-fsa asr-models tarball), download_parakeet_model (mirrors the whisper-turbo flow: range-resume download + SHA256 verify + tar.bz2 extract + readiness check), the download_parakeet_model command, and its registration. SHA256 is a REPLACE_WITH_ placeholder with a fail-loud test, per the existing maintainer-fill convention. --- src-tauri/src/commands.rs | 14 +++ src-tauri/src/download.rs | 160 ++++++++++++++++++++++++++- src-tauri/src/lib.rs | 1 + src-tauri/src/transcription/model.rs | 45 ++++++++ 4 files changed, 219 insertions(+), 1 deletion(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 9e111c0..a7a7543 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -620,6 +620,20 @@ pub async fn download_diarization_model( .await } +#[tauri::command] +pub async fn download_parakeet_model( + data_dir: tauri::State<'_, DataDir>, + cancel_flag: tauri::State<'_, download::DownloadCancelFlag>, + on_event: Channel, +) -> Result<(), String> { + crate::download::download_parakeet_model( + on_event, + data_dir.inner().0.clone(), + cancel_flag.inner().clone(), + ) + .await +} + #[tauri::command] pub async fn cancel_download( data_dir: tauri::State<'_, DataDir>, diff --git a/src-tauri/src/download.rs b/src-tauri/src/download.rs index b543464..836ff20 100644 --- a/src-tauri/src/download.rs +++ b/src-tauri/src/download.rs @@ -20,6 +20,9 @@ const VAD_URL: &str = const WHISPER_TURBO_URL: &str = "https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/sherpa-onnx-whisper-turbo.tar.bz2"; const WHISPER_TURBO_TMP_NAME: &str = "whisper-turbo-model.tar.bz2.tmp"; +const PARAKEET_V3_URL: &str = + "https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/sherpa-onnx-nemo-parakeet-tdt-0.6b-v3-int8.tar.bz2"; +const PARAKEET_V3_TMP_NAME: &str = "parakeet-v3-model.tar.bz2.tmp"; const DIARIZATION_SEGMENTATION_TMP_NAME: &str = "diarization-segmentation-model.tar.bz2.tmp"; const DIARIZATION_EMBEDDING_TMP_NAME: &str = "nemo_en_titanet_small.onnx.tmp"; const PROGRESS_EMIT_STEP_BYTES: u64 = 512 * 1024; @@ -44,6 +47,13 @@ pub(crate) const WHISPER_TURBO_ARCHIVE: ModelArchive = ModelArchive { uncompressed_size: 0, }; +pub(crate) const PARAKEET_V3_ARCHIVE: ModelArchive = ModelArchive { + url: PARAKEET_V3_URL, + sha256: "REPLACE_WITH_SHA256_PARAKEET_V3", + compressed_size: 0, + uncompressed_size: 0, +}; + pub(crate) const DIARIZATION_SEGMENTATION_ARCHIVE: ModelArchive = ModelArchive { url: diarization_model::SEGMENTATION_ARCHIVE_URL, sha256: "REPLACE_WITH_SHA256_DIARIZATION_SEGMENTATION", @@ -712,13 +722,161 @@ pub async fn download_diarization_model( Ok(()) } +pub async fn download_parakeet_model( + on_event: Channel, + data_dir: PathBuf, + cancel_flag: DownloadCancelFlag, +) -> Result<(), String> { + cancel_flag.store(false, Ordering::SeqCst); + + let models_dir = model::models_dir(data_dir.as_path()); + std::fs::create_dir_all(&models_dir) + .map_err(|_| format!("Failed to create models directory: {}", models_dir.display()))?; + + if model::check_parakeet_assets_ready(data_dir.as_path()) { + let _ = on_event.send(DownloadEvent::Complete); + return Ok(()); + } + + // Disk pre-check (D-19): before any HTTP download begins. + if let Err(err) = check_disk_space(data_dir.as_path(), PARAKEET_V3_ARCHIVE.required_free_space()) + { + let msg = err.message(); + send_extract_error(&on_event, &err); + return Err(msg); + } + + let client = Client::new(); + let total_bytes = content_length(&client, PARAKEET_V3_URL).await; + let archive_tmp = models_dir.join(PARAKEET_V3_TMP_NAME); + + match download_to_file( + &client, + DownloadRequest { + url: PARAKEET_V3_URL, + destination_tmp: &archive_tmp, + total_bytes, + base_downloaded: 0, + on_event: &on_event, + cancel_flag: &cancel_flag, + resumable: true, + }, + ) + .await + { + Ok(_) => {} + Err(err) if err == "range_not_satisfiable" => { + cleanup_tmp(&archive_tmp); + match download_to_file( + &client, + DownloadRequest { + url: PARAKEET_V3_URL, + destination_tmp: &archive_tmp, + total_bytes, + base_downloaded: 0, + on_event: &on_event, + cancel_flag: &cancel_flag, + resumable: false, + }, + ) + .await + { + Ok(_) => {} + Err(err) if err == "cancelled" => { + cleanup_tmp(&archive_tmp); + let _ = on_event.send(DownloadEvent::Cancelled); + return Err(err); + } + Err(err) => { + cleanup_tmp(&archive_tmp); + send_error(&on_event, "Unable to download Parakeet model archive. Please retry."); + return Err(err); + } + } + } + Err(err) if err == "cancelled" => { + cleanup_tmp(&archive_tmp); + let _ = on_event.send(DownloadEvent::Cancelled); + return Err(err); + } + Err(err) => { + cleanup_tmp(&archive_tmp); + send_error(&on_event, "Unable to download Parakeet model archive. Please retry."); + return Err(err); + } + } + + let _ = on_event.send(DownloadEvent::Extracting); + + // SHA256 verification (D-18) on the .tmp file, BEFORE extraction. + if let Err(err) = verify_sha256(&archive_tmp, PARAKEET_V3_ARCHIVE.sha256) { + cleanup_tmp(&archive_tmp); + let msg = err.message(); + send_extract_error(&on_event, &err); + return Err(msg); + } + + let archive_path = archive_tmp.clone(); + let dst_path = models_dir.clone(); + let extract_result = + tokio::task::spawn_blocking(move || extract_tar_bz2(&archive_path, &dst_path)).await; + + match extract_result { + Ok(Ok(())) => { + cleanup_tmp(&archive_tmp); + } + Ok(Err(err)) => { + cleanup_tmp(&archive_tmp); + let msg = err.message(); + send_extract_error(&on_event, &err); + return Err(msg); + } + Err(join_err) => { + cleanup_tmp(&archive_tmp); + let err = ExtractError::Unknown(std::io::Error::other(format!( + "parakeet extraction task panicked: {join_err}" + ))); + let msg = err.message(); + send_extract_error(&on_event, &err); + return Err(msg); + } + } + + if !model::check_parakeet_assets_ready(data_dir.as_path()) { + send_error( + &on_event, + "Parakeet model files are incomplete after download. Please retry.", + ); + return Err("download finished but parakeet readiness check failed".to_string()); + } + + let _ = on_event.send(DownloadEvent::Complete); + Ok(()) +} + #[cfg(test)] // These tests intentionally compare const placeholder fields against `0` and use string // literals that are constant-foldable. They are a fail-loud maintainer gate per Plan 03 // Task 0 + CONTEXT.md D-17 — the asserts MUST run at test time, not be elided by clippy. #[allow(clippy::absurd_extreme_comparisons, clippy::assertions_on_constants)] mod model_archive_consts_tests { - use super::{DIARIZATION_SEGMENTATION_ARCHIVE, WHISPER_TURBO_ARCHIVE}; + use super::{DIARIZATION_SEGMENTATION_ARCHIVE, PARAKEET_V3_ARCHIVE, WHISPER_TURBO_ARCHIVE}; + + #[test] + fn parakeet_archive_sha256_is_filled_in() { + assert!( + !PARAKEET_V3_ARCHIVE.sha256.contains("REPLACE_WITH_"), + "PARAKEET_V3_ARCHIVE.sha256 still contains REPLACE_WITH_ — maintainer must fill the real lowercase-hex SHA256 of the downloaded archive" + ); + assert_eq!(PARAKEET_V3_ARCHIVE.sha256.len(), 64); + assert!(PARAKEET_V3_ARCHIVE + .sha256 + .chars() + .all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase())); + assert!( + PARAKEET_V3_ARCHIVE.compressed_size > 0 && PARAKEET_V3_ARCHIVE.uncompressed_size > 0, + ); + } #[test] fn whisper_archive_sha256_is_filled_in() { diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index f88ea1f..5708657 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -261,6 +261,7 @@ pub fn run() { commands::rename_speaker, commands::check_diarization_model_ready, commands::download_diarization_model, + commands::download_parakeet_model, commands::get_diarization_data, commands::cancel_download, commands::list_audio_input_devices, diff --git a/src-tauri/src/transcription/model.rs b/src-tauri/src/transcription/model.rs index 68c4468..4fe6746 100644 --- a/src-tauri/src/transcription/model.rs +++ b/src-tauri/src/transcription/model.rs @@ -1,6 +1,7 @@ use std::path::{Path, PathBuf}; pub const WHISPER_TURBO_DIR_NAME: &str = "sherpa-onnx-whisper-turbo"; +pub const PARAKEET_V3_DIR_NAME: &str = "sherpa-onnx-nemo-parakeet-tdt-0.6b-v3-int8"; pub fn models_dir(data_dir: &Path) -> PathBuf { data_dir.join("models") @@ -10,6 +11,10 @@ pub fn whisper_turbo_model_dir(data_dir: &Path) -> PathBuf { models_dir(data_dir).join(WHISPER_TURBO_DIR_NAME) } +pub fn parakeet_model_dir(data_dir: &Path) -> PathBuf { + models_dir(data_dir).join(PARAKEET_V3_DIR_NAME) +} + pub fn vad_model_path(data_dir: &Path) -> PathBuf { models_dir(data_dir).join("silero_vad.onnx") } @@ -26,6 +31,46 @@ pub fn check_transcription_assets_ready(data_dir: &Path) -> bool { .all(|path| path.exists()) } +pub fn check_parakeet_assets_ready(data_dir: &Path) -> bool { + let dir = parakeet_model_dir(data_dir); + + [ + dir.join("encoder.int8.onnx"), + dir.join("decoder.int8.onnx"), + dir.join("joiner.int8.onnx"), + dir.join("tokens.txt"), + ] + .iter() + .all(|path| path.exists()) +} + pub fn check_model_ready(data_dir: &Path) -> bool { check_transcription_assets_ready(data_dir) && vad_model_path(data_dir).exists() } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parakeet_ready_only_when_all_four_files_present() { + let tmp = std::env::temp_dir().join(format!("on-parakeet-ready-{}", std::process::id())); + let dir = parakeet_model_dir(&tmp); + std::fs::create_dir_all(&dir).unwrap(); + let files = [ + "encoder.int8.onnx", + "decoder.int8.onnx", + "joiner.int8.onnx", + "tokens.txt", + ]; + + assert!(!check_parakeet_assets_ready(&tmp)); + for (idx, name) in files.iter().enumerate() { + std::fs::write(dir.join(name), b"x").unwrap(); + let all_present = idx == files.len() - 1; + assert_eq!(check_parakeet_assets_ready(&tmp), all_present); + } + + std::fs::remove_dir_all(&tmp).ok(); + } +}