fix: preserve Unicode token chunk offsets - #646
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
📝 WalkthroughWalkthrough
ChangesUnicode-safe token chunking
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant TokenChunker
participant TokenizerAdapter
participant SourceText
participant Chunk
TokenChunker->>TokenizerAdapter: encode_with_offsets(text)
TokenizerAdapter-->>TokenChunker: token IDs and character offsets
TokenChunker->>SourceText: select safe token-boundary spans
SourceText-->>TokenChunker: sliced chunk text
TokenChunker->>Chunk: store text, token count, indices, and overlap
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/chonkie/tokenizer.py (1)
676-680: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueCast
.idstolistfor consistency with the other adapters.
TiktokenAutoTokenizer.encode_with_offsetsandTransformersAutoTokenizer.encode_with_offsetsboth explicitly wrap their token IDs inlist(...).TokenizersAutoTokenizer.encode_with_offsets(Line 679) andTokieAutoTokenizer.encode_with_offsets(Line 697) returnencoding.idsdirectly. The consumer inchunker/token.py(_chunk_tokens) doeslist(tokens) == offset_tokens. If either binding's.idsis not a plainlist, this equality silently fails and the offset path silently falls back to decode-based chunking on every call, defeating the feature without any error.Cast explicitly for consistency and defensive robustness.
♻️ Proposed fix
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 + return list(encoding.ids), list(encoding.offsets)Also applies to: 694-698
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/chonkie/tokenizer.py` around lines 676 - 680, Update encode_with_offsets in TokenizersAutoTokenizer and TokieAutoTokenizer to wrap encoding.ids in list(...) before returning it, while preserving the existing offsets unchanged and maintaining the declared tuple return shape.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/chonkie/chunker/token.py`:
- Around line 106-169: Restructure the tokenization flow so offset-aware
tokenizers are attempted first, using the tokens and offsets returned by
encode_with_offsets as the working data without re-encoding. Update chunk() and
_process_batch to avoid unconditional encode/encode_batch calls, while
preserving batched tokenization where offsets are unavailable or fail. Adjust
_chunk_tokens to consume the successful offset result directly and retain the
existing decode-based fallback.
In `@tests/chunkers/test_token_chunker.py`:
- Around line 278-280: Extend the tokenizer parametrization for
test_unicode_boundary_offsets to include the tokie fixture if
TokieAutoTokenizer.encode_with_offsets supports the expected offset API;
otherwise explicitly mark or document the fixture as unsupported for this test.
---
Nitpick comments:
In `@src/chonkie/tokenizer.py`:
- Around line 676-680: Update encode_with_offsets in TokenizersAutoTokenizer and
TokieAutoTokenizer to wrap encoding.ids in list(...) before returning it, while
preserving the existing offsets unchanged and maintaining the declared tuple
return shape.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 2b54765b-c8ef-4189-b493-141bca4f9e49
📒 Files selected for processing (3)
src/chonkie/chunker/token.pysrc/chonkie/tokenizer.pytests/chunkers/test_token_chunker.py
| 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) | ||
|
|
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift
Offset-supporting tokenizers now tokenize the text twice per chunk, and batching is lost for the offset path.
chunk() (Line 194) always calls self.tokenizer.encode(text) first. Then, when the tokenizer supports offsets, _chunk_tokens calls self.tokenizer.encode_with_offsets(text) again (Line 156). Every offset-aware adapter re-runs a full tokenization internally: TiktokenAutoTokenizer.encode_with_offsets calls self.encode(text) plus a full decode_with_offsets; TransformersAutoTokenizer.encode_with_offsets and TokenizersAutoTokenizer.encode_with_offsets each re-invoke the underlying tokenizer. The first-pass tokens value is then only used for an equality check, not as the working token list on the offset path.
This doubles tokenization cost for exactly the adapters this PR targets (tiktoken, tokenizers, transformers fast). For batch chunking it is worse: _process_batch calls self.tokenizer.encode_batch(texts) once (batched), then _chunk_tokens calls encode_with_offsets per text individually — the batching benefit of encode_batch is lost for every text that takes the offset path.
Consider restructuring so encode_with_offsets is tried first as the primary tokenization step, and the plain encode/encode_batch call is only made when offsets are unavailable or fail, instead of always doing both.
♻️ Illustrative restructuring (adjust call sites accordingly)
- 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 _chunk_text_with_tokens(self, text: str, tokens: Optional[Sequence[int]]) -> tuple[list[Chunk], int]:
+ """Chunk text, preferring a single offset-aware tokenization pass."""
+ 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):
+ pass
+ else:
+ if len(offsets) == len(offset_tokens):
+ return self._create_chunks_with_offsets(text, offset_tokens, offsets), len(offset_tokens)
+
+ tokens = tokens if tokens is not None else self.tokenizer.encode(text)
+ 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), len(tokens)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/chonkie/chunker/token.py` around lines 106 - 169, Restructure the
tokenization flow so offset-aware tokenizers are attempted first, using the
tokens and offsets returned by encode_with_offsets as the working data without
re-encoding. Update chunk() and _process_batch to avoid unconditional
encode/encode_batch calls, while preserving batched tokenization where offsets
are unavailable or fail. Adjust _chunk_tokens to consume the successful offset
result directly and retain the existing decode-based fallback.
| @pytest.mark.parametrize( | ||
| "tokenizer_fixture", ["tiktokenizer", "tokenizer", "transformers_tokenizer"] | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check for any existing tokie-backed test coverage of encode_with_offsets / byte-offset conversion.
rg -n -C5 'tokie' --type=py -g '!*/node_modules/*'Repository: feyninc/chonkie
Length of output: 153
Add tokie coverage for Unicode byte-offset conversion.
No Python tests reference tokie, so TokieAutoTokenizer.encode_with_offsets is not covered here. Add a tokie tokenizer fixture parametrization for test_unicode_boundary_offsets if that adapter supports the same offset API, otherwise document/exclude it explicitly.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/chunkers/test_token_chunker.py` around lines 278 - 280, Extend the
tokenizer parametrization for test_unicode_boundary_offsets to include the tokie
fixture if TokieAutoTokenizer.encode_with_offsets supports the expected offset
API; otherwise explicitly mark or document the fixture as unsupported for this
test.
1 participant
Summary
Root cause
TokenChunker decoded each token group independently and accumulated decoded-string lengths. A group beginning in the middle of a UTF-8 character could decode as empty, which lost text and desynchronized later indices.
Validation
Summary by CodeRabbit