diff --git a/apps/desktop-gpui/Cargo.lock b/apps/desktop-gpui/Cargo.lock index 469e8a1469..66e5795c54 100644 --- a/apps/desktop-gpui/Cargo.lock +++ b/apps/desktop-gpui/Cargo.lock @@ -1563,6 +1563,7 @@ dependencies = [ "gpui", "gpui_platform", "gpui_tokio", + "hound", "image 0.25.10", "jpeg-decoder", "kameo", diff --git a/apps/desktop-gpui/Cargo.toml b/apps/desktop-gpui/Cargo.toml index 50fccce29a..d8af892904 100644 --- a/apps/desktop-gpui/Cargo.toml +++ b/apps/desktop-gpui/Cargo.toml @@ -238,6 +238,7 @@ cidre = { git = "https://github.com/CapSoftware/cidre", rev = "bf84b67079a8", fe # Whisper caption models only. [target.'cfg(not(all(target_os = "macos", target_arch = "x86_64")))'.dependencies] parakeet-rs = "0.3.4" +hound = "3.5" # Mirror of the root workspace patch: cap-recording's converter stack is built # and tested against the vendored wgpu-hal, not the crates.io release. diff --git a/apps/desktop-gpui/src/transcription.rs b/apps/desktop-gpui/src/transcription.rs index 9f23b43524..bc7ad223b0 100644 --- a/apps/desktop-gpui/src/transcription.rs +++ b/apps/desktop-gpui/src/transcription.rs @@ -1557,6 +1557,38 @@ fn process_with_whisper( // Parakeet // --------------------------------------------------------------------------- +/// Parakeet's encoder fails with an ONNX broadcast error past 400 s of audio +/// (5,000 frames of 80 ms), and below that its time and memory per second of +/// audio grow with the input's length. Minute-long windows, cut at the quietest +/// 100 ms near each boundary so words stay whole, keep cost linear in the +/// recording's length and lift the 400 s cap. +#[cfg(not(all(target_os = "macos", target_arch = "x86_64")))] +const PARAKEET_CHUNK_SECONDS: usize = 60; +#[cfg(not(all(target_os = "macos", target_arch = "x86_64")))] +const PARAKEET_CUT_SEARCH_SECONDS: usize = 5; + +#[cfg(not(all(target_os = "macos", target_arch = "x86_64")))] +fn parakeet_chunk_ranges(samples: &[f32], sample_rate: usize) -> Vec> { + let target = PARAKEET_CHUNK_SECONDS * sample_rate; + let search = PARAKEET_CUT_SEARCH_SECONDS * sample_rate; + let window = sample_rate / 10; + let energy = |at: usize| -> f32 { samples[at..at + window].iter().map(|s| s * s).sum() }; + let mut ranges = Vec::new(); + let mut start = 0; + while samples.len() - start > target { + let end = start + target; + let quietest = (end - search..=end - window) + .step_by(sample_rate / 100) + .min_by(|a, b| energy(*a).total_cmp(&energy(*b))) + .unwrap_or(end - window); + let cut = quietest + window / 2; + ranges.push(start..cut); + start = cut; + } + ranges.push(start..samples.len()); + ranges +} + /// `process_with_parakeet` (`captions.rs:1202-1302`). #[cfg(not(all(target_os = "macos", target_arch = "x86_64")))] fn process_with_parakeet( @@ -1602,22 +1634,51 @@ fn process_with_parakeet( } }; - let result = { + let mut reader = hound::WavReader::open(audio_path) + .map_err(|e| format!("Failed to open audio file: {e} at path: {audio_path:?}"))?; + let spec = reader.spec(); + if spec.channels != 1 || spec.sample_rate != WHISPER_SAMPLE_RATE { + return Err(format!( + "Expected 16 kHz mono audio, got {} Hz with {} channels", + spec.sample_rate, spec.channels + )); + } + let sample_rate = spec.sample_rate; + let samples = reader + .samples::() + .map(|sample| sample.map(|sample| f32::from(sample) / 32768.0)) + .collect::, _>>() + .map_err(|e| format!("Failed to read audio file: {e}"))?; + if samples.is_empty() { + return Err("No speech detected in the audio".to_string()); + } + + let mut tokens = Vec::new(); + { let mut parakeet = model_arc .lock() .map_err(|e| format!("Failed to lock Parakeet model: {e}"))?; - parakeet - .transcribe_file(audio_path, Some(TimestampMode::Words)) - .map_err(|e| format!("Parakeet transcription failed: {e}"))? - }; + for range in parakeet_chunk_ranges(&samples, sample_rate as usize) { + let offset = range.start as f32 / sample_rate as f32; + let result = parakeet + .transcribe_samples( + samples[range].to_vec(), + sample_rate, + 1, + Some(TimestampMode::Words), + ) + .map_err(|e| format!("Parakeet transcription failed: {e}"))?; + tokens.extend(result.tokens.into_iter().map(|token| CaptionWord { + text: token.text, + start: token.start + offset, + end: token.end + offset, + })); + } + } - tracing::info!("Parakeet produced {} timed tokens", result.tokens.len()); + tracing::info!("Parakeet produced {} timed tokens", tokens.len()); - caption_segments_from_tokens(result.tokens.iter().map(|token| CaptionWord { - text: token.text.clone(), - start: token.start, - end: token.end, - })) + caption_segments_from_tokens(tokens) } /// Parakeet's timed word tokens -> caption segments: drop empty tokens, @@ -2204,6 +2265,66 @@ mod verify_tests; mod tests { use super::*; + #[cfg(not(all(target_os = "macos", target_arch = "x86_64")))] + #[test] + fn parakeet_chunks_cut_in_the_quiet_gap_and_cover_every_sample() { + let rate = 16_000; + let mut samples: Vec = (0..rate * 150).map(|i| (i as f32 * 0.05).sin()).collect(); + samples[rate * 57..rate * 57 + rate / 2].fill(0.0); + + let ranges = parakeet_chunk_ranges(&samples, rate); + + assert_eq!(ranges.len(), 3); + assert!((rate * 57..rate * 57 + rate / 2).contains(&ranges[0].end)); + assert_eq!(ranges[0].start, 0); + assert_eq!(ranges.last().unwrap().end, samples.len()); + for pair in ranges.windows(2) { + assert_eq!(pair[0].end, pair[1].start); + } + assert!( + ranges + .iter() + .all(|r| r.len() <= rate * PARAKEET_CHUNK_SECONDS) + ); + } + + #[cfg(not(all(target_os = "macos", target_arch = "x86_64")))] + #[test] + fn parakeet_keeps_a_one_minute_recording_in_one_chunk() { + let samples = vec![0.5; 16_000 * PARAKEET_CHUNK_SECONDS]; + + assert_eq!( + parakeet_chunk_ranges(&samples, 16_000), + vec![0..samples.len()] + ); + } + + #[cfg(not(all(target_os = "macos", target_arch = "x86_64")))] + #[test] + #[ignore = "needs CAP_PARAKEET_BENCH_MODEL and CAP_PARAKEET_BENCH_WAV"] + fn benchmark_parakeet_transcription() { + let model = std::env::var("CAP_PARAKEET_BENCH_MODEL").expect("CAP_PARAKEET_BENCH_MODEL"); + let wav = + PathBuf::from(std::env::var("CAP_PARAKEET_BENCH_WAV").expect("CAP_PARAKEET_BENCH_WAV")); + let started = std::time::Instant::now(); + let result = process_with_parakeet(&wav, &model); + let seconds = started.elapsed().as_secs_f64(); + match result { + Ok(segments) => { + let words: Vec<&CaptionWord> = segments.iter().flat_map(|s| &s.words).collect(); + println!("parakeet_seconds={seconds:.2} words={}", words.len()); + if let Ok(path) = std::env::var("CAP_PARAKEET_BENCH_WORDS") { + let dump: String = words + .iter() + .map(|w| format!("{}\t{:.2}\t{:.2}\n", w.text, w.start, w.end)) + .collect(); + std::fs::write(path, dump).expect("write words"); + } + } + Err(error) => panic!("parakeet_seconds={seconds:.2} error={error}"), + } + } + fn word(text: &str, index: usize) -> CaptionWord { CaptionWord { text: text.to_string(),