Merge branch 'main' into feature/swift-format

Author: FL33TW00D

Sources/Tokenizers/BertTokenizer.swift

@@ -45,10 +45,13 @@ public class BertTokenizer {
         self.fuseUnknownTokens = fuseUnknownTokens
     }
 
-    public required convenience init(tokenizerConfig: Config, tokenizerData: Config, addedTokens: [String: Int]) throws {
-        guard let vocab = tokenizerData.model?.vocab?.dictionary as? [String: Int] else {
-            throw TokenizerError.missingVocab
+    public required convenience init(tokenizerConfig: Config, tokenizerData: Config, addedTokens: [String : Int]) throws {
+        guard var vocab = tokenizerData.model?.vocab?.dictionary as? [String: Int] else { throw TokenizerError.missingVocab }
+        if let addedTokens = tokenizerData.added_tokens?.dictionary["value"] as? [[String: Any]],
+           let pairs = addedTokens.compactMap({ ($0["content"] as? String, $0["id"] as? Int) }) as? [(String, Int)] {
+            vocab.merge(pairs, uniquingKeysWith: {$1})
         }
+        vocab.merge(addedTokens, uniquingKeysWith: {$1})
         let merges = tokenizerData.model?.merges?.value as? [String]
         let tokenizeChineseChars = tokenizerConfig.handleChineseChars?.boolValue ?? true
         let eosToken = tokenizerConfig.eosToken?.stringValue

Tests/TokenizersTests/BertTokenizerTests.swift

@@ -7,7 +7,8 @@
 //
 
 @testable import Tokenizers
-import XCTest
+@testable import Hub
+
 
 class BertTokenizerTests: XCTestCase {
     override func setUp() {
@@ -175,4 +176,35 @@ class BertTokenizerTests: XCTestCase {
             XCTAssertEqual(decoded, String(expected))
         }
     }
+
+    func testBertTokenizerAddedTokensRecognized() async throws {
+        let base: URL = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first!.appending(component: "huggingface-tests")
+        let hubApi = HubApi(downloadBase: base)
+        let configuration = LanguageModelConfigurationFromHub(modelName: "google-bert/bert-base-uncased", hubApi: hubApi)
+        guard let tokenizerConfig = try await configuration.tokenizerConfig else { fatalError("missing tokenizer config") }
+        let tokenizerData = try await configuration.tokenizerData
+        let addedTokens = [
+            "[ROAD]": 60_001,
+            "[RIVER]": 60_002,
+            "[BUILDING]": 60_003,
+            "[PARK]": 60_004,
+            "[BUFFER]": 60_005,
+            "[INTERSECT]": 60_006,
+            "[UNION]": 60_007,
+        ]
+        let tokenizer = try BertTokenizer(tokenizerConfig: tokenizerConfig, tokenizerData: tokenizerData, addedTokens: addedTokens)
+        for (token, idx) in addedTokens {
+            XCTAssertEqual(tokenizer.convertTokenToId(token), idx)
+        }
+        for (token, idx) in addedTokens {
+            XCTAssertEqual(tokenizer.convertIdToToken(idx), token)
+        }
+
+        // Reading added_tokens from tokenizer.json
+        XCTAssertEqual(tokenizer.convertTokenToId("[PAD]"), 0)
+        XCTAssertEqual(tokenizer.convertTokenToId("[UNK]"), 100)
+        XCTAssertEqual(tokenizer.convertTokenToId("[CLS]"), 101)
+        XCTAssertEqual(tokenizer.convertTokenToId("[SEP]"), 102)
+        XCTAssertEqual(tokenizer.convertTokenToId("[MASK]"), 103)
+    }
 }