fix: preserve Unicode token chunk offsets

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,

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))

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)