Skip to content

Fix SyntheticDataKit.chunk_data emitting chunks over max_tokens - #7073

Merged
danielhanchen merged 2 commits into
unslothai:mainfrom
winklemad:fix/synthetic-chunk-max-tokens
Jul 12, 2026
Merged

danielhanchen merged 2 commits into
unslothai:mainfrom
winklemad:fix/synthetic-chunk-max-tokens

Conversation

@winklemad

@winklemad winklemad commented Jul 11, 2026 •

Copy link
Copy Markdown
Contributor

Summary

SyntheticDataKit.chunk_data splits a document into chunks bounded by max_tokens (max_seq_length - 2*max_generation_tokens - 128). The multi-chunk path builds boundaries with np.linspace(0, length - overlap, n_chunks) — but n_chunks is the desired chunk count, while linspace receives it as the number of points. Pairing boundaries[:-1] with boundaries[1:] turns N points into N-1 ranges, so the function emits one fewer, oversized chunk. Every chunk ends up larger than max_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_tokens case, leaving the multi-chunk path broken.

Example (max_tokens=400, overlap=64, 1 token per word):

doc tokens current fixed
500 1 chunk of 500 (unsplit) 2 chunks of 282/282
1000 2 chunks of 532 3 chunks of 376
2000 5 chunks of ~451 6 chunks of ~387

Every multi-chunk case currently exceeds max_tokens.

Fix

Use n_chunks + 1 linspace points so the [:-1]/[1:] pairing yields exactly n_chunks ranges, each within max_tokens. Overlap is preserved (n_chunks >= 2 in this branch, so n_chunks + 1 >= 3 points is always valid for np.stack).

Test

Extended tests/test_synthetic_chunk_data.py with test_chunk_data_chunks_do_not_exceed_max_tokens, asserting every produced chunk is <= max_tokens for 500- and 2000-token documents. It fails before the fix (chunks over budget / a doc returned unsplit) and passes after. The existing test_chunk_data_still_splits_long_document only checked len(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 check is clean.

No linked issue — found by reading the code.

Update (3e49790)

Also tightened the chunk count in the same path (raised in review): n_chunks was ceil(length / stride), but because consecutive chunks overlap by overlap, covering length needs ceil((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 within max_tokens. Added test_chunk_data_does_not_over_split (fails before at 3 chunks, passes after at 2). The advertised 500/1000/2000 examples are unchanged.

@winklemad
winklemad requested a review from danielhanchen as a code owner July 11, 2026 08:21

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
@winklemad
winklemad force-pushed the fix/synthetic-chunk-max-tokens branch from f2629df to 3e49790 Compare July 11, 2026 08:35
@danielhanchen

Copy link
Copy Markdown
Member

Pushed a small follow-up (comment-only) on top of your change:

  • Condensed the added comments in chunk_data and the two new tests so they stay tight while keeping the reasoning.
  • Adjusted the test_chunk_data_does_not_over_split note: the 3-chunk case it describes is the intermediate state (n_chunks+1 points with the loose count). On current main the two old bugs cancel and 673 already yields 2 chunks, so this test really guards the tightened count rather than failing before the fix. test_chunk_data_chunks_do_not_exceed_max_tokens is the true fail-before test.

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.

@danielhanchen

Copy link
Copy Markdown
Member

@codex review

1 similar comment
@danielhanchen

Copy link
Copy Markdown
Member

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Swish!

Reviewed commit: 45aed2d054

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

@danielhanchen

Copy link
Copy Markdown
Member

Thanks!

@danielhanchen
danielhanchen merged commit 935474c into unslothai:main Jul 12, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants