Replace with enum for future extensibility

Author: pcuenca

Sources/Tokenizers/PreTokenizer.swift

@@ -8,26 +8,32 @@
 import Foundation
 import Hub
 
+public enum PreTokenizerOption: String {
+    case firstSection
+}
+
+public typealias PreTokenizerOptions = Set<PreTokenizerOption>
+
 public protocol PreTokenizer {
-    func preTokenize(text: String, firstSection: Bool) -> [String]
-    func preTokenize(texts: [String], firstSection: Bool) -> [String]
-    func callAsFunction(texts: [String], firstSection: Bool) -> [String]
-    func callAsFunction(text: String, firstSection: Bool) -> [String]
+    func preTokenize(text: String, options: PreTokenizerOptions) -> [String]
+    func preTokenize(texts: [String], options: PreTokenizerOptions) -> [String]
+    func callAsFunction(texts: [String], options: PreTokenizerOptions) -> [String]
+    func callAsFunction(text: String, options: PreTokenizerOptions) -> [String]
 
     init(config: Config)
 }
 
 extension PreTokenizer {
-    func preTokenize(texts: [String], firstSection: Bool = true) -> [String] {
-        texts.flatMap { preTokenize(text: $0, firstSection: firstSection) }
+    func preTokenize(texts: [String], options: PreTokenizerOptions = [.firstSection]) -> [String] {
+        texts.flatMap { preTokenize(text: $0, options: options) }
     }
 
-    func callAsFunction(texts: [String], firstSection: Bool = true) -> [String] {
-        return preTokenize(texts: texts, firstSection: firstSection)
+    func callAsFunction(texts: [String], options: PreTokenizerOptions = [.firstSection]) -> [String] {
+        return preTokenize(texts: texts, options: options)
     }
 
-    func callAsFunction(text: String, firstSection: Bool = true) -> [String] {
-        return preTokenize(text: text, firstSection: firstSection)
+    func callAsFunction(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
+        return preTokenize(text: text, options: options)
     }
 }
 
@@ -70,9 +76,9 @@ class PreTokenizerSequence: PreTokenizer {
         preTokenizers = configs.compactMap { PreTokenizerFactory.fromConfig(config: $0) }
     }
 
-    func preTokenize(text: String, firstSection: Bool = true) -> [String] {
+    func preTokenize(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
         preTokenizers.reduce([text]) { current, preTokenizer in
-            preTokenizer(texts: current, firstSection: firstSection)
+            preTokenizer(texts: current, options: options)
         }
     }
 }
@@ -84,7 +90,7 @@ class WhitespacePreTokenizer: PreTokenizer {
         re = #"\S+"#
     }
 
-    func preTokenize(text: String, firstSection: Bool = true) -> [String] {
+    func preTokenize(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
         return text.ranges(of: re).map { String(text[$0]) }
     }
 }
@@ -124,7 +130,7 @@ class MetaspacePreTokenizer: PreTokenizer {
 
     // https://github.com/huggingface/tokenizers/blob/accd0650b802f2180df40ef1def3bce32156688e/tokenizers/src/pre_tokenizers/metaspace.rs#L114
     // https://github.com/xenova/transformers.js/blob/b07336d8f7ff57453cc164cc68aead2a79cbd57e/src/tokenizers.js#L2153
-    func preTokenize(text: String, firstSection: Bool = true) -> [String] {
+    func preTokenize(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
         let normalized = text.replacingOccurrences(of: " ", with: stringReplacement)
 
         // We add a prefix space if:
@@ -140,7 +146,7 @@ class MetaspacePreTokenizer: PreTokenizer {
             if prependScheme == .always {
                 prepend = stringReplacement
             }
-            if prependScheme == .first && firstSection {
+            if prependScheme == .first && options.contains(.firstSection) {
                 prepend = stringReplacement
             }
         }
@@ -163,7 +169,7 @@ class ByteLevelPreTokenizer: PreTokenizer {
         useRegex = config.useRegex?.boolValue ?? true
     }
 
-    func preTokenize(text: String, firstSection: Bool = true) -> [String] {
+    func preTokenize(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
         // Split on whitespace and punctuation
         let tokens = useRegex ? text.ranges(of: RE).map({ String(text[$0]) }) : [text]
         return tokens.map { token in
@@ -185,7 +191,7 @@ class PunctuationPreTokenizer: PreTokenizer {
         re = "[^\(PUNCTUATION_REGEX)]+|[\(PUNCTUATION_REGEX)]+"
     }
 
-    func preTokenize(text: String, firstSection: Bool = true) -> [String] {
+    func preTokenize(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
         // Ref: https://github.com/xenova/transformers.js/blob/27920d84831e323275b38f0b5186644b7936e1a2/src/tokenizers.js#L1138
         return text.ranges(of: re).map { String(text[$0]) }
     }
@@ -199,7 +205,7 @@ class DigitsPreTokenizer: PreTokenizer {
         re = "[^\\d]+|\\d\(individualDigits ? "" : "+")"
     }
 
-    func preTokenize(text: String, firstSection: Bool = true) -> [String] {
+    func preTokenize(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
         return text.ranges(of: re).map { String(text[$0]) }
     }
 }
@@ -213,7 +219,7 @@ class SplitPreTokenizer: PreTokenizer {
         invert = config.invert?.boolValue ?? false
     }
 
-    func preTokenize(text: String, firstSection: Bool = true) -> [String] {
+    func preTokenize(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
         guard let pattern = pattern else { return [text] }
         return pattern.split(text, invert: invert)
     }

Sources/Tokenizers/Tokenizer.swift

@@ -185,9 +185,9 @@ public class PreTrainedTokenizer: Tokenizer {
         model = try TokenizerModel.from(tokenizerConfig: tokenizerConfig, tokenizerData: tokenizerData, addedTokens: addedTokens)
     }
 
-    func preTokenize(_ text: String, firstSection: Bool) -> [String] {
+    func preTokenize(_ text: String, options: PreTokenizerOptions) -> [String] {
         guard let preTokenizer = preTokenizer else { return [text] }
-        return preTokenizer(text: text, firstSection: firstSection)
+        return preTokenizer(text: text, options: options)
     }
 
     func normalize(_ text: String) -> String {
@@ -231,7 +231,7 @@ public class PreTrainedTokenizer: Tokenizer {
         }
         return sections.enumerated().map { section, x in
             if addedTokens.contains(x) { return [x] }
-            return preTokenize(normalize(x), firstSection: section == 0).flatMap { model($0) }
+            return preTokenize(normalize(x), options: section == 0 ? [.firstSection] : []).flatMap { model($0) }
         }.flatMap { $0 }
     }