Fix silent truncation of long phoneme segments

ARCHITECTURE.md

@@ -81,10 +81,16 @@ flowchart TD
 3. **Phonemize.** Unless the input is already phonemes, `Phonemizer::phonemize(text, language)` dispatches to that
    language's `PhonemizerBase`. It passes along the number language set by `kokoro_set_number_language`. The result is
    one UTF-8 phoneme string in Kokoro's symbol set.
-4. **Chunk.** `split_phonemes` (`src/PhonemeChunker.cpp`) splits at `.,!?;`, trims the segments, and packs them into
-   chunks shorter than `MAX_PHONEME_LENGTH` = 510 **bytes**. A single over-long segment stays one chunk.
-5. **Encode.** `Kokoro::_create_audio` truncates the chunk to 510 bytes. Then `PhonemeEncoder::encode` maps each UTF-8
-   character to its id from `dict/vocab.txt` and **silently drops characters that are not in the vocab**.
+4. **Chunk.** `split_phonemes` (`src/PhonemeChunker.cpp`) splits at `.,!?;`, keeping each mark with the text before
+   it, trims the segments, and packs them into chunks of at most `MAX_PHONEME_LENGTH` = 510 UTF-8 characters, counted
+   as `PhonemeEncoder` splits them (so at most 510 tokens). Punctuation groups are packed whole; an oversized group
+   starts a new chunk before splitting at spaces, so an earlier punctuation boundary takes precedence. A space-free
+   word over the limit (unpunctuated Chinese) is cut after its last tone mark (`→ ↗ ↘ ↓`) within the limit, else at
+   exactly 510 characters. A cut that would strand trailing punctuation retreats to the preceding tone boundary, or
+   reserves the final phoneme and its tone if no earlier tone boundary exists. Nothing is dropped.
+5. **Encode.** `PhonemeEncoder::encode` maps each UTF-8 character to its id from `dict/vocab.txt` and **silently drops
+   characters that are not in the vocab**. More than 510 tokens (unreachable after chunking) are capped at 510 with a
+   `kokoro:` warning on stderr.
    - If the result is empty, no audio is returned and the model is not run. This mirrors upstream, which skips empty
      phoneme strings.
    - Otherwise token 0 is added at both ends.
@@ -325,7 +331,5 @@ The full contract is in [docs/thread-safety.md](docs/thread-safety.md).
 ## Known sharp edges
 
 - `trim_audio` (`src/Kokoro.cpp`) is a no-op, so `trim=true` changes nothing.
-- A single phoneme segment longer than 510 bytes is truncated to 510 bytes. The limit counts bytes, and the encoder
-  drops a UTF-8 character cut at the boundary.
 - Models with an `input_ids` input receive the speed as int32, so its fractional part is lost.
 - Phonemizer output symbols that are missing from `vocab.txt` vanish without a warning.

CMakeLists.txt

@@ -112,6 +112,7 @@ add_library(kokoro_core STATIC
     src/chinese_english/ZHFrontend.cpp
     src/NeuralG2P.cpp
     src/PhonemeText.cpp
+    src/Utf8.cpp
     src/german/GermanG2P.cpp
     src/german/GermanLexicon.cpp
     src/german/GermanPhonemizer.cpp

ISSUES.md

@@ -1,26 +1,3 @@
 # Known issues
 
-## 1. Chunk limit and truncation count UTF-8 bytes, not phonemes
-
-**Where:**
-- `split_phonemes`, `src/PhonemeChunker.cpp:20`: `current_batch.length() + part.length() + 1 >= max_length`.
-- `Kokoro::_create_audio`, `src/Kokoro.cpp:218-219`: `phonemes.substr(0, MAX_PHONEME_LENGTH)`.
-
-`MAX_PHONEME_LENGTH` (510, `src/Kokoro.h:20`) is meant to be a count of phonemes: the model takes at most 510 phoneme tokens plus two padding tokens, and a voice has one style row per length up to 510. Upstream Kokoro counts `len(ps)`, which is Python code points. Both places above use `std::string::length()` and `substr`, which count **bytes**.
-
-**Consequences:**
-- **Chunks are much shorter than they need to be.** Most IPA symbols take more than one byte in UTF-8 (`ʈ ʂ ɕ ŋ` take 2 bytes, tone arrows `→ ↗ ↘ ↓` take 3). For Chinese a 510-byte chunk holds far fewer than 510 phonemes, so text is split into more, shorter chunks than upstream would use. That changes prosody at chunk boundaries and selects style rows for shorter lengths.
-- **Truncation can split a character.** `substr(0, 510)` can cut a multi-byte UTF-8 sequence in half. The tokenizer (`Tokenizer::tokenize` / `split_utf8`) then sees an invalid trailing fragment, which is not in the vocabulary and gets dropped. Truncation only happens when a single chunk is longer than 510 bytes, which is only possible for a segment that `split_phonemes` could not split (one with no `[.,!?;]` for 509+ bytes).
-- **The two limits disagree.** The row index (fixed in `fbf9a6b`) uses the token count, while chunking uses bytes.
-
-**Suggested fix:**
-- Measure chunk length in phoneme tokens (or UTF-8 code points, matching upstream `len(ps)`) instead of bytes, both when packing in `split_phonemes` and when truncating in `_create_audio`.
-- Truncate on a code-point boundary. It is simplest to tokenize first and cap the token vector at `MAX_PHONEME_LENGTH`.
-- A segment longer than the limit should be split at word boundaries (spaces) before any hard truncation. That is what current kokoro-onnx does: its `split_phonemes` "prefer[s] splitting at punctuation marks, then at word boundaries" (`kokoro_onnx/chunker.py`).
-
-**Testing notes:**
-- `tests/test_chunker.cpp` can assert, for a long Chinese phoneme string:
-  - every chunk has at most 510 code points or tokens;
-  - chunks are not cut short at 510 bytes;
-  - no chunk ends in a partial UTF-8 sequence.
-- Expect the `g2p` goldens to be unaffected, because chunking happens after G2P. `synthesis` durations for long inputs will change.
+None open.

src/Kokoro.cpp

@@ -193,16 +193,19 @@ std::pair<std::vector<float>, int> Kokoro::_create_audio(
     const std::vector<float>& voice,
     float speed
 ) {
-    std::string truncated_phonemes = phonemes;
-    if (phonemes.length() > MAX_PHONEME_LENGTH) {
-        truncated_phonemes = phonemes.substr(0, MAX_PHONEME_LENGTH);
-    }
-
-    std::vector<int> tokens_raw = encoder_->encode(truncated_phonemes);
+    std::vector<int> tokens_raw = encoder_->encode(phonemes);
     // Nothing to say: skip the model rather than run it on the padding tokens alone, as upstream
     // Kokoro skips empty phoneme strings (`if not ps: continue`).
     if (tokens_raw.empty()) return {{}, SAMPLE_RATE};
 
+    // split_phonemes keeps chunks within MAX_PHONEME_LENGTH characters, at most one token each: a last resort.
+    const auto max_tokens = static_cast<size_t>(MAX_PHONEME_LENGTH);
+    if (tokens_raw.size() > max_tokens) {
+        std::fprintf(stderr, "kokoro: chunk of %zu phoneme tokens exceeds %zu; dropping the excess\n",
+                     tokens_raw.size(), max_tokens);
+        tokens_raw.resize(max_tokens);
+    }
+
     // Add start and end tokens (0)
     std::vector<int64_t> tokens = {0};
     for (int t : tokens_raw) tokens.push_back(t);

src/Kokoro.h

@@ -17,7 +17,8 @@ class PhonemeEncoder;
 struct KoKoroConfig;
 
 // Constants from config
-const int MAX_PHONEME_LENGTH = 510; // Example value
+// MAX_PHONEME_LENGTH: phoneme tokens per model call; split_phonemes keeps chunks within this many UTF-8 characters.
+const int MAX_PHONEME_LENGTH = 510;
 const int SAMPLE_RATE = 24000;      // Example value
 const int STYLE_DIM = 256;          // floats per style row; a voice is a table of these rows
 

src/PhonemeChunker.cpp

@@ -1,36 +1,140 @@
 #include "PhonemeChunker.h"
 
+#include "Utf8.h"
+
+#include <algorithm>
+#include <array>
 #include <regex>
+#include <string_view>
 
-std::vector<std::string> split_phonemes(const std::string& phonemes, std::size_t max_length) {
-    std::vector<std::string> batches;
-    std::regex re("([.,!?;])");
-    std::sregex_token_iterator it(phonemes.begin(), phonemes.end(), re, {-1, 0}); // -1 for non-match, 0 for match
-    std::sregex_token_iterator end;
-
-    std::string current_batch;
-
-    for (; it != end; ++it) {
-        std::string part = *it;
-        // Removing leading/trailing whitespace
-        part = std::regex_replace(part, std::regex("^\\s+|\\s+$"), "");
-
-        if (part.empty()) continue;
-
-        if (current_batch.length() + part.length() + 1 >= max_length) {
-            if (!current_batch.empty()) batches.push_back(current_batch);
-            current_batch = part;
-        } else {
-             if (std::string(".,!?;").find(part) != std::string::npos) {
-                current_batch += part;
-             } else {
-                if (!current_batch.empty()) current_batch += " ";
-                current_batch += part;
-             }
+namespace {
+
+const std::string kPunctuation = ".,!?;";
+const std::array<std::string_view, 4> kToneMarks = {"→", "↗", "↘", "↓"};
+
+// Regex parts of phonemes split at [.,!?;] (marks kept as their own parts), trimmed, empties dropped.
+std::vector<std::string> segments(const std::string& phonemes) {
+    static const std::regex split_re("([.,!?;])");
+    static const std::regex trim_re("^\\s+|\\s+$");
+    std::vector<std::string> parts;
+    for (std::sregex_token_iterator it(phonemes.begin(), phonemes.end(), split_re, {-1, 0}), end; it != end; ++it) {
+        std::string part = std::regex_replace(it->str(), trim_re, "");
+        if (!part.empty()) parts.push_back(std::move(part));
+    }
+    return parts;
+}
+
+bool is_punctuation(const std::string& segment) {
+    return segment.size() == 1 && kPunctuation.find(segment[0]) != std::string::npos;
+}
+
+// Each punctuation segment joins the group before it; a leading one opens its own group.
+std::vector<std::string> attach_punctuation(const std::vector<std::string>& segments) {
+    std::vector<std::string> groups;
+    for (const std::string& segment : segments) {
+        if (is_punctuation(segment) && !groups.empty()) groups.back() += segment;
+        else groups.push_back(segment);
+    }
+    return groups;
+}
+
+bool is_tone_mark(const std::string& word, std::size_t pos, std::size_t next) {
+    const std::string_view character(word.data() + pos, next - pos);
+    return std::find(kToneMarks.begin(), kToneMarks.end(), character) != kToneMarks.end();
+}
+
+// Prefer the last tone boundary within the window that leaves a phoneme with trailing punctuation. If there is no
+// such boundary, cut at the character limit, moving back to reserve the final phoneme (and its tone) when possible.
+std::size_t hard_cut(const std::string& word, std::size_t begin, std::size_t max_length, std::size_t last_core_start) {
+    std::size_t end = begin;
+    std::size_t after_tone = begin;
+    for (std::size_t n = 0; n < max_length && end < word.size(); ++n) {
+        const std::size_t next = utf8_next(word, end);
+        if (next <= last_core_start && is_tone_mark(word, end, next)) after_tone = next;
+        end = next;
+    }
+    if (end == word.size()) return end;
+    if (after_tone > begin) end = after_tone;
+    if (end > last_core_start && last_core_start > begin) end = last_core_start;
+    return end;
+}
+
+// Byte offset of the final non-punctuation character; a final tone mark also reserves its preceding phoneme.
+std::size_t last_core_char_start(const std::string& word) {
+    const std::size_t core_end = word.find_last_not_of(kPunctuation);
+    if (core_end == std::string::npos) return std::string::npos;
+    std::size_t start = core_end;
+    while (start > 0 && (static_cast<unsigned char>(word[start]) & 0xC0) == 0x80) --start;
+    if (start > 0 && is_tone_mark(word, start, core_end + 1)) {
+        --start;
+        while (start > 0 && (static_cast<unsigned char>(word[start]) & 0xC0) == 0x80) --start;
+    }
+    return start;
+}
+
+// word cut into pieces of at most max_length characters by hard_cut.
+std::vector<std::string> hard_split(const std::string& word, std::size_t max_length) {
+    const std::size_t core_start = last_core_char_start(word);
+    std::vector<std::string> pieces;
+    for (std::size_t begin = 0; begin < word.size();) {
+        const std::size_t end = hard_cut(word, begin, max_length, core_start);
+        pieces.push_back(word.substr(begin, end - begin));
+        begin = end;
+    }
+    return pieces;
+}
+
+// Split an oversized group into space-separated words (empty ones dropped), hard-splitting words over the limit.
+std::vector<std::string> pieces(const std::string& group, std::size_t max_length) {
+    std::vector<std::string> out;
+    std::size_t begin = 0;
+    while (begin <= group.size()) {
+        const std::size_t space = std::min(group.find(' ', begin), group.size());
+        const std::string word = group.substr(begin, space - begin);
+        if (utf8_length(word) > max_length) {
+            const std::vector<std::string> parts = hard_split(word, max_length);
+            out.insert(out.end(), parts.begin(), parts.end());
+        } else if (!word.empty()) {
+            out.push_back(word);
         }
+        begin = space + 1;
     }
-    if (!current_batch.empty()) {
-        batches.push_back(current_batch);
+    return out;
+}
+
+// Append to the current chunk if it fits, else open a new chunk. A zero current_length forces a group boundary.
+void append_piece(std::vector<std::string>& chunks, std::size_t& current_length, const std::string& piece,
+                  std::size_t piece_length, std::size_t max_length) {
+    if (current_length > 0 && current_length + 1 + piece_length <= max_length) {
+        chunks.back() += ' ';
+        chunks.back() += piece;
+        current_length += 1 + piece_length;
+    } else {
+        chunks.push_back(piece);
+        current_length = piece_length;
     }
-    return batches;
+}
+
+// Pack punctuation groups whole when they fit. Start an oversized group in a new chunk before splitting its words.
+std::vector<std::string> pack(const std::vector<std::string>& groups, std::size_t max_length) {
+    std::vector<std::string> chunks;
+    std::size_t current_length = 0;
+    for (const std::string& group : groups) {
+        const std::size_t group_length = utf8_length(group);
+        if (group_length <= max_length) {
+            append_piece(chunks, current_length, group, group_length, max_length);
+            continue;
+        }
+        current_length = 0;
+        for (const std::string& piece : pieces(group, max_length)) {
+            append_piece(chunks, current_length, piece, utf8_length(piece), max_length);
+        }
+    }
+    return chunks;
+}
+
+}  // namespace
+
+std::vector<std::string> split_phonemes(const std::string& phonemes, std::size_t max_length) {
+    return pack(attach_punctuation(segments(phonemes)), max_length);
 }

src/PhonemeChunker.h

@@ -4,6 +4,9 @@
 #include <string>
 #include <vector>
 
-// Splits a phoneme string at [.,!?;] and packs the trimmed segments into chunks, starting a new chunk
-// when the current one would reach max_length bytes. A single segment longer than that stays one chunk.
+// Splits a phoneme string into chunks of at most max_length UTF-8 characters (as PhonemeEncoder counts them, so at
+// most max_length tokens) without dropping any phoneme. Packs [.,!?;] groups whole, keeping each mark with the text
+// before it. An oversized group starts a new chunk before splitting at spaces. A space-free word still too long is
+// cut after its last tone mark (→ ↗ ↘ ↓) within the limit, else at max_length characters. Trailing punctuation
+// reserves the final phoneme and its tone when possible. Segments are trimmed and joined with single spaces.
 std::vector<std::string> split_phonemes(const std::string& phonemes, std::size_t max_length);

src/PhonemeEncoder.cpp

@@ -1,4 +1,5 @@
 #include "PhonemeEncoder.h"
+#include "Utf8.h"
 
 #include <filesystem>
 #include <fstream>
@@ -7,21 +8,6 @@
 
 namespace {
 
-std::vector<std::string> split_utf8(const std::string& str) {
-    std::vector<std::string> chars;
-    for (size_t i = 0; i < str.length();) {
-        const auto c = static_cast<unsigned char>(str[i]);
-        size_t char_len = 1;
-        if ((c & 0xE0) == 0xC0) char_len = 2;
-        else if ((c & 0xF0) == 0xE0) char_len = 3;
-        else if ((c & 0xF8) == 0xF0) char_len = 4;
-        if (i + char_len > str.length()) char_len = str.length() - i;
-        chars.push_back(str.substr(i, char_len));
-        i += char_len;
-    }
-    return chars;
-}
-
 void unescape(std::string& token, const std::string& escape, const char* plain) {
     for (size_t pos = token.find(escape); pos != std::string::npos; pos = token.find(escape, pos + 1)) {
         token.replace(pos, escape.size(), plain);

src/Utf8.cpp

@@ -0,0 +1,28 @@
+#include "Utf8.h"
+
+#include <algorithm>
+
+std::size_t utf8_next(const std::string& s, std::size_t pos) {
+    const auto c = static_cast<unsigned char>(s[pos]);
+    std::size_t len = 1;
+    if ((c & 0xE0) == 0xC0) len = 2;
+    else if ((c & 0xF0) == 0xE0) len = 3;
+    else if ((c & 0xF8) == 0xF0) len = 4;
+    return std::min(pos + len, s.size());
+}
+
+std::size_t utf8_length(const std::string& s) {
+    std::size_t count = 0;
+    for (std::size_t pos = 0; pos < s.size(); pos = utf8_next(s, pos)) ++count;
+    return count;
+}
+
+std::vector<std::string> split_utf8(const std::string& s) {
+    std::vector<std::string> chars;
+    for (std::size_t pos = 0; pos < s.size();) {
+        const std::size_t next = utf8_next(s, pos);
+        chars.push_back(s.substr(pos, next - pos));
+        pos = next;
+    }
+    return chars;
+}

src/Utf8.h

@@ -0,0 +1,17 @@
+#pragma once
+
+#include <cstddef>
+#include <string>
+#include <vector>
+
+// UTF-8 characters as PhonemeEncoder tokenizes them: the lead byte alone sets the length (1-4 bytes), sequences are
+// not validated, and a sequence truncated by the end of the string ends there.
+
+// Byte offset just past the character starting at pos (pos < s.size()).
+std::size_t utf8_next(const std::string& s, std::size_t pos);
+
+// Number of characters in s.
+std::size_t utf8_length(const std::string& s);
+
+// s split into its characters.
+std::vector<std::string> split_utf8(const std::string& s);