diff --git a/crates/buzz-voice/src/pocket.rs b/crates/buzz-voice/src/pocket.rs index d18a30608..7e15c44cb 100644 --- a/crates/buzz-voice/src/pocket.rs +++ b/crates/buzz-voice/src/pocket.rs @@ -136,6 +136,15 @@ pub fn load_text_to_speech(model_dir: &str) -> Result { impl PocketTts { /// Split text into synthesis units that satisfy the bundle's exact /// 50-token input limit. + /// + /// The first sentence remains its own unit when it fits. Oversized + /// sentences fall back to clause, word, and UTF-8 scalar boundaries, while + /// later sentences pack into the largest natural unit that fits. + /// + /// Chunks are contiguous substrings of the prepared model prompt and may + /// retain boundary whitespace. Concatenating them with `chunks.concat()` + /// reconstructs that prompt exactly, and each chunk's prepared token count + /// is at most 50. pub fn split_text_into_chunks(&self, text: &str) -> Result, String> { if SynthesisCallGuard::is_active(self as *const Self as usize) { return Err("Pocket TTS callback re-entered the active engine".to_string()); @@ -146,7 +155,7 @@ impl PocketTts { self.inner .lock() .map_err(|_| "Pocket TTS engine lock poisoned".to_string())? - .split_prompt(&prepared) + .split_playback_prompt(&prepared) } /// Synthesize text with the supplied reference voice. diff --git a/crates/buzz-voice/src/pocket_april.rs b/crates/buzz-voice/src/pocket_april.rs index 3d85a3bdc..11344072d 100644 --- a/crates/buzz-voice/src/pocket_april.rs +++ b/crates/buzz-voice/src/pocket_april.rs @@ -36,6 +36,13 @@ const DECODER_CHUNK_FRAMES: usize = 12; const TOKENS_PER_SECOND_ESTIMATE: f32 = 3.0; const GENERATION_SECONDS_PADDING: f32 = 2.0; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum TextBoundary { + Sentence, + Clause, + Word, +} + #[derive(Debug, Deserialize)] struct Bundle { schema_version: u32, @@ -251,62 +258,27 @@ impl AprilPocketTts { &self, prepared: &AprilPreparedPrompt, ) -> Result, String> { - if self.token_count(&prepared.text)? <= self.bundle.max_token_per_chunk { + if self.prepared_token_count(&prepared.text)? <= self.bundle.max_token_per_chunk { return Ok(vec![prepared.text.clone()]); } + split_at_natural_boundaries( + &prepared.text, + self.bundle.max_token_per_chunk, + false, + |text| self.prepared_token_count(text), + ) + } - let mut chunks = Vec::new(); - let mut current = String::new(); - for word in prepared.text.split_whitespace() { - let candidate = if current.is_empty() { - word.to_string() - } else { - format!("{current} {word}") - }; - if self.prepared_token_count(&candidate)? <= self.bundle.max_token_per_chunk { - current = candidate; - continue; - } - if !current.is_empty() { - chunks.push(std::mem::take(&mut current)); - } - - if self.prepared_token_count(word)? <= self.bundle.max_token_per_chunk { - current = word.to_string(); - continue; - } - - let mut fragment = String::new(); - for ch in word.chars() { - let candidate = format!("{fragment}{ch}"); - if !fragment.is_empty() - && self.prepared_token_count(&candidate)? > self.bundle.max_token_per_chunk - { - chunks.push(std::mem::take(&mut fragment)); - } - fragment.push(ch); - } - current = fragment; - } - if !current.is_empty() { - chunks.push(current); - } - - chunks - .into_iter() - .map(|text| { - let chunk = prepare_april_prompt(&text) - .ok_or_else(|| "Pocket TTS prompt chunk became empty".to_string())?; - let token_count = self.token_count(&chunk.text)?; - if token_count > self.bundle.max_token_per_chunk { - return Err(format!( - "Pocket TTS prompt chunk has {token_count} tokens; maximum is {}", - self.bundle.max_token_per_chunk - )); - } - Ok(chunk.text) - }) - .collect() + pub(crate) fn split_playback_prompt( + &self, + prepared: &AprilPreparedPrompt, + ) -> Result, String> { + split_at_natural_boundaries( + &prepared.text, + self.bundle.max_token_per_chunk, + true, + |text| self.prepared_token_count(text), + ) } pub(crate) fn synth_chunk_streaming( @@ -656,6 +628,150 @@ impl AprilPocketTts { } } +fn split_at_natural_boundaries( + text: &str, + max_tokens: usize, + isolate_first_sentence: bool, + mut token_count: F, +) -> Result, String> +where + F: FnMut(&str) -> Result, +{ + if text.is_empty() { + return Ok(Vec::new()); + } + + let mut chunks = Vec::new(); + let mut start = 0; + while start < text.len() { + while text[start..] + .chars() + .next() + .is_some_and(char::is_whitespace) + { + start += text[start..] + .chars() + .next() + .expect("checked above") + .len_utf8(); + } + if start == text.len() { + break; + } + + let mut first_sentence_end = None; + let mut sentence_end = None; + let mut clause_end = None; + let mut word_end = None; + for (offset, ch) in text[start..].char_indices() { + let end = start + offset + ch.len_utf8(); + let at_word_end = + end == text.len() || text[end..].chars().next().is_some_and(char::is_whitespace); + let at_clause_end = matches!(ch, '—' | '–') + && !text[end..] + .chars() + .next() + .is_some_and(is_closing_punctuation); + if (!at_word_end && !at_clause_end) || token_count(&text[start..end])? > max_tokens { + continue; + } + + word_end = Some(end); + match natural_boundary(&text[start..end], end == text.len()) { + TextBoundary::Sentence => { + first_sentence_end.get_or_insert(end); + sentence_end = Some(end); + } + TextBoundary::Clause => clause_end = Some(end), + TextBoundary::Word => {} + } + } + + let preferred_end = if isolate_first_sentence && chunks.is_empty() { + first_sentence_end.or(clause_end).or(word_end) + } else { + sentence_end.or(clause_end).or(word_end) + }; + let end = if let Some(end) = preferred_end { + end + } else { + // A single word can itself exceed the model limit. Preserve a + // scalar boundary as the final safety case without losing UTF-8. + let mut scalar_end = None; + for (offset, ch) in text[start..].char_indices() { + if ch.is_whitespace() { + break; + } + let end = start + offset + ch.len_utf8(); + if token_count(&text[start..end])? <= max_tokens { + scalar_end = Some(end); + } + } + scalar_end.ok_or_else(|| { + format!( + "Pocket TTS prompt cannot fit one character within the {max_tokens}-token limit" + ) + })? + }; + + let mut next_start = end; + while text[next_start..] + .chars() + .next() + .is_some_and(char::is_whitespace) + { + next_start += text[next_start..] + .chars() + .next() + .expect("checked above") + .len_utf8(); + } + chunks.push(text[start..next_start].to_string()); + start = next_start; + } + + debug_assert_eq!(chunks.concat(), text); + Ok(chunks) +} + +fn natural_boundary(candidate: &str, is_end_of_text: bool) -> TextBoundary { + if is_end_of_text { + return TextBoundary::Sentence; + } + + let mut chars = candidate.chars().rev(); + let mut last = chars.next(); + while last.is_some_and(is_closing_punctuation) { + last = chars.next(); + } + match last { + Some('.' | '!' | '?') if !looks_like_abbreviation(candidate) => TextBoundary::Sentence, + Some(',' | ';' | ':' | '—' | '–') => TextBoundary::Clause, + _ => TextBoundary::Word, + } +} + +fn is_closing_punctuation(ch: char) -> bool { + matches!(ch, '"' | '\'' | '”' | '’' | ')' | ']' | '}') +} + +fn looks_like_abbreviation(candidate: &str) -> bool { + const ABBREVIATIONS: &[&str] = &[ + "Dr.", "Mr.", "Mrs.", "Ms.", "Prof.", "Sr.", "Jr.", "St.", "Ave.", "Rd.", "Blvd.", "Dept.", + "Inc.", "Ltd.", "Co.", "Corp.", "etc.", "vs.", "i.e.", "e.g.", "Ph.D.", + ]; + + let candidate = candidate.trim_end_matches(is_closing_punctuation); + let last_word = candidate + .rsplit_once(char::is_whitespace) + .map_or(candidate, |(_, word)| word); + ABBREVIATIONS.contains(&last_word) + || (last_word.ends_with('.') + && last_word[..last_word.len() - 1] + .chars() + .all(|ch| ch.is_ascii_digit())) +} + fn load_session(path: PathBuf, num_threads: usize) -> Result { if !path.is_file() { return Err(format!("missing Pocket TTS file: {}", path.display())); @@ -879,6 +995,89 @@ mod tests { assert_eq!(shape_len(&[2, 1, 8, 1000, 64]).expect("shape"), 1_024_000); } + fn whitespace_token_count(text: &str) -> Result { + Ok(text.split_whitespace().count()) + } + + #[test] + fn natural_split_keeps_first_sentence_separate_then_packs_the_remainder() { + let text = "One two. Three four. Five six."; + let chunks = split_at_natural_boundaries(text, 4, true, whitespace_token_count).unwrap(); + assert_eq!(chunks, ["One two. ", "Three four. Five six."]); + assert_eq!(chunks.concat(), text); + } + + #[test] + fn model_split_packs_multiple_sentences_within_limit() { + let text = "One two. Three four. Five six."; + let chunks = split_at_natural_boundaries(text, 4, false, whitespace_token_count).unwrap(); + assert_eq!(chunks, ["One two. Three four. ", "Five six."]); + assert_eq!(chunks.concat(), text); + } + + #[test] + fn natural_split_prefers_preceding_sentence_boundary() { + let text = "One two. Three four five six."; + let chunks = split_at_natural_boundaries(text, 5, true, whitespace_token_count).unwrap(); + assert_eq!(chunks, ["One two. ", "Three four five six."]); + assert_eq!(chunks.concat(), text); + } + + #[test] + fn oversized_sentence_uses_clause_then_word_fallback() { + let clause_text = "One two three, four five six seven."; + let clause_chunks = + split_at_natural_boundaries(clause_text, 5, true, whitespace_token_count).unwrap(); + assert_eq!(clause_chunks, ["One two three, ", "four five six seven."]); + assert_eq!(clause_chunks.concat(), clause_text); + + let word_text = "One two three four five six."; + let word_chunks = + split_at_natural_boundaries(word_text, 4, true, whitespace_token_count).unwrap(); + assert_eq!(word_chunks, ["One two three four ", "five six."]); + assert_eq!(word_chunks.concat(), word_text); + } + + #[test] + fn natural_split_preserves_unicode_punctuation_and_abbreviations() { + let text = "“Café naïve?” Maybe—yes, definitely; 東京 speaks."; + let chunks = split_at_natural_boundaries(text, 3, true, whitespace_token_count).unwrap(); + assert_eq!( + chunks, + ["“Café naïve?” ", "Maybe—yes, definitely; ", "東京 speaks."] + ); + assert_eq!(chunks.concat(), text); + + let abbreviation = "Dr. Smith waits. Then leaves."; + let chunks = + split_at_natural_boundaries(abbreviation, 3, true, whitespace_token_count).unwrap(); + assert_eq!(chunks, ["Dr. Smith waits. ", "Then leaves."]); + assert_eq!(chunks.concat(), abbreviation); + + let unspaced_clause = "alpha beta—gamma delta"; + let chunks = + split_at_natural_boundaries(unspaced_clause, 2, true, whitespace_token_count).unwrap(); + assert_eq!(chunks, ["alpha beta—", "gamma delta"]); + assert_eq!(chunks.concat(), unspaced_clause); + } + + #[test] + fn natural_split_does_not_treat_numeric_punctuation_as_unspaced_clauses() { + let text = "Meet at 12:30 with 1,000 guests onward."; + let chunks = split_at_natural_boundaries(text, 3, true, whitespace_token_count).unwrap(); + assert_eq!(chunks, ["Meet at 12:30 ", "with 1,000 guests ", "onward."]); + assert_eq!(chunks.concat(), text); + } + + #[test] + fn oversized_word_uses_utf8_scalar_boundary_without_loss() { + let text = "éééé"; + let chunks = + split_at_natural_boundaries(text, 3, true, |chunk| Ok(chunk.chars().count())).unwrap(); + assert_eq!(chunks, ["ééé", "é"]); + assert_eq!(chunks.concat(), text); + } + #[test] fn normal_noise_has_requested_length() { let mut rng = rand::rng(); @@ -928,8 +1127,10 @@ mod tests { assert!(chunks.len() > 1); assert!(chunks.iter().all(|chunk| { - engine.token_count(chunk).expect("tokenize chunk") <= engine.bundle.max_token_per_chunk + engine.prepared_token_count(chunk).expect("tokenize chunk") + <= engine.bundle.max_token_per_chunk })); + assert_eq!(chunks.concat(), prepared.text); } #[test] @@ -943,16 +1144,20 @@ mod tests { let chunks = engine.split_prompt(&prepared).expect("split long sentence"); let token_counts: Vec<_> = chunks .iter() - .map(|chunk| engine.token_count(chunk).expect("count tokens")) + .map(|chunk| engine.prepared_token_count(chunk).expect("count tokens")) .collect(); - assert_eq!( - chunks, - [ - "And sometimes, when I am certain the reader is rested, I will engage him with a sentence of considerable length, a sentence that burns with energy and builds with all the.", - "Impetus of a crescendo, the roll of the drums, the crash of the cymbals–sounds that say listen to this, it is important.", - ] - ); - assert_eq!(token_counts, [48, 44]); + assert!(token_counts + .iter() + .all(|&count| count <= engine.bundle.max_token_per_chunk)); + assert_eq!(chunks.concat(), prepared.text); + assert!(chunks.len() > 1); + assert!(chunks[..chunks.len() - 1].iter().all(|chunk| { + chunk + .trim_end() + .chars() + .last() + .is_some_and(|ch| ['.', '!', '?', ',', ';', ':', '—', '–'].contains(&ch)) + })); } }