Make tokenizer merge rules type-safe

Author: DePasqualeOrg

Sources/Hub/Hub.swift

@@ -27,6 +27,18 @@ public enum TokenizerVocab: @unchecked Sendable {
     case unigram(NSArray)
 }
 
+/// Merge rules extracted from tokenizer.json for fast BPE initialization.
+///
+/// - Note: `@unchecked Sendable` is safe because the underlying data is immutable after extraction from JSON.
+public struct TokenizerMerges: @unchecked Sendable {
+    /// The raw merge rules as extracted from JSON.
+    public let rules: [Any]
+
+    public init(_ rules: [Any]) {
+        self.rules = rules
+    }
+}
+
 public extension Hub {
     /// Errors that can occur during Hub client operations.
     ///
@@ -142,7 +154,7 @@ public actor LanguageModelConfigurationFromHub {
     private var _tokenizerConfig: Config?
     private var _tokenizerData: Config?
     private var _tokenizerVocab: TokenizerVocab?
-    private var _tokenizerMerges: [Any]?
+    private var _tokenizerMerges: TokenizerMerges?
 
     /// Initializes configuration loading from a remote Hub repository.
     ///
@@ -257,8 +269,8 @@ public actor LanguageModelConfigurationFromHub {
         }
     }
 
-    /// Raw merges array extracted directly from JSON for fast BPE tokenizer initialization.
-    public var tokenizerMerges: [Any]? {
+    /// Merge rules extracted directly from JSON for fast BPE tokenizer initialization.
+    public var tokenizerMerges: TokenizerMerges? {
         get async throws {
             try await ensureLoaded()
             return _tokenizerMerges
@@ -281,7 +293,7 @@ public actor LanguageModelConfigurationFromHub {
         var tokenizerConfig: Config?
         var tokenizerData: Config
         var tokenizerVocab: TokenizerVocab?
-        var tokenizerMerges: [Any]?
+        var tokenizerMerges: TokenizerMerges?
     }
 
     /// Resolves tokenizerConfig with fallback logic.
@@ -373,7 +385,7 @@ public actor LanguageModelConfigurationFromHub {
 
             // Extract vocab/merges for fast tokenizer initialization (BPE and Unigram)
             var tokenizerVocab: TokenizerVocab? = nil
-            var tokenizerMerges: [Any]? = nil
+            var tokenizerMerges: TokenizerMerges? = nil
 
             if let modelDict = parsed["model"] as? NSDictionary {
                 let model = NSMutableDictionary(dictionary: modelDict)
@@ -382,7 +394,9 @@ public actor LanguageModelConfigurationFromHub {
                 // Only extract and strip for BPE and Unigram models
                 if modelType == "BPE", let vocab = model["vocab"] as? NSDictionary {
                     tokenizerVocab = .bpe(vocab)
-                    tokenizerMerges = model["merges"] as? [Any]
+                    if let merges = model["merges"] as? [Any] {
+                        tokenizerMerges = TokenizerMerges(merges)
+                    }
 
                     // Only strip if opted in (for backward compatibility)
                     if stripVocabForPerformance {

Sources/Tokenizers/Tokenizer.swift

@@ -228,7 +228,7 @@ enum TokenizerModel {
         tokenizerData: Config,
         addedTokens: [String: Int],
         tokenizerVocab: TokenizerVocab?,
-        tokenizerMerges: [Any]?,
+        tokenizerMerges: TokenizerMerges?,
         strict: Bool = true
     ) throws -> TokenizingModel {
         guard let tokenizerClassName = tokenizerConfig.tokenizerClass.string() else {
@@ -250,7 +250,7 @@ enum TokenizerModel {
         // Note: includes empty subclasses (creates BPETokenizer instance)
         if tokenizerClass is BPETokenizer.Type,
             case .bpe(let rawVocab) = tokenizerVocab,
-            let rawMerges = tokenizerMerges
+            let rawMerges = tokenizerMerges?.rules
         {
             return try BPETokenizer(
                 tokenizerConfig: tokenizerConfig,
@@ -285,7 +285,7 @@ enum TokenizerModel {
         tokenizerData: Config,
         addedTokens: [String: Int],
         tokenizerVocab: TokenizerVocab?,
-        tokenizerMerges: [Any]?,
+        tokenizerMerges: TokenizerMerges?,
         strict: Bool = true
     ) async throws -> TokenizingModel {
         guard let tokenizerClassName = tokenizerConfig.tokenizerClass.string() else {
@@ -307,7 +307,7 @@ enum TokenizerModel {
         // Note: includes empty subclasses (creates BPETokenizer instance)
         if tokenizerClass is BPETokenizer.Type,
             case .bpe(let rawVocab) = tokenizerVocab,
-            let rawMerges = tokenizerMerges
+            let rawMerges = tokenizerMerges?.rules
         {
             return await BPETokenizer.createAsync(
                 tokenizerConfig: tokenizerConfig,
@@ -628,7 +628,7 @@ public class PreTrainedTokenizer: @unchecked Sendable, Tokenizer {
         tokenizerConfig: Config,
         tokenizerData: Config,
         tokenizerVocab: TokenizerVocab? = nil,
-        tokenizerMerges: [Any]? = nil,
+        tokenizerMerges: TokenizerMerges? = nil,
         strict: Bool = true
     ) throws {
         var addedTokens: [String: Int] = [:]
@@ -751,7 +751,7 @@ public class PreTrainedTokenizer: @unchecked Sendable, Tokenizer {
         tokenizerConfig: Config,
         tokenizerData: Config,
         tokenizerVocab: TokenizerVocab?,
-        tokenizerMerges: [Any]?,
+        tokenizerMerges: TokenizerMerges?,
         strict: Bool = true
     ) async throws -> PreTrainedTokenizer {
         // Parse addedTokens (small data, used for model init)
@@ -1202,7 +1202,7 @@ public extension AutoTokenizer {
         tokenizerConfig: Config,
         tokenizerData: Config,
         tokenizerVocab: TokenizerVocab?,
-        tokenizerMerges: [Any]?,
+        tokenizerMerges: TokenizerMerges?,
         strict: Bool = true
     ) throws -> Tokenizer {
         let tokenizerClass = tokenizerClass(for: tokenizerConfig)
@@ -1220,7 +1220,7 @@ public extension AutoTokenizer {
         tokenizerConfig: Config,
         tokenizerData: Config,
         tokenizerVocab: TokenizerVocab?,
-        tokenizerMerges: [Any]?,
+        tokenizerMerges: TokenizerMerges?,
         strict: Bool = true
     ) async throws -> Tokenizer {
         let selectedClass = tokenizerClass(for: tokenizerConfig)
@@ -1294,7 +1294,7 @@ class LlamaPreTrainedTokenizer: PreTrainedTokenizer, @unchecked Sendable {
         tokenizerConfig: Config,
         tokenizerData: Config,
         tokenizerVocab: TokenizerVocab? = nil,
-        tokenizerMerges: [Any]? = nil,
+        tokenizerMerges: TokenizerMerges? = nil,
         strict: Bool = true
     ) throws {
         isLegacy = tokenizerConfig.legacy.boolean(or: true)
@@ -1328,7 +1328,7 @@ class LlamaPreTrainedTokenizer: PreTrainedTokenizer, @unchecked Sendable {
         tokenizerConfig: Config,
         tokenizerData: Config,
         tokenizerVocab: TokenizerVocab?,
-        tokenizerMerges: [Any]?,
+        tokenizerMerges: TokenizerMerges?,
         strict: Bool = true
     ) async throws -> PreTrainedTokenizer {
         let isLegacy = tokenizerConfig.legacy.boolean(or: true)

Tests/HubTests/ConfigTests.swift

@@ -540,7 +540,7 @@ struct ConfigTests {
         // Merges should be extracted
         let merges = try await loader.tokenizerMerges
         #expect(merges != nil)
-        #expect(merges?.count == 2)
+        #expect(merges?.rules.count == 2)
 
         // tokenizerData.model.vocab should be empty (stripped)
         let tokenizerData = try await loader.tokenizerData