Skip to content

Do not panic in BPE builder when a merge exceeds the longest vocab token#2153

Open
v-code01 wants to merge 1 commit into
huggingface:mainfrom
v-code01:bpe-merge-buffer-no-panic
Open

Do not panic in BPE builder when a merge exceeds the longest vocab token#2153
v-code01 wants to merge 1 commit into
huggingface:mainfrom
v-code01:bpe-merge-buffer-no-panic

Conversation

@v-code01

@v-code01 v-code01 commented Jul 5, 2026

Copy link
Copy Markdown

What

BpeBuilder::build builds each merge pair (a, b) into a fixed scratch buffer sized to max_len (the longest vocabulary token), then looks the concatenation up in the vocab:

let mut buffer: Vec<u8> = vec![0; max_len];
// ...
buffer[0..a.len()].copy_from_slice(a.as_bytes());
let b_len = b.len() - prefix_len;
let merge_len = a.len() + b_len;
buffer[a.len()..merge_len].copy_from_slice(&b.as_bytes()[prefix_len..]); // <-- write before the vocab check
let new_token = unsafe { from_utf8_unchecked(&buffer[..merge_len]) };
let new_id = vocab.get(new_token).ok_or_else(|| Error::MergeTokenOutOfVocabulary(...))?;

The write happens before the vocab lookup, and its length merge_len is only bounded by max_len when the merged token is itself a vocab entry. Two crafted inputs break that assumption:

  1. Over-long merge — pair two in-vocab tokens whose concatenation is not in the vocab and is longer than any vocab entry. Minimal repro: vocab {"a": 0, "b": 1} (so max_len == 1), merge ("a", "b"). Then merge_len == 2, and buffer[1..2] on a length-1 buffer panics with range end index 2 out of range for slice of length 1.
  2. Prefix underflowb.len() - prefix_len underflows (panic) when a right token is shorter than continuing_subword_prefix.

This is reachable from untrusted input: Deserialize for BPE routes through BpeBuilder::new().…vocab_and_merges(…).build(), so a crafted tokenizer.json panics 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 the MergeTokenOutOfVocabulary condition 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:

let b_tail = b.as_bytes().get(prefix_len..)
    .ok_or_else(|| Error::MergeTokenOutOfVocabulary(format!("{a}{b}")))?;
let merge_len = a.len() + b_tail.len();
if merge_len > buffer.len() {
    return Err(Error::MergeTokenOutOfVocabulary(format!("{a}{b}")).into());
}
buffer[0..a.len()].copy_from_slice(a.as_bytes());
buffer[a.len()..merge_len].copy_from_slice(b_tail);

Behavior is unchanged for valid inputs (same byte-level concatenation, same lookup); only the crafted cases now return Err instead of panicking.

Tests

  • merge_concat_longer_than_vocab_does_not_panic — the minimal over-long case now returns MergeTokenOutOfVocabulary.
  • merge_right_token_shorter_than_prefix_does_not_panic — the prefix-underflow case now errors gracefully.

Full tokenizers lib suite green (204 passed); cargo fmt/clippy clean on the changed file.

cc @ArthurZucker @McPatate

`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.
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.

1 participant