diff --git a/src/chonkie/chunker/token.py b/src/chonkie/chunker/token.py index 095bcca43..b23418b86 100644 --- a/src/chonkie/chunker/token.py +++ b/src/chonkie/chunker/token.py @@ -4,6 +4,7 @@ """ +from bisect import bisect_right from typing import Generator, Sequence, Union from tqdm import trange @@ -102,6 +103,73 @@ def _create_chunks( return chunks + def _create_chunks_with_offsets( + self, + text: str, + tokens: Sequence[int], + offsets: Sequence[tuple[int, int]], + ) -> list[Chunk]: + """Create character-safe chunks from tokenizer offsets.""" + safe_boundaries = [0] + for index in range(1, len(tokens)): + previous_start, previous_end = offsets[index - 1] + current_start, _ = offsets[index] + if previous_end > previous_start and previous_end == current_start: + safe_boundaries.append(index) + safe_boundaries.append(len(tokens)) + + chunks = [] + start = 0 + while start < len(tokens): + requested_end = min(start + self.chunk_size, len(tokens)) + end = safe_boundaries[bisect_right(safe_boundaries, requested_end) - 1] + if end <= start: + end = safe_boundaries[bisect_right(safe_boundaries, start)] + + start_index = offsets[start][0] + end_index = offsets[end - 1][1] + chunks.append( + Chunk( + text=text[start_index:end_index], + start_index=start_index, + end_index=end_index, + token_count=end - start, + ) + ) + + if end == len(tokens): + break + + requested_start = max(0, end - self.chunk_overlap) + next_start = safe_boundaries[bisect_right(safe_boundaries, requested_start) - 1] + if next_start <= start: + next_start = safe_boundaries[bisect_right(safe_boundaries, start)] + start = next_start + + return chunks + + def _chunk_with_offsets(self, text: str) -> tuple[list[Chunk], int] | None: + """Chunk text with source offsets when the tokenizer provides them.""" + encode_with_offsets = getattr(self.tokenizer, "encode_with_offsets", None) + if callable(encode_with_offsets): + try: + offset_tokens, offsets = encode_with_offsets(text) + except (NotImplementedError, ValueError): + return None + else: + if offset_tokens and len(offsets) == len(offset_tokens): + chunks = self._create_chunks_with_offsets(text, offset_tokens, offsets) + return chunks, len(offset_tokens) + + return None + + def _chunk_tokens(self, tokens: Sequence[int]) -> list[Chunk]: + """Chunk token IDs using the decode-based fallback.""" + token_groups = list(self._token_group_generator(tokens)) + token_counts = [len(token_group) for token_group in token_groups] + chunk_texts = self.tokenizer.decode_batch(token_groups) + return self._create_chunks(chunk_texts, token_groups, token_counts) + def _token_group_generator(self, tokens: Sequence[int]) -> Generator[list[int], None, None]: """Generate chunks from a list of tokens.""" for start in range(0, len(tokens), self.chunk_size - self.chunk_overlap): @@ -125,47 +193,37 @@ def chunk(self, text: str) -> list[Chunk]: logger.debug(f"Chunking text of length {len(text)} with chunk_size={self.chunk_size}") - # Encode full text - text_tokens = self.tokenizer.encode(text) - - # Calculate token groups and counts - token_groups = list(self._token_group_generator(text_tokens)) - token_counts = [len(toks) for toks in token_groups] - - # decode the token groups into the chunk texts - chunk_texts = self.tokenizer.decode_batch(token_groups) - - # Create the chunks from the token groups and token counts - chunks = self._create_chunks(chunk_texts, token_groups, token_counts) + offset_chunks = self._chunk_with_offsets(text) + if offset_chunks is not None: + chunks, token_count = offset_chunks + else: + text_tokens = self.tokenizer.encode(text) + chunks = self._chunk_tokens(text_tokens) + token_count = len(text_tokens) - logger.info(f"Created {len(chunks)} chunks from {len(text_tokens)} tokens") + logger.info(f"Created {len(chunks)} chunks from {token_count} tokens") return chunks def _process_batch(self, texts: list[str]) -> list[list[Chunk]]: """Process a batch of texts.""" - # encode the texts into tokens in a batch - tokens_list = self.tokenizer.encode_batch(texts) - result: list = [] - - for tokens in tokens_list: - if not tokens: - result.append([]) - continue - - # get the token groups - token_groups = list(self._token_group_generator(tokens)) - - # get the token counts - token_counts = [len(token_group) for token_group in token_groups] - - # decode the token groups into the chunk texts - chunk_texts = self.tokenizer.decode_batch(token_groups) - - # create the chunks from the token groups and token counts - chunks = self._create_chunks(chunk_texts, token_groups, token_counts) - result.append(chunks) - - return result + result: list[list[Chunk] | None] = [None] * len(texts) + fallback_indices = [] + fallback_texts = [] + + for index, text in enumerate(texts): + offset_chunks = self._chunk_with_offsets(text) + if offset_chunks is not None: + result[index] = offset_chunks[0] + else: + fallback_indices.append(index) + fallback_texts.append(text) + + if fallback_texts: + tokens_list = self.tokenizer.encode_batch(fallback_texts) + for index, tokens in zip(fallback_indices, tokens_list): + result[index] = self._chunk_tokens(tokens) if tokens else [] + + return [chunks if chunks is not None else [] for chunks in result] def chunk_batch( # ty: ignore[invalid-method-override] self, diff --git a/src/chonkie/tokenizer.py b/src/chonkie/tokenizer.py index 9a4880520..ea106a1bc 100644 --- a/src/chonkie/tokenizer.py +++ b/src/chonkie/tokenizer.py @@ -2,6 +2,7 @@ import inspect from abc import ABC, abstractmethod +from bisect import bisect_right from collections import defaultdict from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, Any, Protocol @@ -25,6 +26,20 @@ } +def _byte_offsets_to_character_offsets( + text: str, byte_offsets: Sequence[tuple[int, int]] +) -> list[tuple[int, int]]: + """Convert UTF-8 byte offsets into offsets into a Python string.""" + character_boundaries = [0] + for character in text: + character_boundaries.append(character_boundaries[-1] + len(character.encode("utf-8"))) + + def character_offset(byte_offset: int) -> int: + return bisect_right(character_boundaries, byte_offset) - 1 + + return [(character_offset(start), character_offset(end)) for start, end in byte_offsets] + + class TokenizerProtocol(Protocol): """Protocol defining the interface for tokenizers. @@ -600,6 +615,15 @@ class TiktokenAutoTokenizer(AutoTokenizer): if TYPE_CHECKING: tokenizer: tiktoken.Encoding + def encode_with_offsets(self, text: str) -> tuple[list[int], list[tuple[int, int]]]: + """Encode text and return character offsets for each token.""" + token_ids = list(self.encode(text)) + decoded, starts = self.tokenizer.decode_with_offsets(token_ids) + if decoded != text: + raise ValueError("Tokenizer did not round-trip the input text.") + ends = [*starts[1:], len(text)] + return token_ids, list(zip(starts, ends)) + class TransformersAutoTokenizer(AutoTokenizer): """Adapter for HuggingFace `transformers` tokenizers.""" @@ -618,6 +642,11 @@ def encode_batch(self, texts: Sequence[str]) -> Sequence[Sequence[int]]: encoded = self.tokenizer(texts, add_special_tokens=False) return encoded["input_ids"] + def encode_with_offsets(self, text: str) -> tuple[list[int], list[tuple[int, int]]]: + """Encode text and return character offsets for each token.""" + encoded = self.tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) + return list(encoded["input_ids"]), [tuple(offset) for offset in encoded["offset_mapping"]] + def decode_batch(self, token_sequences: Sequence[Sequence[int]]) -> Sequence[str]: """Batch decode using batch_decode method.""" return self.tokenizer.batch_decode( @@ -644,6 +673,11 @@ def encode_batch(self, texts: Sequence[str]) -> Sequence[Sequence[int]]: for encoding in self.tokenizer.encode_batch(texts, add_special_tokens=False) ] + def encode_with_offsets(self, text: str) -> tuple[list[int], list[tuple[int, int]]]: + """Encode text and return character offsets for each token.""" + encoding = self.tokenizer.encode(text, add_special_tokens=False) + return list(encoding.ids), list(encoding.offsets) + class TokieAutoTokenizer(AutoTokenizer): """Adapter for tokie tokenizers.""" @@ -657,6 +691,11 @@ def encode(self, text: str) -> list[int]: """Encode text and extract token IDs.""" return self.tokenizer.encode(text, add_special_tokens=False).ids + def encode_with_offsets(self, text: str) -> tuple[list[int], list[tuple[int, int]]]: + """Encode text and return character offsets for each token.""" + encoding = self.tokenizer.encode_with_offsets(text, add_special_tokens=False) + return list(encoding.ids), _byte_offsets_to_character_offsets(text, encoding.offsets) + def decode(self, tokens: Sequence[int]) -> str: """Decode token IDs back to text.""" result = self.tokenizer.decode(list(tokens)) diff --git a/tests/chunkers/test_token_chunker.py b/tests/chunkers/test_token_chunker.py index 37e04ba36..206d1d7bd 100644 --- a/tests/chunkers/test_token_chunker.py +++ b/tests/chunkers/test_token_chunker.py @@ -8,6 +8,7 @@ import tiktoken from tiktoken import Encoding from tokenizers import Tokenizer +from tokie import Tokenizer as TokieTokenizer from transformers import AutoTokenizer, PreTrainedTokenizerFast from chonkie import Chunk, TokenChunker @@ -37,6 +38,15 @@ def tokenizer() -> Tokenizer: pytest.skip(f"Could not load tokenizers tokenizer: {e}") +@pytest.fixture +def tokie_tokenizer() -> TokieTokenizer: + """Fixture that returns a GPT-2 tokenizer from the tokie library.""" + try: + return TokieTokenizer.from_pretrained("openai-community/gpt2") + except (OSError, ValueError) as e: + pytest.skip(f"Could not load tokie tokenizer: {e}") + + @pytest.fixture def sample_text() -> str: """Fixture that returns a sample text for testing the TokenChunker.""" @@ -275,6 +285,30 @@ def test_token_chunker_indices_complex_md(sample_complex_markdown_text: str) -> verify_chunk_indices(chunks, sample_complex_markdown_text) +@pytest.mark.parametrize( + "tokenizer_fixture", ["tiktokenizer", "tokenizer", "tokie_tokenizer", "transformers_tokenizer"] +) +def test_token_chunker_preserves_unicode_source_offsets( + request: pytest.FixtureRequest, tokenizer_fixture: str +) -> None: + """Test that token boundaries cannot split a multi-byte character.""" + text = "a🩺 hello world" + tokenizer = request.getfixturevalue(tokenizer_fixture) + chunker = TokenChunker(tokenizer=tokenizer, chunk_size=2, chunk_overlap=0) + + chunks = chunker.chunk(text) + + assert [chunk.text for chunk in chunks] == ["a", "🩺", " hello world"] + assert [chunk.token_count for chunk in chunks] == [1, 3, 2] + assert all(chunk.end_index > chunk.start_index for chunk in chunks) + verify_chunk_indices(chunks, text) + + for batch_chunks in chunker.chunk_batch([text, text]): + assert [chunk.text for chunk in batch_chunks] == ["a", "🩺", " hello world"] + assert all(chunk.end_index > chunk.start_index for chunk in batch_chunks) + verify_chunk_indices(batch_chunks, text) + + def test_token_chunker_token_counts(tiktokenizer: Encoding, sample_text: str) -> None: """Test that the TokenChunker correctly calculates token counts.""" chunker = TokenChunker(tokenizer=tiktokenizer, chunk_size=512, chunk_overlap=128)