Skip to content

fix(Tokenizers): avoid quadratic regex matching for added tokens without whitespace padding - #386

Merged
pcuenca merged 4 commits into
huggingface:mainfrom
shubhransh-gupta:fix/added-tokens-regex-perf
Oct 29, 2026
Merged

pcuenca merged 4 commits into
huggingface:mainfrom
shubhransh-gupta:fix/added-tokens-regex-perf

Conversation

@shubhransh-gupta

Copy link
Copy Markdown
Contributor

Fixes #383.

Summary

In PreTrainedTokenizer, addedTokensRegex was constructed by wrapping every added token in a capturing group (\(token)) even when neither lstrip nor rstrip was requested.

For tokenizers with large added token vocabularies (e.g. Google Gemma-3 with 6,415 added tokens), this caused an $O(N^2)$ quadratic slowdown in Foundation's ICU regular expression engine due to capturing group table allocations, slowing down single newline \n encoding from ~0.02 ms to ~29 ms (~1.1s for a 40-line prompt on device).

Changes

  1. Non-capturing groups by default: When $0.prefix and $0.suffix are both false, compile the token using a non-capturing group (?:\(token)) instead of (\(token)).
  2. Preserve whitespace stripping: Added tokens with lstrip or rstrip still retain their capturing group so split(by: NSRegularExpression) continues to isolate and strip surrounding whitespace as expected.
  3. Empty check guard: Avoid compiling an empty regex pattern when unwrappedAddedTokens is empty.
  4. Regression test: Added addedTokensRegexPerformanceAndCorrectness in Tests/TokenizersTests/TokenizerTests.swift.

Verification

  • Executed swift test across all targets (50/50 tests passing).
  • Validated code style with swift-format lint --strict --recursive Sources Tests (clean, 0 warnings).

@pcuenca pcuenca left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Hi @shubhransh-gupta, thanks a lot for your contribution, and sorry for the delay 🙏

Your optimization seems plausible, but the tests you propose are not checking for correctness or performance. I'm happy to merge without them and rely on the existing test corpus, as long as you can provide some indication about performance improvements you measured. Would you happen to have something handy?

#expect(tokenizerOpt != nil)
let tokenizer = tokenizerOpt!

// Verify that added tokens and newlines are matched and encoded properly without quadratic explosion

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This doesn't really verify the quadratic explosion claim.

@pcuenca

pcuenca commented Oct 29, 2026

Copy link
Copy Markdown
Member

Actually, #383 provides sufficient evidence, I'll add @xocialize as a co-author to this PR.

Comment thread Tests/TokenizersTests/TokenizerTests.swift Outdated
@pcuenca
pcuenca merged commit deaec6c into huggingface:main Oct 29, 2026
4 checks 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.

Tokenizer.encode is quadratic in the number of added tokens (~29 ms per newline with Gemma-3)

2 participants