feat: load GPTNeoXTokenizer checkpoints (OLMoE)

Sources/SwiftLM/Server.swift

@@ -50,7 +50,8 @@ private struct TransformersTokenizerLoader: TokenizerLoader, Sendable {
     init(modelId: String = "this model") { self.modelId = modelId }
 
     func load(from directory: URL) async throws -> any MLXLMCommon.Tokenizer {
-        let effectiveDirectory = Self.remappingUnigramTokenizerClassIfNeeded(in: directory, modelId: modelId)
+        let unigramRemapped = Self.remappingUnigramTokenizerClassIfNeeded(in: directory, modelId: modelId)
+        let effectiveDirectory = remappingBPECompatibleTokenizerClassIfNeeded(in: unigramRemapped, modelId: modelId)
         let t = try await AutoTokenizer.from(modelFolder: effectiveDirectory)
         return TransformersTokenizerBridge(t, modelId: modelId)
     }
@@ -132,6 +133,61 @@ private struct TransformersTokenizerLoader: TokenizerLoader, Sendable {
     }
 }
 
+/// `tokenizer_class` values that swift-transformers does not list in its `knownTokenizers`
+/// table but whose `tokenizer.json` is the same byte-level BPE as GPT-2, so `BPETokenizer`
+/// loads them correctly. Loading such a checkpoint otherwise fails with
+/// `unsupportedTokenizer("GPTNeoXTokenizer")` before the server starts (OLMoE, GPT-NeoX
+/// and Pythia derivatives).
+let bpeCompatibleTokenizerClasses: Set<String> = ["GPTNeoXTokenizer"]
+
+/// Returns `directory` unchanged unless its `tokenizer_class` is one of
+/// `bpeCompatibleTokenizerClasses` *and* `tokenizer.json` declares a BPE model. In that
+/// case it returns a scratch copy, original files symlinked so the HF cache is untouched,
+/// whose `tokenizer_config.json` says `GPT2Tokenizer`, which swift-transformers maps to
+/// `BPETokenizer`. Checking `model.type` keeps a mislabelled Unigram or WordPiece
+/// checkpoint from being forced into the BPE implementation.
+///
+/// Same mechanism as `TransformersTokenizerLoader.remappingUnigramTokenizerClassIfNeeded`;
+/// delete both once swift-transformers selects the implementation from `model.type`.
+func remappingBPECompatibleTokenizerClassIfNeeded(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 == "BPE",
+        let configData = try? Data(contentsOf: tokenizerConfigURL),
+        var configJSON = try? JSONSerialization.jsonObject(with: configData) as? [String: Any],
+        let declared = configJSON["tokenizer_class"] as? String,
+        bpeCompatibleTokenizerClasses.contains(declared.replacingOccurrences(of: "Fast", with: ""))
+    else {
+        return directory
+    }
+
+    let scratchDirectory = FileManager.default.temporaryDirectory
+        .appendingPathComponent("swiftlm-bpe-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"] = "GPT2Tokenizer"
+        let rewritten = try JSONSerialization.data(withJSONObject: configJSON, options: [.prettyPrinted])
+        try rewritten.write(to: scratchDirectory.appendingPathComponent("tokenizer_config.json"))
+    } catch {
+        // Best-effort, as for the Unigram remap: fall back to the original directory and
+        // let the loader report its own error.
+        return directory
+    }
+
+    print("[SwiftLM] \(modelId): tokenizer_class \"\(declared)\" is not in swift-transformers' table; "
+        + "tokenizer.json is byte-level BPE, so loading it as GPT2Tokenizer.")
+    return scratchDirectory
+}
+
 /// A chat template that the Jinja parser rejected.
 ///
 /// Worth its own type because the underlying error says only what the parser choked on

tests/SwiftLMTests/TokenizerClassRemapTests.swift

@@ -0,0 +1,73 @@
+import XCTest
+import Foundation
+@testable import SwiftLM
+
+// MARK: - tokenizer_class values swift-transformers does not know
+//
+// OLMoE ships `tokenizer_class: GPTNeoXTokenizer`. swift-transformers' table has no such
+// entry and throws before the server starts. The checkpoint's tokenizer.json is plain
+// byte-level BPE, so SwiftLM loads it from a scratch copy that says GPT2Tokenizer.
+final class TokenizerClassRemapTests: XCTestCase {
+
+    private var dirs: [URL] = []
+
+    override func tearDown() {
+        for d in dirs { try? FileManager.default.removeItem(at: d) }
+        dirs = []
+        super.tearDown()
+    }
+
+    private func makeModelDir(tokenizerClass: String?, modelType: String) throws -> URL {
+        let dir = FileManager.default.temporaryDirectory
+            .appendingPathComponent("tok-remap-\(UUID().uuidString)", isDirectory: true)
+        try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
+        dirs.append(dir)
+        var config: [String: Any] = ["model_max_length": 1024]
+        if let tokenizerClass { config["tokenizer_class"] = tokenizerClass }
+        try JSONSerialization.data(withJSONObject: config).write(to: dir.appendingPathComponent("tokenizer_config.json"))
+        try JSONSerialization.data(withJSONObject: ["model": ["type": modelType]])
+            .write(to: dir.appendingPathComponent("tokenizer.json"))
+        try Data("weights".utf8).write(to: dir.appendingPathComponent("model.safetensors"))
+        return dir
+    }
+
+    private func tokenizerClass(in dir: URL) throws -> String? {
+        let data = try Data(contentsOf: dir.appendingPathComponent("tokenizer_config.json"))
+        return (try JSONSerialization.jsonObject(with: data) as? [String: Any])?["tokenizer_class"] as? String
+    }
+
+    func testGPTNeoXBPEIsRemappedToGPT2() throws {
+        let dir = try makeModelDir(tokenizerClass: "GPTNeoXTokenizer", modelType: "BPE")
+        let out = remappingBPECompatibleTokenizerClassIfNeeded(in: dir, modelId: "test")
+        dirs.append(out)
+        XCTAssertNotEqual(out, dir)
+        XCTAssertEqual(try tokenizerClass(in: out), "GPT2Tokenizer")
+        XCTAssertEqual(try tokenizerClass(in: dir), "GPTNeoXTokenizer", "the original checkpoint is not modified")
+        XCTAssertEqual(try Data(contentsOf: out.appendingPathComponent("model.safetensors")), Data("weights".utf8),
+                       "other files are reachable through the scratch copy")
+    }
+
+    func testFastSuffixIsStripped() throws {
+        let dir = try makeModelDir(tokenizerClass: "GPTNeoXTokenizerFast", modelType: "BPE")
+        let out = remappingBPECompatibleTokenizerClassIfNeeded(in: dir, modelId: "test")
+        dirs.append(out)
+        XCTAssertEqual(try tokenizerClass(in: out), "GPT2Tokenizer")
+    }
+
+    func testKnownClassIsLeftAlone() throws {
+        let dir = try makeModelDir(tokenizerClass: "Qwen2Tokenizer", modelType: "BPE")
+        XCTAssertEqual(remappingBPECompatibleTokenizerClassIfNeeded(in: dir, modelId: "test"), dir)
+    }
+
+    /// A GPTNeoX-labelled checkpoint whose tokenizer.json is not BPE must not be forced
+    /// into the BPE implementation: the loader should report its own error instead.
+    func testNonBPEModelIsLeftAlone() throws {
+        let dir = try makeModelDir(tokenizerClass: "GPTNeoXTokenizer", modelType: "Unigram")
+        XCTAssertEqual(remappingBPECompatibleTokenizerClassIfNeeded(in: dir, modelId: "test"), dir)
+    }
+
+    func testMissingTokenizerClassIsLeftAlone() throws {
+        let dir = try makeModelDir(tokenizerClass: nil, modelType: "BPE")
+        XCTAssertEqual(remappingBPECompatibleTokenizerClassIfNeeded(in: dir, modelId: "test"), dir)
+    }
+}