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,70 @@ 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_tokens(self, text: str, tokens: Sequence[int]) -> list[Chunk]:
+        """Chunk tokens, using 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):
+                # Some tokenizer backends do not expose offset mappings. Keep
+                # their existing decode-based behavior in that case.
+                pass
+            else:
+                if list(tokens) == offset_tokens and len(offsets) == len(tokens):
+                    return self._create_chunks_with_offsets(text, tokens, offsets)
+
+        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):
@@ -128,15 +193,7 @@ def chunk(self, text: str) -> list[Chunk]:
         # 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)
+        chunks = self._chunk_tokens(text, text_tokens)
 
         logger.info(f"Created {len(chunks)} chunks from {len(text_tokens)} tokens")
         return chunks
@@ -147,23 +204,12 @@ def _process_batch(self, texts: list[str]) -> list[list[Chunk]]:
         tokens_list = self.tokenizer.encode_batch(texts)
         result: list = []
 
-        for tokens in tokens_list:
+        for text, tokens in zip(texts, 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)
+            result.append(self._chunk_tokens(text, tokens))
 
         return result
 

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 encoding.ids, 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 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

@@ -275,6 +275,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", "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)