fix(xpacks/llm): stop TokenCountSplitter dropping characters at punctuation cuts

[linhongyu510]

TokenCountSplitter can silently drop source characters when a punctuation cut falls inside a token. For example, splitting "a.b.c.d.e.f.g.h.i.j.k.l." with min_tokens=1, max_tokens=3 and rejoining the chunks loses f, i, and l on the base implementation.

The splitter now locates the punctuation cut by accumulating decode_single_token_bytes() lengths and consumes exactly the whole source tokens it emits. The token straddling the cut stays for the next chunk. If no whole token fits before the cut, the full window is retained, preserving forward progress and text such as ...).

This follows the byte-offset approach requested in review. It avoids decoding partial UTF-8 prefixes and removes the growing-prefix decoding loop. Public arguments and metadata propagation are unchanged. The changelog entry is under Unreleased.

Validation:

• Full splitter module: 8 passed with the real Pathway engine and tiktoken 0.14.0.
• Added Russian and Chinese reconstruction cases: both fail on the previous PR head and pass with this change.
• Added a first-token punctuation case: the previous head drops ) from ...) tail; this change preserves the text and metadata.
• The original punctuation-loss regression continues to pass.
• Black 24.10.0, isort, Flake8, and git diff --check pass. The test run reports one existing Pydantic deprecation warning from prompts.py.

Local performance comparison on a fixed 128,000-character English corpus, seven warm runs per case (median, milliseconds):

max_tokens	Base implementation	Previous PR head	This change
128	7.145	26.029	6.524
500 (default)	6.374	67.690	4.745
2000	5.655	230.502	5.146

All timed outputs reconstruct that corpus exactly. These are local microbenchmark results on macOS arm64 / Python 3.11.15, not end-to-end RAG timings. The comparison uses base accf4a40 and previous PR head 24287e57.

AI assistance was used for implementation and local verification. The change stays focused on punctuation-cut behavior; arbitrary hard token cuts through a multibyte character are outside this fix.

[zxqfd555]

Thanks for the PR! The bug is real, but the fix has two problems, both coming from kept.startswith(tokenizer.decode(chunk_tokens[:n])).

1. It breaks on non-ASCII text. A token boundary can fall in the middle of a multi-byte character; decoding such a prefix yields U+FFFD, startswith fails and the loop stops early. The cursor then advances by one token while the whole kept text is emitted, so the output repeats the input many times over. The current code handles these inputs correctly, so this is a regression. Both cases below pass on main and fail on this branch:

@pytest.mark.parametrize(
    "txt",
    [
        "Привет, мир. Это проверка разбиения текста! Работает ли оно? Да. " * 5,
        "你好,世界。这是一个测试!它有效吗?是的. " * 10,
    ],
    ids=["russian", "chinese"],
)
def test_tokencount_does_not_duplicate_non_ascii(txt):
    import unicodedata

    splitter = TokenCountSplitter()
    chunks = [chunk for chunk, _ in splitter.chunk(txt)]

    assert "".join(chunks) == unicodedata.normalize("NFKC", txt)

2. It is quadratic in max_tokens. Every prefix of the window is decoded for every chunk. With the default parameters, splitting plain English text takes about 7x longer than the current code (128k characters: 28 ms -> 200 ms).

Both go away if the cut is located by byte offsets: accumulate len(tokenizer.decode_single_token_bytes(t)) over the window's tokens and stop at the last one that still fits in len(kept.encode()). Could you rework the fix this way and add the test above?

[linhongyu510]

Implemented the requested byte-offset approach in df560d3d4152. The splitter accumulates decode_single_token_bytes() lengths up to the UTF-8 byte length of the punctuation cut, avoiding partial-prefix decoding and the growing-prefix loop.

The regression coverage includes your Russian and Chinese reconstruction cases, plus a first-token punctuation case (...) tail). If no whole token fits before the punctuation cut, the full window is emitted so the consumed tokens still match the emitted text.

All 8 reported checks on this head currently pass. The PR description records the earlier local tests, the fixed-corpus timing comparison, and the remaining hard-token-cut boundary. AI assistance is disclosed there. Thanks for identifying both regressions.

[linhongyu510]

Addressed in df560d3d. Both points from this review are fixed by locating the cut with byte offsets instead of kept.startswith(...):

1. Non-ASCII duplication regression — fixed. The cut is now found by accumulating len(tokenizer.decode_single_token_bytes(t)) over the window's tokens and stopping at the last whole token that still fits in len(kept.encode("utf-8")). No partial-UTF-8 prefix is decoded, so a token boundary inside a multi-byte character can no longer yield U+FFFD and stop the loop early.

2. Quadratic behavior — fixed. The growing-prefix decode loop is gone; each window walks its tokens once (linear in max_tokens).

Regression tests added (exactly the RU/CN cases from the review, plus a first-token ...) case):

• test_tokencount_does_not_duplicate_non_ascii[russian]
• test_tokencount_does_not_duplicate_non_ascii[chinese]
• test_tokencount_preserves_token_straddling_first_punctuation_cut

Before/after evidence (local, tiktoken 0.14.0, Python 3.11):

• On the previous PR head 24287e57 (with these tests applied): 3 failed, 5 passed — the two non-ASCII cases and the ...) case fail.
• On this head df560d3d: 8 passed, 0 failed.

black 24.10.0 --check, isort --check-only, flake8, and git diff --check are clean on the two changed source files. Public arguments and metadata propagation are unchanged; only how the punctuation cut is located changed.

Could you take another look? No approval assumed — awaiting re-review. AI assistance was used for implementation and local verification.