diff --git a/crates/project/src/caption_words.rs b/crates/project/src/caption_words.rs index 8fdb6c39fc..3b419da39f 100644 --- a/crates/project/src/caption_words.rs +++ b/crates/project/src/caption_words.rs @@ -1,5 +1,8 @@ use crate::CaptionWord; +#[cfg(test)] +mod normalization_tests; + /// Whisper/Parakeet sometimes stretch a trailing word's end across a following /// silence (e.g. a 16s "seconds."), which leaves the rendered caption stuck on /// screen and duplicates the word across timeline cuts once projected. Real @@ -79,34 +82,33 @@ pub fn caption_boundary_word_is_weak(word: &CaptionWord) -> bool { ) } -pub fn normalize_caption_words(words: Vec) -> Vec { - let mut normalized: Vec = Vec::with_capacity(words.len()); - - for word in words { - let text = word.text.trim(); - if text.is_empty() { - continue; +pub fn normalize_caption_words(mut words: Vec) -> Vec { + words.retain_mut(|word| { + let trimmed = word.text.trim(); + if trimmed.is_empty() { + return false; } - - if caption_token_attaches_to_previous(text) - && let Some(previous) = normalized.last_mut() - { - previous.text.push_str(text); + let start = word.text.len() - word.text.trim_start().len(); + let end = start + trimmed.len(); + word.text.truncate(end); + drop(word.text.drain(..start)); + true + }); + words.dedup_by(|word, previous| { + if caption_token_attaches_to_previous(&word.text) { + previous.text.push_str(&word.text); previous.end = word.end; + true } else { - normalized.push(CaptionWord { - text: text.to_string(), - start: word.start, - end: word.end, - }); + false } - } + }); - for word in &mut normalized { + for word in &mut words { word.end = word.end.min(word.start + MAX_CAPTION_WORD_DURATION); } - normalized + words } pub fn caption_text_from_words<'a>(words: impl IntoIterator) -> String { diff --git a/crates/project/src/caption_words/normalization_tests.rs b/crates/project/src/caption_words/normalization_tests.rs new file mode 100644 index 0000000000..73fcdcd79d --- /dev/null +++ b/crates/project/src/caption_words/normalization_tests.rs @@ -0,0 +1,196 @@ +use super::*; +use std::{hint::black_box, time::Instant}; + +fn baseline_normalize(words: Vec) -> Vec { + let mut normalized: Vec = Vec::with_capacity(words.len()); + for word in words { + let text = word.text.trim(); + if text.is_empty() { + continue; + } + if caption_token_attaches_to_previous(text) + && let Some(previous) = normalized.last_mut() + { + previous.text.push_str(text); + previous.end = word.end; + } else { + normalized.push(CaptionWord { + text: text.to_string(), + start: word.start, + end: word.end, + }); + } + } + for word in &mut normalized { + word.end = word.end.min(word.start + MAX_CAPTION_WORD_DURATION); + } + normalized +} + +fn assert_words_equal(actual: &[CaptionWord], expected: &[CaptionWord]) { + assert_eq!(actual.len(), expected.len()); + for (actual, expected) in actual.iter().zip(expected) { + assert_eq!(actual.text, expected.text); + assert_eq!(actual.start.to_bits(), expected.start.to_bits()); + assert_eq!(actual.end.to_bits(), expected.end.to_bits()); + } +} + +#[test] +fn normalization_matches_previous_text_and_float_bits() { + let texts = [ + "", + " ", + "\t\n", + "\u{2003}\u{a0}", + "hello", + " hello\t", + "é日", + " नमस्ते ", + ",", + " . ", + "!?", + "’s", + "、", + "。", + "[", + "(test)", + "a\u{300}", + "\u{200b}", + ]; + let times = [ + 0.0, + -0.0, + 1.0, + -1.0, + f32::MAX, + f32::MIN, + f32::INFINITY, + f32::NEG_INFINITY, + f32::from_bits(0x7fc00123), + ]; + for first in texts { + for second in texts { + for (index, start) in times.into_iter().enumerate() { + let words = [first, second, " ", ",", "word", " . ", "。"] + .into_iter() + .enumerate() + .map(|(offset, text)| CaptionWord { + text: text.to_string(), + start, + end: times[(index + offset + 1) % times.len()], + }) + .collect::>(); + let expected = baseline_normalize(words.clone()); + let actual = normalize_caption_words(words); + assert_words_equal(&actual, &expected); + assert_eq!( + caption_text_from_words(&actual), + caption_text_from_words(&expected) + ); + let chunk_text = |words: &[CaptionWord]| { + caption_word_chunks(words) + .into_iter() + .map(caption_text_from_words) + .collect::>() + }; + assert_eq!(chunk_text(&actual), chunk_text(&expected)); + } + } + } + assert!(normalize_caption_words(Vec::new()).is_empty()); +} + +#[test] +fn normalization_retains_owned_word_and_vector_allocations() { + let words: Vec<_> = [" first ", "second", "\u{2003}third\u{2003}"] + .into_iter() + .map(|text| CaptionWord { + text: text.to_string(), + ..Default::default() + }) + .collect(); + let vector = words.as_ptr(); + let strings: Vec<_> = words.iter().map(|word| word.text.as_ptr()).collect(); + let words = normalize_caption_words(words); + assert_eq!(words.as_ptr(), vector); + for (word, pointer) in words.iter().zip(strings) { + assert_eq!(word.text.as_ptr(), pointer); + } + assert_eq!(caption_text_from_words(&words), "first second third"); +} + +fn fixture(count: usize, padded: bool) -> Vec { + let tokens = [ + "Recording", + "a", + "clear", + "demonstration", + ",", + "with", + "Cap", + ".", + ]; + (0..count) + .map(|index| { + let token = tokens[index % tokens.len()]; + CaptionWord { + text: if padded { + format!(" {token} ") + } else { + token.into() + }, + start: index as f32 * 0.4, + end: index as f32 * 0.4 + 0.3, + } + }) + .collect() +} + +fn measure(words: &[CaptionWord], normalize: fn(Vec) -> Vec) -> f64 { + let iterations = (120_000 / words.len()).clamp(8, 500); + let inputs: Vec<_> = (0..iterations).map(|_| words.to_vec()).collect(); + let started = Instant::now(); + for input in inputs { + black_box(normalize(black_box(input))); + } + started.elapsed().as_secs_f64() * 1e6 / iterations as f64 +} + +#[test] +#[ignore] +fn benchmark_caption_word_normalization() { + println!("words,padded,before_us,after_us,speedup"); + for count in [32, 1_000, 12_000] { + for padded in [false, true] { + let words = fixture(count, padded); + assert_words_equal( + &normalize_caption_words(words.clone()), + &baseline_normalize(words.clone()), + ); + for _ in 0..3 { + black_box(baseline_normalize(words.clone())); + black_box(normalize_caption_words(words.clone())); + } + let mut before = Vec::new(); + let mut after = Vec::new(); + for run in 0..9 { + if run % 2 == 0 { + before.push(measure(&words, baseline_normalize)); + after.push(measure(&words, normalize_caption_words)); + } else { + after.push(measure(&words, normalize_caption_words)); + before.push(measure(&words, baseline_normalize)); + } + } + before.sort_by(f64::total_cmp); + after.sort_by(f64::total_cmp); + println!( + "{count},{padded},{:.3},{:.3},{:.2}", + before[4], + after[4], + before[4] / after[4] + ); + } + } +}