Skip to content

Improve Unigram tokenization performance for long inputs#34

Open
bepsvpt wants to merge 1 commit into
huggingface:mainfrom
bepsvpt:fix/unigram-quadratic-tokenization
Open

Improve Unigram tokenization performance for long inputs#34
bepsvpt wants to merge 1 commit into
huggingface:mainfrom
bepsvpt:fix/unigram-quadratic-tokenization

Conversation

@bepsvpt

@bepsvpt bepsvpt commented Jul 16, 2026

Copy link
Copy Markdown

Summary

When Unigram tokenizes text, it creates the remaining string for every character:

chars.slice(begin_pos).join("")

For long input, this copies a lot of characters and the total work becomes O(n²).

This PR changes common_prefix_search to use the original character array with a start offset. It avoids creating a new suffix string every time.

The token scores and Viterbi logic are not changed.

Benchmark

Tested on the same machine with Node.js 26.5.0. Each case ran 3 times and the table shows the median time.

Input length Before After Speedup
2,000 16.336 ms 1.390 ms 11.8x
4,000 55.560 ms 3.089 ms 18.0x
8,000 202.157 ms 5.047 ms 40.1x
16,000 773.224 ms 10.182 ms 75.9x
32,000 3,063.974 ms 24.541 ms 124.9x

Before this change, the time is close to 4x when the input size doubles. After this change, it is around 2x.

Tests

  • CharTrie offset search unit test, including surrogate-pair tokens and out-of-range offsets
  • Unigram Viterbi test with overlapping multi-code-point tokens
  • A 50,000-character regression test that fails (5s timeout) if quadratic suffix materialization is reintroduced

AI assistance

AI tools were used for this PR:

  • Development: Codex (gpt-5.6-sol, xhigh)
  • Code review and tests: Claude Code (claude-fable-5, xhigh)

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