Implement addedTokens in BertTokenizer (#193)

Author: shavit

Sources/Tokenizers/BertTokenizer.swift

@@ -46,9 +46,12 @@ public class BertTokenizer {
     }
 
     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
+        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

@@ -8,7 +8,7 @@
 
 import XCTest
 @testable import Tokenizers
-
+@testable import Hub
 
 
 class BertTokenizerTests: XCTestCase {
@@ -178,4 +178,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)
+    }
 }