Use BinaryDistinctString for consistent hashing on Linux and macOS

Sources/Tokenizers/BPETokenizer.swift

@@ -101,8 +101,10 @@ private struct BPEMergeCandidate: Comparable {
 /// including GPT-2, RoBERTa, and other transformer models.
 class BPETokenizer: PreTrainedTokenizerModel, @unchecked Sendable {
     let bpeRanks: [BytePair: Int]
-    private let tokensToIds: [NSString: Int]
-    private let idsToTokens: [Int: NSString]
+    // NSString equality and hashing disagree for some canonically equivalent strings on Linux.
+    // Preserve the exact vocabulary keys on every platform.
+    private let tokensToIds: [BinaryDistinctString: Int]
+    private let idsToTokens: [Int: String]
 
     /// The total number of tokens in the vocabulary.
     var vocabCount: Int { tokensToIds.count }
@@ -174,28 +176,26 @@ class BPETokenizer: PreTrainedTokenizerModel, @unchecked Sendable {
         let addedTokens = addedTokens.reduce(into: [BinaryDistinctString: Config]()) { result, element in
             result[BinaryDistinctString(element.key)] = .init(element.value)
         }
-        tokensToIds = vocab.merging(addedTokens) { $1 }.reduce(into: [NSString: Int]()) { result, element in
-            result[element.key.nsString] = element.value.integer()
-        }
+        tokensToIds = vocab.merging(addedTokens) { $1 }.compactMapValues { $0.integer() }
 
-        idsToTokens = tokensToIds.reduce(into: [Int: NSString]()) { result, element in
-            result[element.value] = element.key
+        idsToTokens = tokensToIds.reduce(into: [Int: String]()) { result, element in
+            result[element.value] = element.key.string
         }
 
         // Populate tokens
         if let unknownToken = TokenizerModel.unknownToken(from: tokenizerConfig) {
             self.unknownToken = unknownToken
-            unknownTokenId = tokensToIds[unknownToken as NSString]
+            unknownTokenId = tokensToIds[BinaryDistinctString(unknownToken)]
         } else {
             unknownToken = nil
             unknownTokenId = nil
         }
 
         eosToken = addedTokenAsString(tokenizerConfig.eosToken)
-        eosTokenId = eosToken == nil ? nil : tokensToIds[eosToken! as NSString]
+        eosTokenId = eosToken == nil ? nil : tokensToIds[BinaryDistinctString(eosToken!)]
 
         bosToken = addedTokenAsString(tokenizerConfig.bosToken)
-        bosTokenId = bosToken == nil ? nil : tokensToIds[bosToken! as NSString]
+        bosTokenId = bosToken == nil ? nil : tokensToIds[BinaryDistinctString(bosToken!)]
 
         // `fuse_unk` is a model setting in tokenizer.json; tokenizer_config.json files don't carry it.
         fuseUnknownTokens = tokenizerData.model["fuseUnk"].boolean() ?? tokenizerConfig.fuseUnk.boolean(or: false)
@@ -208,15 +208,15 @@ class BPETokenizer: PreTrainedTokenizerModel, @unchecked Sendable {
     /// - Parameter token: The token string to convert
     /// - Returns: The numeric ID, or the unknown token ID if not found
     func convertTokenToId(_ token: String) -> Int? {
-        tokensToIds[token as NSString] ?? unknownTokenId
+        tokensToIds[BinaryDistinctString(token)] ?? unknownTokenId
     }
 
     /// Converts a numeric token ID back to its string representation.
     ///
     /// - Parameter id: The numeric token ID to convert
     /// - Returns: The token string, or nil if the ID is invalid
     func convertIdToToken(_ id: Int) -> String? {
-        idsToTokens[id] as String?
+        idsToTokens[id]
     }
 
     func byteEncode(text: String) -> [String] {
@@ -336,17 +336,17 @@ class BPETokenizer: PreTrainedTokenizerModel, @unchecked Sendable {
     /// - Returns: An array of BPE token strings
     func tokenize(text: String) -> [String] {
         // As in `tokenizers`, the whole word is looked up before any merge.
-        if ignoreMerges, tokensToIds[text as NSString] != nil {
+        if ignoreMerges, tokensToIds[BinaryDistinctString(text)] != nil {
             return [text]
         }
         var tokens: [String] = []
         let bpeTokens = bpe(token: text)
         for token in bpeTokens {
-            if tokensToIds[token as NSString] != nil {
+            if tokensToIds[BinaryDistinctString(token)] != nil {
                 tokens.append(token)
             } else {
                 let bytePieces = byteFallback ? hexaEncode(text: token) : []
-                if byteFallback && bytePieces.allSatisfy({ tokensToIds[$0 as NSString] != nil }) {
+                if byteFallback && bytePieces.allSatisfy({ tokensToIds[BinaryDistinctString($0)] != nil }) {
                     tokens.append(contentsOf: bytePieces)
                 } else if let unknownToken {
                     tokens.append(unknownToken)

Sources/Tokenizers/UnigramTokenizer.swift

@@ -39,8 +39,8 @@ class UnigramTokenizer: PreTrainedTokenizerModel, @unchecked Sendable {
     /// The minimum score found in the vocabulary (used for score calculations).
     let minScore: Float
 
-    /// Mapping from token strings to their numeric IDs.
-    let tokensToIds: [NSString: Int]
+    /// Mapping from exact token code units to their numeric IDs, without canonical equivalence.
+    let tokensToIds: [BinaryDistinctString: Int]
 
     /// The beginning-of-sequence token (hardcoded as space for Unigram).
     let bosToken: String? = " "
@@ -105,11 +105,11 @@ class UnigramTokenizer: PreTrainedTokenizerModel, @unchecked Sendable {
         byteFallback = tokenizerData.model["byteFallback"].boolean() ?? false
         unknownPiece = SentencePieceToken(token: vocab[unknownTokenId].token, score: minScore - 10)
 
-        tokensToIds = Dictionary(uniqueKeysWithValues: vocab.map { $0.token as NSString }.enumerated().map { ($1, $0) })
-        bosTokenId = tokensToIds[bosToken! as NSString] // May be nil
+        tokensToIds = Dictionary(uniqueKeysWithValues: vocab.enumerated().map { (BinaryDistinctString($0.element.token), $0.offset) })
+        bosTokenId = tokensToIds[BinaryDistinctString(bosToken!)] // May be nil
 
         eosToken = tokenizerConfig.eosToken.string()
-        eosTokenId = eosToken == nil ? nil : tokensToIds[eosToken! as NSString]
+        eosTokenId = eosToken == nil ? nil : tokensToIds[BinaryDistinctString(eosToken!)]
 
         trie = Trie()
         trie.append(contentsOf: vocab.map { $0.token.unicodeScalars })
@@ -120,7 +120,7 @@ class UnigramTokenizer: PreTrainedTokenizerModel, @unchecked Sendable {
     /// - Parameter token: The token string to convert
     /// - Returns: The numeric ID, or the unknown token ID if not found
     func convertTokenToId(_ token: String) -> Int? {
-        tokensToIds[token as NSString] ?? unknownTokenId
+        tokensToIds[BinaryDistinctString(token)] ?? unknownTokenId
     }
 
     /// Converts a numeric token ID back to its string representation.
@@ -153,7 +153,7 @@ class UnigramTokenizer: PreTrainedTokenizerModel, @unchecked Sendable {
             for tokenChars in trie.commonPrefixSearchIterator(suffix) {
                 let tokenLength = tokenChars.count
                 let token = String(String.UnicodeScalarView(tokenChars))
-                guard let tokenId = tokensToIds[token as NSString] else { fatalError("Token not in vocab: \(token)") }
+                guard let tokenId = tokensToIds[BinaryDistinctString(token)] else { fatalError("Token not in vocab: \(token)") }
                 let tokenScore = vocab[tokenId].score
                 lattice.insert(startOffset: beginPos, length: tokenLength, score: tokenScore, tokenId: tokenId)
                 if !hasSingleNode, tokenLength == mblen {
@@ -168,9 +168,9 @@ class UnigramTokenizer: PreTrainedTokenizerModel, @unchecked Sendable {
 
         guard byteFallback else { return lattice.tokens }
         return lattice.tokens.flatMap { piece -> [String] in
-            if tokensToIds[piece as NSString] != nil { return [piece] }
+            if tokensToIds[BinaryDistinctString(piece)] != nil { return [piece] }
             let bytePieces = piece.utf8.map { byteFallbackEncoderTable[Int($0)] }
-            guard bytePieces.allSatisfy({ tokensToIds[$0 as NSString] != nil })
+            guard bytePieces.allSatisfy({ tokensToIds[BinaryDistinctString($0)] != nil })
             else { return [piece] }
             return bytePieces
         }

Tests/TokenizersTests/BinaryDistinctStringTests.swift

@@ -111,7 +111,7 @@ struct BinaryDistinctStringTests {
         #expect(BinaryDistinctString(ns).value == Legacy.value(ns))
     }
 
-    @Test("A lone-surrogate token in a BPE vocabulary maps exactly as before")
+    @Test("A lone-surrogate vocabulary key stays distinct from the replacement character")
     func loneSurrogateInVocab() throws {
         // yyjson rejects a lone `\ud800` escape ("no low surrogate in string"), so a tokenizer.json
         // cannot carry one; a `Config` built in code can.
@@ -130,13 +130,9 @@ struct BinaryDistinctStringTests {
         #expect(tokenizer.vocabCount == 4)
         #expect(tokenizer.convertTokenToId("a") == 1)
         #expect(tokenizer.convertTokenToId("b") == 3)
-        #if canImport(Darwin)
-        // The lone-surrogate key is not the U+FFFD key: this lookup misses and falls back to <unk> (id 0).
+        // Vocabulary keys retain their original code units on every platform.
+        // A lookup for U+FFFD misses and falls back to <unk> (id 0).
         #expect(tokenizer.convertTokenToId("\u{FFFD}") == 0)
-        #else
-        // The NSString key is already repaired to U+FFFD, so the lookup finds the token.
-        #expect(tokenizer.convertTokenToId("\u{FFFD}") == 2)
-        #endif
         // Decoding the token yields the repaired string.
         #expect(tokenizer.convertIdToToken(2).map { Array($0.utf16) } == [0xFFFD])
     }

Tests/TokenizersTests/VocabularyIdentityTests.swift

@@ -0,0 +1,37 @@
+import Foundation
+import Testing
+
+@testable import Hub
+@testable import Tokenizers
+
+@Suite("Vocabulary identity")
+struct VocabularyIdentityTests {
+    @Test("Canonically equivalent tokens keep distinct IDs", arguments: ["BPE", "Unigram"])
+    func distinctTokenIds(modelType: String) throws {
+        // Parse JSON directly so neither Swift String dictionaries nor Foundation bridging
+        // can collapse canonically equivalent vocabulary keys while constructing the fixture.
+        let model =
+            modelType == "BPE"
+            ? #"{"type":"BPE","vocab":{"<unk>":0," ":1,";":2,"\u037e":3,"\u00e9":4,"e\u0301":5,"\u00f6":6},"merges":[],"ignore_merges":true}"#
+            : #"{"type":"Unigram","unk_id":0,"vocab":[["<unk>",0],[" ",-1],[";",-1],["\u037e",-1],["\u00e9",-1],["e\u0301",-1],["\u00f6",-1]]}"#
+        let data = try YYJSONParser.parseToConfig(Data("{\"model\":\(model)}".utf8))
+        let config = try YYJSONParser.parseToConfig(Data(#"{"unk_token":"<unk>","bos_token":";","eos_token":"\u037e"}"#.utf8))
+        let tokenizer: any PreTrainedTokenizerModel
+        if modelType == "BPE" {
+            tokenizer = try BPETokenizer(tokenizerConfig: config, tokenizerData: data, addedTokens: [:])
+        } else {
+            tokenizer = try UnigramTokenizer(tokenizerConfig: config, tokenizerData: data, addedTokens: [:])
+        }
+
+        let tokens = ["<unk>", " ", ";", "\u{037E}", "\u{00E9}", "e\u{0301}", "\u{00F6}"]
+        for (id, token) in tokens.enumerated() {
+            #expect(tokenizer.convertTokenToId(token) == id)
+            #expect(tokenizer.convertIdToToken(id).map { Array($0.utf8) } == Array(token.utf8))
+            #expect(tokenizer.tokenize(text: token).map { tokenizer.convertTokenToId($0) } == [id])
+        }
+        // A canonically equivalent spelling absent from the vocabulary must stay unknown.
+        #expect(tokenizer.convertTokenToId("o\u{0308}") == 0)
+        #expect(tokenizer.eosTokenId == 3)
+        #expect(tokenizer.bosTokenId == (modelType == "BPE" ? 2 : 1))
+    }
+}