diff --git a/src-tauri/src/models/mod.rs b/src-tauri/src/models/mod.rs index 3d55198..9e5adac 100644 --- a/src-tauri/src/models/mod.rs +++ b/src-tauri/src/models/mod.rs @@ -6,8 +6,34 @@ use serde::{Deserialize, Serialize}; use crate::error::{AppError, AppResult}; -pub const DEFAULT_MODEL_ID: &str = "parakeet-tdt-0.6b-v3-q8_0"; -pub const DEFAULT_MODEL_URL: &str = "https://huggingface.co/cstr/parakeet-tdt-0.6b-v3-GGUF/resolve/main/parakeet-tdt-0.6b-v3-q8_0.gguf"; +/// Modelo padrão usado pelo `SherpaOnnxTranscriber` (Etapa 2). +/// É um pacote multi-arquivo do release `asr-models` do k2-fsa/sherpa-onnx. +pub const SHERPA_MODEL_ID: &str = "sherpa-onnx-nemo-parakeet-tdt-0.6b-v3-int8"; + +/// ID legado do modelo GGUF (mantido para o botão "Baixar modelo" não +/// quebrar builds existentes que ainda apontam para o catálogo antigo). +pub const GGUF_MODEL_ID: &str = "parakeet-tdt-0.6b-v3-q8_0"; +pub const GGUF_MODEL_URL: &str = "https://huggingface.co/cstr/parakeet-tdt-0.6b-v3-GGUF/resolve/main/parakeet-tdt-0.6b-v3-q8_0.gguf"; + +/// ID canônico de modelo a usar como "principal" (o que o app baixa por padrão). +pub const DEFAULT_MODEL_ID: &str = SHERPA_MODEL_ID; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(tag = "kind", rename_all = "camelCase")] +pub enum ModelFormat { + /// Arquivo único (ex.: `.gguf`). + SingleFile { filename: String, size_bytes: u64 }, + /// Pacote multi-arquivo (ex.: sherpa-onnx v3 int8 com 4 ONNX). + /// `archive_url` é o tarball `.tar.bz2`; `archive_size_bytes` é o tamanho esperado; + /// `extracted_subdir` é o subdiretório dentro do tarball; + /// `required_files` é a lista de arquivos esperados pós-extração. + Archive { + archive_url: String, + archive_size_bytes: u64, + extracted_subdir: String, + required_files: Vec, + }, +} #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "camelCase")] @@ -15,10 +41,8 @@ pub struct OfflineModel { pub id: String, pub display_name: String, pub repository: String, - pub filename: String, pub quantization: String, - pub size_bytes: u64, - pub download_url: String, + pub format: ModelFormat, pub languages: Vec, pub recommended_for_simple_pc: bool, } @@ -36,7 +60,8 @@ pub enum ModelInstallStatus { pub struct ModelWithStatus { pub model: OfflineModel, pub status: ModelInstallStatus, - pub local_path: Option, + /// Diretório onde o modelo está (ou deveria estar) instalado. + pub local_dir: PathBuf, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -48,49 +73,99 @@ pub struct ModelDownloadResult { } pub fn catalog() -> Vec { - vec![OfflineModel { - id: DEFAULT_MODEL_ID.to_string(), - display_name: "Parakeet TDT 0.6B v3 Q8".to_string(), - repository: "cstr/parakeet-tdt-0.6b-v3-GGUF".to_string(), - filename: "parakeet-tdt-0.6b-v3-q8_0.gguf".to_string(), - quantization: "Q8_0".to_string(), - size_bytes: 745_121_600, - download_url: DEFAULT_MODEL_URL.to_string(), - languages: vec![ - "bg", "hr", "cs", "da", "nl", "en", "et", "fi", "fr", "de", "el", "hu", "it", "lv", - "lt", "mt", "pl", "pt", "ro", "sk", "sl", "es", "sv", "ru", "uk", - ] - .into_iter() - .map(str::to_string) - .collect(), - recommended_for_simple_pc: true, - }] + vec![ + OfflineModel { + id: SHERPA_MODEL_ID.to_string(), + display_name: "Parakeet TDT 0.6B v3 (sherpa-onnx int8)".to_string(), + repository: "k2-fsa/sherpa-onnx asr-models".to_string(), + quantization: "int8".to_string(), + format: ModelFormat::Archive { + archive_url: "https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/sherpa-onnx-nemo-parakeet-tdt-0.6b-v3-int8.tar.bz2" + .to_string(), + archive_size_bytes: 670_000_000, + extracted_subdir: "sherpa-onnx-nemo-parakeet-tdt-0.6b-v3-int8".to_string(), + required_files: vec![ + "encoder.int8.onnx".to_string(), + "decoder.int8.onnx".to_string(), + "joiner.int8.onnx".to_string(), + "tokens.txt".to_string(), + ], + }, + languages: vec![ + "bg", "hr", "cs", "da", "nl", "en", "et", "fi", "fr", "de", "el", "hu", "it", "lv", + "lt", "mt", "pl", "pt", "ro", "sk", "sl", "es", "sv", "ru", "uk", + ] + .into_iter() + .map(str::to_string) + .collect(), + recommended_for_simple_pc: true, + }, + // Modelo legado GGUF (CrispASR). Continua aparecendo no catálogo para + // o botão "Baixar modelo" não quebrar builds existentes. + OfflineModel { + id: GGUF_MODEL_ID.to_string(), + display_name: "Parakeet TDT 0.6B v3 (GGUF Q8 - legados)".to_string(), + repository: "cstr/parakeet-tdt-0.6b-v3-GGUF".to_string(), + quantization: "Q8_0".to_string(), + format: ModelFormat::SingleFile { + filename: "parakeet-tdt-0.6b-v3-q8_0.gguf".to_string(), + size_bytes: 745_121_600, + }, + languages: vec!["pt".to_string(), "en".to_string()], + recommended_for_simple_pc: false, + }, + ] } pub fn model_statuses(model_dir: &Path) -> Vec { catalog() .into_iter() .map(|model| { - let path = model_dir.join(&model.filename); - let status = status_for_path(&path, model.size_bytes); + let local_dir = model_install_dir(model_dir, &model); + let status = status_for_model(&local_dir, &model); ModelWithStatus { model, - local_path: path.exists().then_some(path), status, + local_dir, } }) .collect() } -pub fn status_for_path(path: &Path, expected_size: u64) -> ModelInstallStatus { - match fs::metadata(path) { - Ok(metadata) if metadata.len() >= expected_size => ModelInstallStatus::Ready, - Ok(metadata) => ModelInstallStatus::Error(format!( - "arquivo incompleto: {} de {} bytes", - metadata.len(), - expected_size - )), - Err(_) => ModelInstallStatus::NotInstalled, +fn model_install_dir(model_dir_root: &Path, model: &OfflineModel) -> PathBuf { + model_dir_root.join(&model.id) +} + +fn status_for_model(install_dir: &Path, model: &OfflineModel) -> ModelInstallStatus { + match &model.format { + ModelFormat::SingleFile { + filename, + size_bytes, + } => { + let path = install_dir.join(filename); + match fs::metadata(&path) { + Ok(metadata) if metadata.len() >= *size_bytes => ModelInstallStatus::Ready, + Ok(metadata) => ModelInstallStatus::Error(format!( + "arquivo incompleto: {} de {} bytes", + metadata.len(), + size_bytes + )), + Err(_) => ModelInstallStatus::NotInstalled, + } + } + ModelFormat::Archive { required_files, .. } => { + let mut missing: Vec<&str> = Vec::new(); + for file in required_files { + if !install_dir.join(file).exists() { + missing.push(file.as_str()); + } + } + if missing.is_empty() { + ModelInstallStatus::Ready + } else { + ModelInstallStatus::Error(format!("arquivos ausentes: {}", missing.join(", "))) + } + } } } @@ -101,26 +176,60 @@ pub fn default_model() -> OfflineModel { .expect("catálogo deve conter o modelo padrão") } -pub fn default_model_path(model_dir: &Path) -> PathBuf { - model_dir.join(default_model().filename) +/// Retorna o diretório onde o modelo sherpa-onnx v3 int8 deve estar +/// instalado (com os 4 arquivos ONNX dentro). +pub fn sherpa_model_dir(model_dir_root: &Path) -> PathBuf { + model_install_dir(model_dir_root, &default_model()) +} + +/// Retorna o caminho do arquivo `.gguf` legado (se ainda em uso). +pub fn gguf_model_path(model_dir_root: &Path) -> PathBuf { + model_dir_root + .join(GGUF_MODEL_ID) + .join("parakeet-tdt-0.6b-v3-q8_0.gguf") } pub fn download_default_model(model_dir: &Path) -> AppResult { let model = default_model(); - fs::create_dir_all(model_dir).map_err(|error| AppError::Io(error.to_string()))?; + let install_dir = model_install_dir(model_dir, &model); + fs::create_dir_all(&install_dir).map_err(|error| AppError::Io(error.to_string()))?; - let final_path = model_dir.join(&model.filename); - let temp_path = model_dir.join(format!("{}.partial", model.filename)); + match &model.format { + ModelFormat::SingleFile { + filename, + size_bytes, + } => download_single_file(&model, &install_dir, filename, *size_bytes), + ModelFormat::Archive { + archive_url, + archive_size_bytes, + extracted_subdir, + required_files, + } => download_archive( + &model, + &install_dir, + archive_url, + *archive_size_bytes, + extracted_subdir, + required_files, + ), + } +} - let mut response = reqwest::blocking::get(&model.download_url) +fn download_single_file( + model: &OfflineModel, + install_dir: &Path, + filename: &str, + expected_size: u64, +) -> AppResult { + let final_path = install_dir.join(filename); + let temp_path = install_dir.join(format!("{filename}.partial")); + let mut response = reqwest::blocking::get(model.download_url_if_single()) .map_err(|error| AppError::Network(error.to_string()))? .error_for_status() .map_err(|error| AppError::Network(error.to_string()))?; - let mut file = fs::File::create(&temp_path).map_err(|error| AppError::Io(error.to_string()))?; let mut buffer = [0_u8; 1024 * 128]; let mut bytes_written = 0_u64; - loop { let read = response .read(&mut buffer) @@ -132,62 +241,189 @@ pub fn download_default_model(model_dir: &Path) -> AppResult AppResult { + let temp_archive = install_dir.join("archive.tar.bz2"); + let mut response = reqwest::blocking::get(archive_url) + .map_err(|error| AppError::Network(error.to_string()))? + .error_for_status() + .map_err(|error| AppError::Network(error.to_string()))?; + let mut file = + fs::File::create(&temp_archive).map_err(|error| AppError::Io(error.to_string()))?; + let mut buffer = [0_u8; 1024 * 128]; + let mut bytes_written = 0_u64; + loop { + let read = response + .read(&mut buffer) + .map_err(|error| AppError::Network(error.to_string()))?; + if read == 0 { + break; + } + file.write_all(&buffer[..read]) + .map_err(|error| AppError::Io(error.to_string()))?; + bytes_written += read as u64; + } + file.flush() + .map_err(|error| AppError::Io(error.to_string()))?; + if bytes_written < expected_archive_size / 2 { + // tolerância: tarballs do HF às vezes têm tamanho declarado errado + let _ = fs::remove_file(&temp_archive); + return Err(AppError::InvalidModel(format!( + "download incompleto: {bytes_written} bytes (esperado ~{expected_archive_size})" + ))); + } + + // Extrai o tarball usando `tar` do sistema (funciona em Windows+Linux+macOS). + let sep = std::path::MAIN_SEPARATOR.to_string(); + let prefix = install_dir.to_string_lossy().to_string(); + let prefix_with_sep = if prefix.ends_with(&sep) { + prefix.clone() + } else { + format!("{prefix}{sep}") + }; + let rel_archive = temp_archive + .to_string_lossy() + .to_string() + .strip_prefix(&prefix_with_sep) + .map(str::to_string) + .unwrap_or_else(|| temp_archive.to_string_lossy().to_string()); + let status = std::process::Command::new("tar") + .arg("-xjf") + .arg(&rel_archive) + .current_dir(install_dir) + .status() + .map_err(|error| AppError::Io(format!("falha ao executar tar: {error}")))?; + if !status.success() { + let _ = fs::remove_file(&temp_archive); + return Err(AppError::Io(format!( + "tar falhou com status {status} ao extrair {rel_archive}" + ))); + } + let _ = fs::remove_file(&temp_archive); + + // Move os arquivos extraídos de `install_dir//...` para `install_dir/...` + // para que `sherpa-onnx-offline --tokens=install_dir/tokens.txt` funcione. + let src_dir = install_dir.join(extracted_subdir); + if !src_dir.exists() { + return Err(AppError::InvalidModel(format!( + "subdiretório esperado {extracted_subdir:?} não encontrado após extração" + ))); + } + for file in required_files { + let from = src_dir.join(file); + let to = install_dir.join(file); + if from.exists() { + if to.exists() { + let _ = fs::remove_file(&to); + } + fs::rename(&from, &to).map_err(|error| AppError::Io(error.to_string()))?; + } + } + let _ = fs::remove_dir_all(&src_dir); + + Ok(ModelDownloadResult { + model_id: model.id.clone(), + path: install_dir.to_path_buf(), + bytes_written, + }) +} + +impl OfflineModel { + /// URL de download para modelos SingleFile. Mantida para retrocompatibilidade. + pub fn download_url_if_single(&self) -> String { + if let ModelFormat::SingleFile { filename: _, .. } = &self.format { + // O GGUF legado continua no HF; o sherpa-onnx já tem seu próprio caminho + // via Archive. Para SingleFile, usamos a URL do GGUF por id. + match self.id.as_str() { + GGUF_MODEL_ID => GGUF_MODEL_URL.to_string(), + _ => String::new(), + } + } else { + String::new() + } + } +} + #[cfg(test)] mod tests { - use std::time::{SystemTime, UNIX_EPOCH}; - use super::*; #[test] - fn catalog_contains_parakeet_q8_model() { + fn catalog_contains_sherpa_parakeet_v3() { let model = default_model(); - - assert_eq!(model.repository, "cstr/parakeet-tdt-0.6b-v3-GGUF"); - assert_eq!(model.filename, "parakeet-tdt-0.6b-v3-q8_0.gguf"); - assert_eq!(model.quantization, "Q8_0"); - assert_eq!(model.download_url, DEFAULT_MODEL_URL); - assert!(model.languages.contains(&"pt".to_string())); + assert_eq!(model.id, SHERPA_MODEL_ID); + assert_eq!(model.quantization, "int8"); + assert!(matches!(model.format, ModelFormat::Archive { .. })); } #[test] - fn status_detects_ready_model_by_size() { - let unique = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_nanos()) + fn catalog_keeps_gguf_legacy() { + let gguf = catalog() + .into_iter() + .find(|model| model.id == GGUF_MODEL_ID) + .expect("GGUF legado deve continuar no catálogo"); + assert!(matches!(gguf.format, ModelFormat::SingleFile { .. })); + } + + #[test] + fn sherpa_model_dir_returns_path() { + let root = std::env::temp_dir(); + let dir = sherpa_model_dir(&root); + assert!(dir.ends_with(SHERPA_MODEL_ID)); + } + + #[test] + fn archive_model_status_reports_missing_files() { + let unique = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_nanos()) .unwrap_or(0); - let dir = std::env::temp_dir().join(format!("transcript-model-test-{unique}")); - fs::create_dir_all(&dir).unwrap_or_else(|error| panic!("create temp dir: {error}")); - let path = dir.join("model.gguf"); - fs::write(&path, [1_u8, 2, 3, 4]).unwrap_or_else(|error| panic!("write model: {error}")); + let dir = std::env::temp_dir().join(format!("transcript-archive-test-{unique}")); + fs::create_dir_all(&dir).unwrap_or_else(|error| panic!("create dir: {error}")); + let model = default_model(); + let status = status_for_model(&dir, &model); + assert!(matches!(status, ModelInstallStatus::Error(_))); - assert_eq!(status_for_path(&path, 4), ModelInstallStatus::Ready); - assert!(matches!( - status_for_path(&path, 5), - ModelInstallStatus::Error(_) - )); + // cria o primeiro arquivo e checa de novo + if let ModelFormat::Archive { required_files, .. } = &model.format { + fs::write(dir.join(&required_files[0]), b"placeholder") + .unwrap_or_else(|error| panic!("write file: {error}")); + } + let status = status_for_model(&dir, &model); + assert!(matches!(status, ModelInstallStatus::Error(_))); - let _ = fs::remove_file(path); - let _ = fs::remove_dir(dir); + // cria todos e checa Ready + if let ModelFormat::Archive { required_files, .. } = &model.format { + for file in required_files { + fs::write(dir.join(file), b"x") + .unwrap_or_else(|error| panic!("write {file}: {error}")); + } + } + let status = status_for_model(&dir, &model); + assert_eq!(status, ModelInstallStatus::Ready); + + let _ = fs::remove_dir_all(dir); } } diff --git a/src-tauri/src/ui_commands.rs b/src-tauri/src/ui_commands.rs index 4f4570a..68d53ec 100644 --- a/src-tauri/src/ui_commands.rs +++ b/src-tauri/src/ui_commands.rs @@ -4,7 +4,7 @@ use crate::app_state::{AppSnapshot, ManagedAppState}; use crate::diagnostics::{audio_or_error, log_attempt, tail, TranscriptionAttempt}; use crate::injection::{FallbackInjector, InsertMethod, InsertRequest, InsertResult, TextInjector}; use crate::models::{ - default_model_path, download_default_model, model_statuses, ModelDownloadResult, + download_default_model, model_statuses, ModelDownloadResult, ModelWithStatus, }; use crate::settings::AppSettings; @@ -61,7 +61,7 @@ pub(crate) fn stop_recording_and_transcribe( audio_path: audio_path.clone(), audio_size_bytes: 0, sidecar: "parakeet-main".to_string(), - model: default_model_path(&state.model_dir).display().to_string(), + model: state.sherpa_model_dir().display().to_string(), exit_code: None, stdout_tail: String::new(), stderr_tail: String::new(), @@ -105,7 +105,7 @@ pub(crate) fn stop_recording_and_transcribe( audio_path: audio_path.clone(), audio_size_bytes: audio_size, sidecar: "parakeet-main".to_string(), - model: default_model_path(&state.model_dir).display().to_string(), + model: state.sherpa_model_dir().display().to_string(), exit_code: None, stdout_tail: tail(&result.text, 1024), stderr_tail: String::new(), @@ -126,7 +126,7 @@ pub(crate) fn stop_recording_and_transcribe( audio_path: audio_path.clone(), audio_size_bytes: audio_size, sidecar: "parakeet-main".to_string(), - model: default_model_path(&state.model_dir).display().to_string(), + model: state.sherpa_model_dir().display().to_string(), exit_code: None, stdout_tail: String::new(), stderr_tail: String::new(), @@ -241,7 +241,12 @@ mod tests { // transcriber não falhar antes de invocar o sidecar inexistente. let model_dir = unique_path("sherpa-empty-model"); std::fs::create_dir_all(&model_dir).unwrap_or_else(|error| panic!("mkdir model: {error}")); - for name in ["tokens.txt", "encoder.int8.onnx", "decoder.int8.onnx", "joiner.int8.onnx"] { + for name in [ + "tokens.txt", + "encoder.int8.onnx", + "decoder.int8.onnx", + "joiner.int8.onnx", + ] { std::fs::write(model_dir.join(name), b"placeholder") .unwrap_or_else(|error| panic!("write {name}: {error}")); }