Export bosToken and eosToken (#70)

Author: pcuenca

Sources/Tokenizers/BPETokenizer.swift

@@ -36,9 +36,13 @@ class BPETokenizer: PreTrainedTokenizerModel {
     private let tokensToIds: [String: Int]
     private let idsToTokens: [Int: String]
 
+    public let bosToken: String?
+    public let bosTokenId: Int?
+    public let eosToken: String?
+    public let eosTokenId: Int?
     public let unknownToken: String?
     public let unknownTokenId: Int?
-    
+
     required init(tokenizerConfig: Config, tokenizerData: Config, addedTokens: [String : Int]) throws {
         guard let merges = tokenizerData.model?.merges?.value as? [String] else { fatalError("BPETokenizer requires merges") }
         guard let vocab = tokenizerData.model?.vocab?.dictionary as? [String: Int] else {
@@ -55,14 +59,20 @@ class BPETokenizer: PreTrainedTokenizerModel {
         self.tokensToIds = vocab.merging(addedTokens) { $1 }
         self.idsToTokens = Utils.invert(self.tokensToIds)
 
-        // Populate unknown token
+        // Populate tokens
         if let unknownToken = TokenizerModel.unknownToken(from: tokenizerConfig) {
             self.unknownToken = unknownToken
             self.unknownTokenId = self.tokensToIds[unknownToken]
         } else {
             self.unknownToken = nil
             self.unknownTokenId = nil
         }
+
+        eosToken = tokenizerConfig.eosToken?.stringValue
+        eosTokenId = eosToken == nil ? nil : tokensToIds[eosToken!]
+
+        bosToken = tokenizerConfig.bosToken?.stringValue
+        bosTokenId = bosToken == nil ? nil : tokensToIds[bosToken!]
     }
 
     func convertTokenToId(_ token: String) -> Int? {

Sources/Tokenizers/BertTokenizer.swift

@@ -17,14 +17,26 @@ class BertTokenizer {
 
     private let vocab: [String: Int]
     private let ids_to_tokens: [Int: String]
-    
+
+    var bosToken: String?
+    var bosTokenId: Int?
+    var eosToken: String?
+    var eosTokenId: Int?
+
     init(vocab: [String: Int],
          merges: [String]?,
-         tokenizeChineseChars: Bool = true) {
+         tokenizeChineseChars: Bool = true,
+         bosToken: String? = nil,
+         eosToken: String? = nil
+    ) {
         self.vocab = vocab
         self.ids_to_tokens = Utils.invert(vocab)
         self.wordpieceTokenizer = WordpieceTokenizer(vocab: self.vocab)
         self.tokenizeChineseChars = tokenizeChineseChars
+        self.bosToken = bosToken
+        self.bosTokenId = bosToken == nil ? nil : vocab[bosToken!]
+        self.eosToken = eosToken
+        self.eosTokenId = eosToken == nil ? nil : vocab[eosToken!]
     }
 
     required convenience init(tokenizerConfig: Config, tokenizerData: Config, addedTokens: [String : Int]) throws {
@@ -33,7 +45,9 @@ class BertTokenizer {
         }
         let merges = tokenizerData.model?.merges?.value as? [String]
         let tokenizeChineseChars = tokenizerConfig.handleChineseChars?.boolValue ?? true
-        self.init(vocab: vocab, merges: merges, tokenizeChineseChars: tokenizeChineseChars)
+        let eosToken = tokenizerConfig.eosToken?.stringValue
+        let bosToken = tokenizerConfig.bosToken?.stringValue
+        self.init(vocab: vocab, merges: merges, tokenizeChineseChars: tokenizeChineseChars, bosToken: bosToken, eosToken: eosToken)
     }
 
 

Sources/Tokenizers/Tokenizer.swift

@@ -30,6 +30,10 @@ public protocol TokenizingModel {
     func convertIdToToken(_ id: Int) -> String?
     func convertIdsToTokens(_ ids: [Int]) -> [String?]
 
+    var bosToken: String? { get }
+    var bosTokenId: Int? { get }
+    var eosToken: String? { get }
+    var eosTokenId: Int? { get }
     var unknownToken: String? { get }
     var unknownTokenId: Int? { get }
 }
@@ -103,6 +107,10 @@ public protocol Tokenizer {
     func convertIdToToken(_ id: Int) -> String?
     func convertIdsToTokens(_ ids: [Int]) -> [String?]
 
+    var bosToken: String? { get }
+    var bosTokenId: Int? { get }
+    var eosToken: String? { get }
+    var eosTokenId: Int? { get }
     var unknownToken: String? { get }
     var unknownTokenId: Int? { get }
 }
@@ -124,6 +132,10 @@ public extension Tokenizer {
 public class PreTrainedTokenizer: Tokenizer {
     let model: TokenizingModel
 
+    public var bosToken: String? { model.bosToken }
+    public var bosTokenId: Int? { model.bosTokenId }
+    public var eosToken: String? { model.eosToken }
+    public var eosTokenId: Int? { model.eosTokenId }
     public var unknownToken: String? { model.unknownToken }
     public var unknownTokenId: Int? { model.unknownTokenId }
 

Sources/Tokenizers/UnigramTokenizer.swift

@@ -24,7 +24,7 @@ class UnigramTokenizer: PreTrainedTokenizerModel {
     let minScore: Float
     let tokensToIds: [String: Int]
 
-    let bosToken = " "
+    let bosToken: String? = " "
     let bosTokenId: Int?
     let eosToken: String?
     let eosTokenId: Int?
@@ -51,7 +51,7 @@ class UnigramTokenizer: PreTrainedTokenizerModel {
         self.unknownPiece = SentencePieceToken(token: vocab[unknownTokenId].token, score: minScore - 10)
 
         tokensToIds = Dictionary(uniqueKeysWithValues: vocab.map { $0.token }.enumerated().map { ($1, $0) })
-        bosTokenId = tokensToIds[bosToken]      // May be nil
+        bosTokenId = tokensToIds[bosToken!]      // May be nil
 
         eosToken = tokenizerConfig.eosToken?.stringValue
         eosTokenId = eosToken == nil ? nil : tokensToIds[eosToken!]