From db61e7516ca5057c96068404ca6b245d8c781faa Mon Sep 17 00:00:00 2001 From: Pedro Cuenca Date: Thu, 29 Oct 2026 18:03:13 +0100 Subject: [PATCH] Use BinaryDistinctString for consistent hashing in Linux and macOS --- Sources/Tokenizers/BPETokenizer.swift | 30 +++++++-------- Sources/Tokenizers/UnigramTokenizer.swift | 18 ++++----- .../BinaryDistinctStringTests.swift | 10 ++--- .../VocabularyIdentityTests.swift | 37 +++++++++++++++++++ 4 files changed, 64 insertions(+), 31 deletions(-) create mode 100644 Tests/TokenizersTests/VocabularyIdentityTests.swift diff --git a/Sources/Tokenizers/BPETokenizer.swift b/Sources/Tokenizers/BPETokenizer.swift index c2ff1107..b125ad8b 100644 --- a/Sources/Tokenizers/BPETokenizer.swift +++ b/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,7 +208,7 @@ 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. @@ -216,7 +216,7 @@ class BPETokenizer: PreTrainedTokenizerModel, @unchecked Sendable { /// - 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) diff --git a/Sources/Tokenizers/UnigramTokenizer.swift b/Sources/Tokenizers/UnigramTokenizer.swift index 1bf24c7c..1c13e975 100644 --- a/Sources/Tokenizers/UnigramTokenizer.swift +++ b/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 } diff --git a/Tests/TokenizersTests/BinaryDistinctStringTests.swift b/Tests/TokenizersTests/BinaryDistinctStringTests.swift index fcbe4250..3dbe417e 100644 --- a/Tests/TokenizersTests/BinaryDistinctStringTests.swift +++ b/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 (id 0). + // Vocabulary keys retain their original code units on every platform. + // A lookup for U+FFFD misses and falls back to (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]) } diff --git a/Tests/TokenizersTests/VocabularyIdentityTests.swift b/Tests/TokenizersTests/VocabularyIdentityTests.swift new file mode 100644 index 00000000..46d98e32 --- /dev/null +++ b/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":{"":0," ":1,";":2,"\u037e":3,"\u00e9":4,"e\u0301":5,"\u00f6":6},"merges":[],"ignore_merges":true}"# + : #"{"type":"Unigram","unk_id":0,"vocab":[["",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":"","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 = ["", " ", ";", "\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)) + } +}