Repository navigation
Fix SyntheticDataKit.chunk_data emitting chunks over max_tokens - #7073
Conversation
There was a problem hiding this comment.
Code Review
This pull request fixes an off-by-one error in the synthetic data chunking logic where np.linspace was called with n_chunks instead of n_chunks + 1 points, which previously resulted in fewer, oversized chunks that exceeded the maximum token limit. A corresponding unit test test_chunk_data_chunks_do_not_exceed_max_tokens has been added to verify that documents are correctly split and chunks remain within the token limit. There are no review comments, so we have no feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f2629df0c9
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # Emit n_chunks ranges: pairing boundaries[:-1] with boundaries[1:] turns | ||
| # N points into N-1 ranges, so linspace needs n_chunks + 1 points. Using | ||
| # n_chunks points produced one fewer, oversized chunk (exceeding max_tokens). | ||
| boundaries = np.ceil(np.linspace(0, length - self.overlap, n_chunks + 1)).astype(int) |
There was a problem hiding this comment.
Avoid emitting an extra chunk near stride boundaries
Because n_chunks is still computed as ceil(length / (max_tokens - overlap)), changing linspace to emit exactly that many ranges over-splits documents whose length is within one overlap window after a stride multiple. For example, with max_tokens=400, overlap=64, and length=673, this now writes 3 chunks of about 267 tokens even though 2 chunks of at most 400 tokens cover the whole document; each extra chunk becomes another downstream generation input with duplicated context. If the intent is the minimum number of bounded chunks, the count needs to be based on the non-overlapped span, e.g. (length - overlap) / (max_tokens - overlap), when using n_chunks + 1 boundary points.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch — addressed in 3e49790. n_chunks now uses ceil((length - overlap) / stride) so the count is the minimum number of overlap-bounded chunks rather than ceil(length / stride). The length=673 case now yields 2 chunks (~369/368) instead of 3 of ~267.
I swept lengths 401–3000 (max_tokens=400, overlap=64) and confirmed every chunk stays <= max_tokens with full coverage and no gaps, and that n_chunks >= 2 still holds so the np.stack pairing is valid. The advertised examples (500/1000/2000) are unchanged. Added test_chunk_data_does_not_over_split as a regression (fails before at 3 chunks, passes after at 2).
The multi-chunk path built boundaries from np.linspace(..., n_chunks), but pairing boundaries[:-1] with boundaries[1:] turns N points into N-1 ranges, so it produced one fewer, oversized chunk: every chunk exceeded max_tokens and a document just over the threshold came back as a single unsplit chunk. Use n_chunks + 1 points so exactly n_chunks ranges are emitted, each within max_tokens. Also base n_chunks on the non-overlapped span: consecutive chunks overlap by overlap, so covering length needs ceil((length - overlap) / stride) chunks, not ceil(length / stride). The looser count over-counted by one just past a stride multiple (a 673-token doc became 3 chunks of ~267 instead of 2 of ~369), emitting an extra redundant chunk. Coverage and overlap are unchanged and every chunk still stays within max_tokens.
f2629df to
3e49790
Compare
|
Pushed a small follow-up (comment-only) on top of your change:
Verified the fix independently: an exhaustive sweep of the boundary math (hundreds of thousands of (max_tokens, overlap, length) cases plus a full-function fuzz) shows every chunk stays within max_tokens with minimal count, exact overlap and full coverage, and the result is stable across numpy 1.24 through 2.5. AST diff confirms this push is comment-only. Thanks for the clean fix. |
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Didn't find any major issues. Swish! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
Thanks! |
Summary
SyntheticDataKit.chunk_datasplits a document into chunks bounded bymax_tokens(max_seq_length - 2*max_generation_tokens - 128). The multi-chunk path builds boundaries withnp.linspace(0, length - overlap, n_chunks)— butn_chunksis the desired chunk count, whilelinspacereceives it as the number of points. Pairingboundaries[:-1]withboundaries[1:]turns N points into N-1 ranges, so the function emits one fewer, oversized chunk. Every chunk ends up larger thanmax_tokens, and a document just over the threshold is returned as a single unsplit chunk.The single-chunk branch just above already documents this exact off-by-one — "the linspace/stack pairing emits one fewer range than boundary points" — it was only worked around for the
length <= max_tokenscase, leaving the multi-chunk path broken.Example (
max_tokens=400,overlap=64, 1 token per word):Every multi-chunk case currently exceeds
max_tokens.Fix
Use
n_chunks + 1linspace points so the[:-1]/[1:]pairing yields exactlyn_chunksranges, each withinmax_tokens. Overlap is preserved (n_chunks >= 2in this branch, son_chunks + 1 >= 3points is always valid fornp.stack).Test
Extended
tests/test_synthetic_chunk_data.pywithtest_chunk_data_chunks_do_not_exceed_max_tokens, asserting every produced chunk is<= max_tokensfor 500- and 2000-token documents. It fails before the fix (chunks over budget / a doc returned unsplit) and passes after. The existingtest_chunk_data_still_splits_long_documentonly checkedlen(out) > 1, so it missed this.The chunk-boundary logic is pure
numpy; I verified fail-before/pass-after by running the exact algorithm with the test's mock tokenizer standalone (importing the full module requires a GPU environment).ruff checkis clean.No linked issue — found by reading the code.
Update (3e49790)
Also tightened the chunk count in the same path (raised in review):
n_chunkswasceil(length / stride), but because consecutive chunks overlap byoverlap, coveringlengthneedsceil((length - overlap) / stride)chunks. The looser count over-counted by one just past a stride multiple — e.g. a 673-token doc produced 3 chunks of ~267 where 2 of ~369 suffice — emitting an extra redundant generation input. Now uses the non-overlapped span; coverage and overlap are unchanged and every chunk still stays withinmax_tokens. Addedtest_chunk_data_does_not_over_split(fails before at 3 chunks, passes after at 2). The advertised 500/1000/2000 examples are unchanged.