Simplify Pocket TTS segmentation

Signed-off-by: John Tennant <jtennant@squareup.com>

crates/buzz-voice/src/pocket.rs

@@ -136,6 +136,15 @@ pub fn load_text_to_speech(model_dir: &str) -> Result<PocketTts, String> {
 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<Vec<String>, 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.

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<Vec<String>, 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<Vec<String>, 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<F>(
@@ -656,6 +628,150 @@ impl AprilPocketTts {
     }
 }
 
+fn split_at_natural_boundaries<F>(
+    text: &str,
+    max_tokens: usize,
+    isolate_first_sentence: bool,
+    mut token_count: F,
+) -> Result<Vec<String>, String>
+where
+    F: FnMut(&str) -> Result<usize, String>,
+{
+    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<Session, String> {
     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<usize, String> {
+        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))
+        }));
     }
 }