Add BART tokenizer

Sources/Tokenizers/BartTokenizer.swift

@@ -0,0 +1,47 @@
+import Hub
+
+class BartTokenizer {
+    public var bosToken: String?
+    public var bosTokenId: Int?
+    public var eosToken: String?
+    public var eosTokenId: Int?
+    public var unknownToken: String?
+    public var unknownTokenId: Int?
+    public var fuseUnknownTokens: Bool
+
+    public let padToken: String
+    public let sepToken: String
+    public let clsToken: String
+    public let maskToken: String
+
+    private let vocab: [String: Int]
+    private let ids_to_tokens: [Int: String]
+    private let bpe: BPETokenizer
+
+    required public init(tokenizerConfig: Config, tokenizerData: Config, addedTokens: [String : Int]) throws {
+        guard let vocab = tokenizerData.model?.vocab?.dictionary as? [String: Int] else { throw TokenizerError.missingVocab }
+        self.bosToken = tokenizerConfig.bosToken?.stringValue ?? "<s>"
+        self.bosTokenId = bosToken == nil ? nil : vocab[bosToken!]
+        self.eosToken = tokenizerConfig.eosToken?.stringValue ?? "</s>"
+        self.eosTokenId = eosToken == nil ? nil : vocab[eosToken!]
+        self.unknownToken = tokenizerConfig.unkToken?.stringValue ?? "<unk>"
+        self.unknownTokenId = unknownToken == nil ? nil : vocab[unknownToken!]
+        self.fuseUnknownTokens = tokenizerConfig.fuseUnk?.boolValue ?? false
+        self.padToken = tokenizerConfig.padToken?.stringValue ?? "<pad>"
+        self.sepToken = tokenizerConfig.sepToken?.stringValue ?? "</s>"
+        self.clsToken = tokenizerConfig.clsToken?.stringValue ?? "<s>"
+        self.maskToken = tokenizerConfig.maskToken?.stringValue ?? "<mask>"
+        self.vocab = vocab
+        self.ids_to_tokens = Utils.invert(vocab)
+        self.bpe = try BPETokenizer(tokenizerConfig: tokenizerConfig, tokenizerData: tokenizerData, addedTokens: addedTokens)
+    }
+
+    func callAsFunction(_ text: String) -> [String] { bpe.tokenize(text: text) }
+    func unTokenize(tokens: [Int]) -> [String] { tokens.compactMap({ ids_to_tokens[$0] }) }
+}
+
+extension BartTokenizer: PreTrainedTokenizerModel {
+    func convertTokenToId(_ token: String) -> Int? { vocab[token] ?? unknownTokenId }
+    func convertIdToToken(_ id: Int) -> String? { ids_to_tokens[id] }
+    func tokenize(text: String) -> [String] { bpe.tokenize(text: text) }
+}

Sources/Tokenizers/Tokenizer.swift

@@ -77,6 +77,7 @@ public protocol PreTrainedTokenizerModel: TokenizingModel {
 
 struct TokenizerModel {
     static let knownTokenizers: [String : PreTrainedTokenizerModel.Type] = [
+        "BartTokenizer"      : BartTokenizer.self,
         "BertTokenizer"      : BertTokenizer.self,
         "DistilbertTokenizer": BertTokenizer.self,
         "DistilBertTokenizer": BertTokenizer.self,

Tests/TokenizersTests/Resources/bart_large_mnli.json

@@ -0,0 +1 @@
+{"text": "Justin Timberlake and Jessica Biel, welcome to parenthood.\n        The celebrity couple announced the arrival of their son, Silas Randall Timberlake, in statements to People.\n        \"Silas was the middle name of Timberlake's maternal grandfather Bill Bomar, who died in 2012, while Randall is the musician's own middle name, as well as his father's first,\" People reports.\n        The couple announced the pregnancy in January, with an Instagram post. It is the first baby for both.", "token_ids": [0, 29466, 19440, 11985, 8, 7103, 163, 5255, 6, 2814, 7, 181, 40532, 4, 50118, 1437, 1437, 1437, 1437, 1437, 1437, 1437, 20, 6794, 891, 585, 5, 5237, 9, 49, 979, 6, 8897, 281, 16222, 19440, 11985, 6, 11, 1997, 7, 1806, 4, 50118, 1437, 1437, 1437, 1437, 1437, 1437, 1437, 22, 23719, 281, 21, 5, 1692, 766, 9, 19440, 11985, 18, 22835, 10642, 1585, 23707, 271, 6, 54, 962, 11, 1125, 6, 150, 16222, 16, 5, 9613, 18, 308, 1692, 766, 6, 25, 157, 25, 39, 1150, 18, 78, 60, 1806, 690, 4, 50118, 1437, 1437, 1437, 1437, 1437, 1437, 1437, 20, 891, 585, 5, 6690, 11, 644, 6, 19, 41, 1838, 618, 4, 85, 16, 5, 78, 1928, 13, 258, 4, 2], "bpe_tokens": ["<s>", "Justin", "\u0120Timber", "lake", "\u0120and", "\u0120Jessica", "\u0120B", "iel", ",", "\u0120welcome", "\u0120to", "\u0120p", "arenthood", ".", "\u010a", "\u0120", "\u0120", "\u0120", "\u0120", "\u0120", "\u0120", "\u0120", "\u0120The", "\u0120celebrity", "\u0120couple", "\u0120announced", "\u0120the", "\u0120arrival", "\u0120of", "\u0120their", "\u0120son", ",", "\u0120Sil", "as", "\u0120Randall", "\u0120Timber", "lake", ",", "\u0120in", "\u0120statements", "\u0120to", "\u0120People", ".", "\u010a", "\u0120", "\u0120", "\u0120", "\u0120", "\u0120", "\u0120", "\u0120", "\u0120\"", "Sil", "as", "\u0120was", "\u0120the", "\u0120middle", "\u0120name", "\u0120of", "\u0120Timber", "lake", "'s", "\u0120maternal", "\u0120grandfather", "\u0120Bill", "\u0120Bom", "ar", ",", "\u0120who", "\u0120died", "\u0120in", "\u01202012", ",", "\u0120while", "\u0120Randall", "\u0120is", "\u0120the", "\u0120musician", "'s", "\u0120own", "\u0120middle", "\u0120name", ",", "\u0120as", "\u0120well", "\u0120as", "\u0120his", "\u0120father", "'s", "\u0120first", ",\"", "\u0120People", "\u0120reports", ".", "\u010a", "\u0120", "\u0120", "\u0120", "\u0120", "\u0120", "\u0120", "\u0120", "\u0120The", "\u0120couple", "\u0120announced", "\u0120the", "\u0120pregnancy", "\u0120in", "\u0120January", ",", "\u0120with", "\u0120an", "\u0120Instagram", "\u0120post", ".", "\u0120It", "\u0120is", "\u0120the", "\u0120first", "\u0120baby", "\u0120for", "\u0120both", ".", "</s>"], "decoded_text": "<s>Justin Timberlake and Jessica Biel, welcome to parenthood.\n        The celebrity couple announced the arrival of their son, Silas Randall Timberlake, in statements to People.\n        \"Silas was the middle name of Timberlake's maternal grandfather Bill Bomar, who died in 2012, while Randall is the musician's own middle name, as well as his father's first,\" People reports.\n        The couple announced the pregnancy in January, with an Instagram post. It is the first baby for both.</s>"}

Tests/TokenizersTests/TokenizerTests.swift

@@ -72,6 +72,23 @@ class BertUncasedTokenizerTests: TokenizerTests {
     override class var unknownTokenId: Int? { 100 }
 }
 
+class BartLargeMnliTokenizerTests: TokenizerTests {
+    override class var hubModelName: String? { "facebook/bart-large-mnli" }
+    override class var encodedSamplesFilename: String? { "bart_large_mnli" }
+    override class var unknownTokenId: Int? { 3 }
+
+    override func testTokenize() async {
+        let testCases: [(String, [String])] = [
+            ("Justin Timberlake and Jessica Biel, welcome to parenthood.", ["Justin", "\u{0120}Timber", "lake", "\u{0120}and", "\u{0120}Jessica", "\u{0120}B", "iel", ",", "\u{0120}welcome", "\u{0120}to", "\u{0120}p", "arenthood", "."]),
+            ("Silas Randall Timberlake, in statements to People. \"Silas was the middle name of Timberlake's maternal grandfather Bill Bomar, who died in 2012, while Randall is the musician's own middle name, as well as his father's first,\" People reports.", ["Sil", "as", "\u{0120}Randall", "\u{0120}Timber", "lake", ",", "\u{0120}in", "\u{0120}statements", "\u{0120}to", "\u{0120}People", ".", "\u{0120}\"", "Sil", "as", "\u{0120}was", "\u{0120}the", "\u{0120}middle", "\u{0120}name", "\u{0120}of", "\u{0120}Timber", "lake", "'s", "\u{0120}maternal", "\u{0120}grandfather", "\u{0120}Bill", "\u{0120}Bom", "ar", ",", "\u{0120}who", "\u{0120}died", "\u{0120}in", "\u{0120}2012", ",", "\u{0120}while", "\u{0120}Randall", "\u{0120}is", "\u{0120}the", "\u{0120}musician", "'s", "\u{0120}own", "\u{0120}middle", "\u{0120}name", ",", "\u{0120}as", "\u{0120}well", "\u{0120}as", "\u{0120}his", "\u{0120}father", "'s", "\u{0120}first", ",\"", "\u{0120}People", "\u{0120}reports", "."]),
+        ]
+        let tokenizer = await BartLargeMnliTokenizerTests._tester!.tokenizer!
+        for (text, expectTokens) in testCases {
+            XCTAssertEqual(tokenizer.tokenize(text: text), expectTokens)
+        }
+    }
+}
+
 class GemmaTokenizerTests: TokenizerTests {
     override class var hubModelName: String? { "pcuenq/gemma-tokenizer" }
     override class var encodedSamplesFilename: String? { "gemma_encoded" }