fix: Unigram tokenizer checkpoints crash with "BPETokenizer requires merges"

Sources/SwiftLM/Server.swift

@@ -49,9 +49,86 @@ private struct TransformersTokenizerLoader: TokenizerLoader, Sendable {
     init(modelId: String = "this model") { self.modelId = modelId }
 
     func load(from directory: URL) async throws -> any MLXLMCommon.Tokenizer {
-        let t = try await AutoTokenizer.from(modelFolder: directory)
+        let effectiveDirectory = Self.remappingUnigramTokenizerClassIfNeeded(in: directory, modelId: modelId)
+        let t = try await AutoTokenizer.from(modelFolder: effectiveDirectory)
         return TransformersTokenizerBridge(t, modelId: modelId)
     }
+
+    /// 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, e.g. several Japanese LLM projects that build their own vocab and
+    /// save via `AutoTokenizer`/`PreTrainedTokenizerFast` without a dedicated
+    /// subclass), there is no `merges` field and `BPETokenizer` hits a
+    /// `fatalError` instead of throwing (SwiftLM#155, reported by @kei-Optim).
+    ///
+    /// Workaround until the fix lands upstream in `huggingface/swift-transformers`
+    /// (`TokenizerModel.from` should consult `tokenizer.json`'s `model.type` when
+    /// `tokenizer_class` is generic/unknown): if `tokenizer.json` declares a
+    /// Unigram model and `tokenizer_config.json`'s `tokenizer_class` isn't already
+    /// one of the table's Unigram-mapped names, build 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"`, which resolves to `UnigramTokenizer`.
+    /// Its `init(tokenizerConfig:tokenizerData:addedTokens:)` doesn't depend on the
+    /// class name itself, so the rewrite is otherwise inert.
+    private static func remappingUnigramTokenizerClassIfNeeded(
+        in directory: URL, modelId: String
+    ) -> URL {
+        let tokenizerDataURL = directory.appendingPathComponent("tokenizer.json")
+        let tokenizerConfigURL = directory.appendingPathComponent("tokenizer_config.json")
+        guard
+            let tokenizerData = try? Data(contentsOf: tokenizerDataURL),
+            let tokenizerJSON = try? JSONSerialization.jsonObject(with: tokenizerData) as? [String: Any],
+            (tokenizerJSON["model"] as? [String: Any])?["type"] as? String == "Unigram",
+            let configData = try? Data(contentsOf: tokenizerConfigURL),
+            var configJSON = try? JSONSerialization.jsonObject(with: configData) as? [String: Any]
+        else {
+            return directory
+        }
+
+        let unigramMappedClasses: Set<String> = [
+            "XLMRobertaTokenizer", "Xlm-RobertaTokenizer", "T5Tokenizer",
+        ]
+        let currentClass = (configJSON["tokenizer_class"] as? String)?
+            .replacingOccurrences(of: "Fast", with: "")
+        if let currentClass, unigramMappedClasses.contains(currentClass) {
+            return directory
+        }
+
+        let originalClassDescription = (configJSON["tokenizer_class"] as? String) ?? "unset"
+        let scratchDirectory = FileManager.default.temporaryDirectory
+            .appendingPathComponent("swiftlm-unigram-tokenizer-\(UUID().uuidString)", isDirectory: true)
+        do {
+            try FileManager.default.createDirectory(
+                at: scratchDirectory, withIntermediateDirectories: true)
+            let contents = try FileManager.default.contentsOfDirectory(
+                at: directory, includingPropertiesForKeys: nil)
+            for fileURL in contents where fileURL.lastPathComponent != "tokenizer_config.json" {
+                try FileManager.default.createSymbolicLink(
+                    at: scratchDirectory.appendingPathComponent(fileURL.lastPathComponent),
+                    withDestinationURL: fileURL)
+            }
+            configJSON["tokenizer_class"] = "XLMRobertaTokenizer"
+            let rewritten = try JSONSerialization.data(
+                withJSONObject: configJSON, options: [.prettyPrinted])
+            try rewritten.write(to: scratchDirectory.appendingPathComponent("tokenizer_config.json"))
+        } catch {
+            // Scratch-directory setup is best-effort — fall through to the original
+            // directory rather than fail the whole load over a workaround failing.
+            return directory
+        }
+
+        print(
+            "[SwiftLM] \(modelId): tokenizer.json is Unigram but tokenizer_class is generic "
+                + "(\"\(originalClassDescription)\"); loading via a remapped tokenizer_class "
+                + "so it resolves to UnigramTokenizer instead of crashing in BPETokenizer.")
+        return scratchDirectory
+    }
 }
 
 /// A chat template that the Jinja parser rejected.

scripts/make-test-fixtures.py

@@ -77,6 +77,44 @@ def write_tokenizer(out):
     )
 
 
+def write_unigram_tokenizer(out):
+    """SwiftLM#155: a SentencePiece-style Unigram model.json, saved (as several real
+    Japanese-LLM checkpoints do) with the generic tokenizer_class "PreTrainedTokenizerFast"
+    rather than a class swift-transformers' knownTokenizers table maps to UnigramTokenizer
+    (XLMRobertaTokenizer, Xlm-RobertaTokenizer, T5Tokenizer). Before #155's fix, that
+    combination resolved to BPETokenizer, which fatalErrors on the missing "merges" field
+    instead of throwing. Same 288-token vocab as write_tokenizer, just a different model
+    algorithm and score instead of merges, so this drops in for build_dense's config/weights
+    unchanged.
+    """
+    vocab = [(t, 0.0 if t == SPECIALS[0] else -1.0) for t in SPECIALS]
+    for ch in sorted(pre_tokenizers.ByteLevel.alphabet()):
+        vocab.append((ch, -2.0 - len(vocab) * 0.01))
+    while len(vocab) < VOCAB:
+        vocab.append((f"<|unused{len(vocab)}|>", -10.0))
+
+    tok = Tokenizer(models.Unigram(vocab, 0, byte_fallback=False))
+    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+    tok.decoder = decoders.ByteLevel()
+    tok.post_processor = processors.ByteLevel(trim_offsets=False)
+    tok.add_special_tokens(SPECIALS)
+    tok.save(os.path.join(out, "tokenizer.json"))
+
+    json.dump(
+        {
+            # The bug trigger: generic tokenizer_class alongside a Unigram model.json.
+            "tokenizer_class": "PreTrainedTokenizerFast",
+            "bos_token": "<|endoftext|>",
+            "eos_token": "<|im_end|>",
+            "pad_token": "<pad>",
+            "unk_token": "<|endoftext|>",
+            "chat_template": CHAT_TEMPLATE,
+        },
+        open(os.path.join(out, "tokenizer_config.json"), "w"),
+        indent=2,
+    )
+
+
 def rand(rng, *shape):
     return (rng.standard_normal(shape) * 0.02).astype(np.float16)
 
@@ -352,12 +390,15 @@ def build_moe_nested(out):
     "kv-shared-absent": (build_gemma4_kv_shared, {"vestigial": False}),
     "kv-shared-present": (build_gemma4_kv_shared, {"vestigial": True}),
     "moe-nested": (build_moe_nested, {}),
+    "unigram-tokenizer": (build_dense, {}, write_unigram_tokenizer),
 }
 
 if __name__ == "__main__":
     os.makedirs(ROOT, exist_ok=True)
     total = 0
-    for name, (fn, kwargs) in FIXTURES.items():
+    for name, spec in FIXTURES.items():
+        fn, kwargs, *rest = spec
+        tokenizer_writer = rest[0] if rest else write_tokenizer
         out = os.path.join(ROOT, name)
         # Clear this fixture's own directory, and only its own. Writing over an existing
         # one leaves behind files the current builder no longer emits — flip stray-shard
@@ -370,7 +411,7 @@ def build_moe_nested(out):
             shutil.rmtree(out)
         os.makedirs(out)
         n = fn(out, **kwargs)
-        write_tokenizer(out)
+        tokenizer_writer(out)
         size = sum(os.path.getsize(os.path.join(out, f)) for f in os.listdir(out))
         total += size
         print(f"  {name:<20} {n:>3} tensors  {size/1024:>7.1f} KB")

tests/fixtures/unigram-tokenizer/config.json

@@ -0,0 +1,16 @@
+{
+  "model_type": "qwen2",
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "vocab_size": 288,
+  "hidden_size": 64,
+  "intermediate_size": 128,
+  "num_hidden_layers": 2,
+  "num_attention_heads": 4,
+  "num_key_value_heads": 2,
+  "max_position_embeddings": 512,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false
+}