From c7736f9226f91f695713b1a58b663e464e14fb49 Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Mon, 12 Oct 2026 03:27:57 -0700 Subject: [PATCH] feat: load GPTNeoXTokenizer checkpoints (OLMoE) as byte-level BPE swift-transformers has no GPTNeoXTokenizer entry and throws before the server starts. When tokenizer.json is BPE, load it from a scratch copy whose tokenizer_config.json says GPT2Tokenizer. Verified end to end on OLMoE-1B-7B-0125-Instruct-4bit. Co-Authored-By: Claude Sonnet 5.5 --- Sources/SwiftLM/Server.swift | 58 ++++++++++++++- .../TokenizerClassRemapTests.swift | 73 +++++++++++++++++++ 2 files changed, 130 insertions(+), 1 deletion(-) create mode 100644 tests/SwiftLMTests/TokenizerClassRemapTests.swift diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index ba0e42f..253f33e 100644 --- a/Sources/SwiftLM/Server.swift +++ b/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 = ["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 diff --git a/tests/SwiftLMTests/TokenizerClassRemapTests.swift b/tests/SwiftLMTests/TokenizerClassRemapTests.swift new file mode 100644 index 0000000..4b0a0e1 --- /dev/null +++ b/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) + } +}