Do not panic in BPE builder when a merge exceeds the longest vocab token#2153
Open
v-code01 wants to merge 1 commit into
Open
Do not panic in BPE builder when a merge exceeds the longest vocab token#2153v-code01 wants to merge 1 commit into
v-code01 wants to merge 1 commit into
Conversation
`BpeBuilder::build` concatenates each merge pair `(a, b)` into a fixed
scratch buffer sized to `max_len`, the length of the longest vocabulary
token, then looks the result up in the vocab. The write
buffer[a.len()..a.len() + b.len() - prefix_len]
happens *before* that lookup, and its length is unbounded: it is only
correct when the merged token is itself a vocabulary entry (hence no
longer than `max_len`). A crafted merge list — e.g. from an untrusted
`tokenizer.json`, since deserialization routes through this same builder
— can pair two in-vocab tokens whose concatenation is longer than any
vocab entry (minimal case: vocab `{"a","b"}`, merge `("a","b")`), which
indexes past the buffer and panics ("range end index N out of range").
The same line also underflows `b.len() - prefix_len` when a right token
is shorter than `continuing_subword_prefix`.
Both cases mean the merged token cannot exist in the vocabulary, so
surface them as the intended `MergeTokenOutOfVocabulary` error instead of
panicking: strip the prefix with a checked slice and bail before the
buffer write when the concatenation is longer than the scratch buffer.
Adds regression tests for the over-long merge and the short-right-token
cases.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
BpeBuilder::buildbuilds each merge pair(a, b)into a fixed scratch buffer sized tomax_len(the longest vocabulary token), then looks the concatenation up in the vocab:The write happens before the vocab lookup, and its length
merge_lenis only bounded bymax_lenwhen the merged token is itself a vocab entry. Two crafted inputs break that assumption:{"a": 0, "b": 1}(somax_len == 1), merge("a", "b"). Thenmerge_len == 2, andbuffer[1..2]on a length-1 buffer panics withrange end index 2 out of range for slice of length 1.b.len() - prefix_lenunderflows (panic) when a right token is shorter thancontinuing_subword_prefix.This is reachable from untrusted input:
Deserialize for BPEroutes throughBpeBuilder::new().…vocab_and_merges(…).build(), so a craftedtokenizer.jsonpanics the loader instead of returning an error.Fix
Both cases mean the merged token cannot exist in the vocabulary (everything in the vocab is
≤ max_len), which is exactly theMergeTokenOutOfVocabularycondition the code already handles gracefully three lines later. So strip the prefix with a checked slice and bail before the buffer write when the concatenation can't fit:Behavior is unchanged for valid inputs (same byte-level concatenation, same lookup); only the crafted cases now return
Errinstead of panicking.Tests
merge_concat_longer_than_vocab_does_not_panic— the minimal over-long case now returnsMergeTokenOutOfVocabulary.merge_right_token_shorter_than_prefix_does_not_panic— the prefix-underflow case now errors gracefully.Full
tokenizerslib suite green (204 passed);cargo fmt/clippyclean on the changed file.cc @ArthurZucker @McPatate