fix: Unigram tokenizer checkpoints crash with "BPETokenizer requires merges" - #158
Merged
Conversation
…merges" Fixes #155. Reported and fully root-caused by @kei-Optim, whose local fix this implements. swift-transformers' TokenizerModel.from(...) picks a concrete tokenizer implementation purely from the tokenizer_class name string (its knownTokenizers table) — it never inspects tokenizer.json's own model.type. After stripping a "Fast" suffix, "PreTrainedTokenizerFast" becomes "PreTrainedTokenizer", which the table maps explicitly to BPETokenizer. When the checkpoint's actual tokenizer.json is Unigram (SentencePiece-style — common for models trained from scratch with a custom vocab, several Japanese LLM projects among them), there is no merges field and BPETokenizer hits a fatalError instead of throwing. Confirmed affected: mlx-community/Tanuki-8B-dpo-v1.0-8bit, llm-jp/llm-jp-13b-v2.0, llm-jp/llm-jp-3-13b — all Unigram tokenizer.json saved with the generic tokenizer_class "PreTrainedTokenizerFast". Fix: TransformersTokenizerLoader.load(from:) now inspects tokenizer.json's model.type before calling AutoTokenizer.from. If it's Unigram and tokenizer_class isn't already one of the table's Unigram-mapped names (XLMRobertaTokenizer, Xlm-RobertaTokenizer, T5Tokenizer), it builds a scratch copy of the checkpoint directory (original files referenced via symlink, so the HF cache is untouched) with tokenizer_config.json's tokenizer_class rewritten to "XLMRobertaTokenizer", then loads from that instead. UnigramTokenizer's init(tokenizerConfig:tokenizerData:addedTokens:) doesn't depend on the class name itself, so the rewrite is otherwise inert. Adds a synthetic unigram-tokenizer fixture (scripts/make-test-fixtures.py, tests/test-fixtures.sh) reproducing the exact shape — a real Unigram tokenizer.json built via the tokenizers library, saved with the generic tokenizer_class — so this has regression coverage without needing any of the affected (multi-GB) real checkpoints. Red-green verified: reverting the Server.swift change locally reproduces the exact reported crash ("Tokenizers/BPETokenizer.swift:95: Fatal error: BPETokenizer requires merges") against the new fixture; restoring it, the fixture loads, prefills, and returns a well-formed completion.
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
1 participant
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.
Summary
Fixes #155. Reported and fully root-caused by @kei-Optim — this implements their local fix.
swift-transformers'TokenizerModel.from(...)picks a concrete tokenizer implementation purely from thetokenizer_classname string — it never inspectstokenizer.json's ownmodel.type. After stripping a"Fast"suffix,"PreTrainedTokenizerFast"becomes"PreTrainedTokenizer", which maps explicitly toBPETokenizer. When the checkpoint's actualtokenizer.jsonis Unigram (SentencePiece-style, common for from-scratch vocabs — several Japanese LLM projects among them), there's nomergesfield andBPETokenizerhits afatalErrorinstead of throwing:Confirmed affected:
mlx-community/Tanuki-8B-dpo-v1.0-8bit,llm-jp/llm-jp-13b-v2.0,llm-jp/llm-jp-3-13b.Fix
TransformersTokenizerLoader.load(from:)now inspectstokenizer.json'smodel.typebefore callingAutoTokenizer.from. If it's Unigram andtokenizer_classisn't already one of swift-transformers' Unigram-mapped names (XLMRobertaTokenizer,Xlm-RobertaTokenizer,T5Tokenizer), it builds a scratch copy of the checkpoint directory — original files symlinked, so the HF cache is untouched — withtokenizer_config.json'stokenizer_classrewritten to"XLMRobertaTokenizer", then loads from that instead.UnigramTokenizer's init doesn't depend on the class name, so the rewrite is otherwise inert.The proper long-term fix belongs upstream in
huggingface/swift-transformers(TokenizerModel.fromshould consulttokenizer.json'smodel.typewhentokenizer_classis generic/unknown) — this is a local workaround in the meantime, same spirit as the reporter's own framing.Test coverage
Adds a synthetic
unigram-tokenizerfixture reproducing the exact shape — a real Unigramtokenizer.jsonbuilt via thetokenizerslibrary, saved with the generictokenizer_class— so this has regression coverage without needing any of the affected multi-GB real checkpoints.Red-green verified locally: reverting the
Server.swiftchange reproduces the exact reported crash against the new fixture; restoring it, the fixture loads, prefills, and returns a well-formed completion.Test plan
bash tests/test-fixtures.sh— 7/7 passed locally, including the new fixture.