Unigram tokenizer checkpoints crash with "BPETokenizer requires merges" (tokenizer_class="PreTrainedTokenizerFast")

[kei-Optim]

Summary

Loading a checkpoint whose tokenizer.json uses a Unigram model (SentencePiece-style) but whose tokenizer_config.json sets tokenizer_class to the generic "PreTrainedTokenizerFast" crashes at load time with:

Tokenizers/BPETokenizer.swift:95: Fatal error: BPETokenizer requires merges

Root cause

swift-transformers' TokenizerModel.from(...) (Tokenizer.swift) 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 actual tokenizer.json is Unigram, there is no merges field, and BPETokenizer hits a fatalError instead of throwing.

This isn't specific to one checkpoint — it hits any model whose tokenizer was trained from scratch (custom vocab/algorithm) and saved via AutoTokenizer/PreTrainedTokenizerFast without a dedicated subclass, which is common for Japanese LLM projects that build their own SentencePiece vocab.

Confirmed affected checkpoints

Repo	tokenizer_class	tokenizer.json model.type
mlx-community/Tanuki-8B-dpo-v1.0-8bit	PreTrainedTokenizerFast	Unigram
llm-jp/llm-jp-13b-v2.0	PreTrainedTokenizerFast	Unigram
llm-jp/llm-jp-3-13b	PreTrainedTokenizerFast	Unigram

Models that inherit an existing architecture (Llama/GPT-NeoX/Qwen-based Japanese LLMs like rinna, ELYZA, Swallow, Sarashina, PLaMo) were not affected — they save a specific tokenizer_class (e.g. LlamaTokenizer) that swift-transformers already maps correctly.

Repro

./SwiftLM --model mlx-community/Tanuki-8B-dpo-v1.0-8bit --port 5413

Fails during tokenizer load, after the model weights finish downloading.

Proposed fix (implemented locally, not yet committed)

In Sources/SwiftLM/Server.swift, TransformersTokenizerLoader.load(from:) now inspects tokenizer.json's model.type before calling AutoTokenizer.from(modelFolder:). If model.type == "Unigram" and the checkpoint's tokenizer_class isn't one of swift-transformers' known 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 directory instead. This causes swift-transformers to select UnigramTokenizer, whose init(tokenizerConfig:tokenizerData:addedTokens:) doesn't depend on the class name itself.

Verified working end-to-end with mlx-community/Tanuki-8B-dpo-v1.0-8bit (Japanese generation succeeds, ~18.5 tok/s on M4 Pro 48GB).

Happy to open a PR with this change if useful. The proper long-term fix likely belongs upstream in huggingface/swift-transformers (have TokenizerModel.from consult tokenizer.json's model.type when tokenizer_class is a generic/unknown name), but this local workaround unblocks affected checkpoints in the meantime.

[solderzzc]

Thank you for this — it's an unusually thorough report.

The root-cause analysis matches what we'd expect: TokenizerModel.from(...) resolving purely from the tokenizer_class string, with "PreTrainedTokenizerFast" → "PreTrainedTokenizer" → BPETokenizer, and then a fatalError rather than a throw once merges is missing.

The most useful part is the table separating the affected checkpoints from the ones that inherit a concrete tokenizer_class — that makes it clear this is a class of checkpoint (tokenizers trained from scratch and saved via AutoTokenizer/PreTrainedTokenizerFast) rather than a single bad model, which isn't at all obvious from the crash alone. That framing is what makes the issue actionable.

A PR would be very welcome, yes please. The symlinked scratch-directory approach sounds like the right shape for a workaround — leaving the HF cache untouched matters, and rewriting only tokenizer_config.json keeps the blast radius small.

And agreed on the long-term fix belonging upstream in huggingface/swift-transformers — having TokenizerModel.from consult tokenizer.json's model.type when tokenizer_class is generic or unknown. If you open one there as well, link it here and we can drop the local workaround once it lands.

Thanks again, and for including the repro and the throughput number.