diff --git a/CHANGELOG.md b/CHANGELOG.md index 8142bc16f..facd8f7f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,9 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm ### Added - `pw.xpacks.llm.rerankers.LLMReranker` now accepts `call_kwargs`, the kwargs passed to each call of the LLM. The default is still `{"temperature": 0}`; pass `call_kwargs={}` to use the reranker with models that accept only the default `temperature`, such as Claude Opus 4.7 and newer (including Claude Opus 5.5) on Bedrock. +### Fixed +- `pathway.xpacks.llm.splitters.TokenCountSplitter` no longer skips source text when a punctuation cut falls inside a token. It locates the cut using token byte lengths and advances by exactly the whole source tokens emitted in the chunk. + ## [0.33.0] - 2026-09-18 ### Changed diff --git a/python/pathway/xpacks/llm/splitters.py b/python/pathway/xpacks/llm/splitters.py index ea8457af7..3d94f70f9 100644 --- a/python/pathway/xpacks/llm/splitters.py +++ b/python/pathway/xpacks/llm/splitters.py @@ -259,6 +259,7 @@ def chunk(self, text: str, metadata: dict = {}, **kwargs) -> list[tuple[str, dic while i < len(tokens): chunk_tokens = tokens[i : i + max_tokens] chunk = tokenizer.decode(chunk_tokens) + consumed = len(chunk_tokens) last_punctuation = max( [chunk.rfind(p) for p in self.PUNCTUATION], default=-1 ) @@ -266,8 +267,24 @@ def chunk(self, text: str, metadata: dict = {}, **kwargs) -> list[tuple[str, dic last_punctuation != -1 and last_punctuation > self.CHARS_PER_TOKEN * min_tokens ): - chunk = chunk[: last_punctuation + 1] - i += len(tokenizer.encode_ordinary(chunk)) + cut_bytes = len(chunk[: last_punctuation + 1].encode("utf-8")) + token_bytes = 0 + cut_tokens = 0 + # Token boundaries can split a UTF-8 character. Count bytes rather + # than decoding every growing prefix to locate the punctuation cut. + for token in chunk_tokens: + token_bytes += len(tokenizer.decode_single_token_bytes(token)) + if token_bytes > cut_bytes: + break + cut_tokens += 1 + # If the first token straddles the cut, keep the full window so + # that the emitted text still matches the tokens we consume. + if cut_tokens: + consumed = cut_tokens + chunk = tokenizer.decode(chunk_tokens[:consumed]) + # Guard against a chunk that would not advance the cursor, which would + # otherwise stall the loop forever. + i += max(consumed, 1) output.append((chunk, metadata)) return output diff --git a/python/pathway/xpacks/llm/tests/test_splitters.py b/python/pathway/xpacks/llm/tests/test_splitters.py index 88bf4b1de..afaa24ebc 100644 --- a/python/pathway/xpacks/llm/tests/test_splitters.py +++ b/python/pathway/xpacks/llm/tests/test_splitters.py @@ -3,6 +3,7 @@ from __future__ import annotations import pandas as pd +import pytest import pathway as pw from pathway.tests.utils import assert_table_equality @@ -31,6 +32,49 @@ def test_tokencount(): assert_table_equality(result, input_table) +def test_tokencount_does_not_drop_characters(): + # Regression: when a chunk was cut at a punctuation mark, the cursor advanced + # by the re-encoded kept text, which could skip the tokens straddling the cut + # and silently drop characters. Concatenating the chunks must reconstruct the + # (unicode-normalized) input exactly. + import unicodedata + + splitter = TokenCountSplitter(min_tokens=1, max_tokens=3) + txt = "a.b.c.d.e.f.g.h.i.j.k.l." + chunks = [chunk for chunk, _ in splitter.chunk(txt)] + + assert "".join(chunks) == unicodedata.normalize("NFKC", txt) + + +@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) + + +def test_tokencount_preserves_token_straddling_first_punctuation_cut(): + # cl100k_base encodes "...)" as a single token. If the + # punctuation prefix has no whole token, emit the window intact. + splitter = TokenCountSplitter(min_tokens=0, max_tokens=1) + txt = "...) tail" + metadata = {"source": "example"} + chunks = splitter.chunk(txt, metadata) + + assert "".join(chunk for chunk, _ in chunks) == txt + assert all(chunk and meta == metadata for chunk, meta in chunks) + + def test_recursive_from_encoding(): splitter = RecursiveSplitter( encoding_name="cl100k_base", chunk_size=30, chunk_overlap=0