Update PreTokenizers so Metaspace can conditionally act

Author: pcuenca

Sources/Tokenizers/PreTokenizer.swift

@@ -9,27 +9,26 @@ import Foundation
 import Hub
 
 public protocol PreTokenizer {
-    func preTokenize(text: String) -> [String]
-    func preTokenize(texts: [String]) -> [String]
-    func callAsFunction(texts: [String]) -> [String]
-    func callAsFunction(text: String) -> [String]
+    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]
 
     init(config: Config)
 }
 
 extension PreTokenizer {
-    func preTokenize(texts: [String]) -> [String] {
-        texts.flatMap { preTokenize(text: $0) }
+    func preTokenize(texts: [String], firstSection: Bool = true) -> [String] {
+        texts.flatMap { preTokenize(text: $0, firstSection: firstSection) }
     }
 
-    func callAsFunction(texts: [String]) -> [String] {
-        return preTokenize(texts: texts)
+    func callAsFunction(texts: [String], firstSection: Bool = true) -> [String] {
+        return preTokenize(texts: texts, firstSection: firstSection)
     }
 
-    func callAsFunction(text: String) -> [String] {
-        return preTokenize(text: text)
+    func callAsFunction(text: String, firstSection: Bool = true) -> [String] {
+        return preTokenize(text: text, firstSection: firstSection)
     }
-    
 }
 
 enum PreTokenizerType: String {
@@ -71,9 +70,9 @@ class PreTokenizerSequence: PreTokenizer {
         preTokenizers = configs.compactMap { PreTokenizerFactory.fromConfig(config: $0) }
     }
 
-    func preTokenize(text: String) -> [String] {
+    func preTokenize(text: String, firstSection: Bool = true) -> [String] {
         preTokenizers.reduce([text]) { current, preTokenizer in
-            preTokenizer(texts: current)
+            preTokenizer(texts: current, firstSection: firstSection)
         }
     }
 }
@@ -85,7 +84,7 @@ class WhitespacePreTokenizer: PreTokenizer {
         re = #"\S+"#
     }
 
-    func preTokenize(text: String) -> [String] {
+    func preTokenize(text: String, firstSection: Bool = true) -> [String] {
         return text.ranges(of: re).map { String(text[$0]) }
     }
 }
@@ -125,7 +124,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) -> [String] {
+    func preTokenize(text: String, firstSection: Bool = true) -> [String] {
         let normalized = text.replacingOccurrences(of: " ", with: stringReplacement)
 
         // We add a prefix space if:
@@ -141,7 +140,7 @@ class MetaspacePreTokenizer: PreTokenizer {
             if prependScheme == .always {
                 prepend = stringReplacement
             }
-            if prependScheme == .first /* && first_section */ {
+            if prependScheme == .first && firstSection {
                 prepend = stringReplacement
             }
         }
@@ -164,7 +163,7 @@ class ByteLevelPreTokenizer: PreTokenizer {
         useRegex = config.useRegex?.boolValue ?? true
     }
 
-    func preTokenize(text: String) -> [String] {
+    func preTokenize(text: String, firstSection: Bool = true) -> [String] {
         // Split on whitespace and punctuation
         let tokens = useRegex ? text.ranges(of: RE).map({ String(text[$0]) }) : [text]
         return tokens.map { token in
@@ -186,7 +185,7 @@ class PunctuationPreTokenizer: PreTokenizer {
         re = "[^\(PUNCTUATION_REGEX)]+|[\(PUNCTUATION_REGEX)]+"
     }
 
-    func preTokenize(text: String) -> [String] {
+    func preTokenize(text: String, firstSection: Bool = true) -> [String] {
         // Ref: https://github.com/xenova/transformers.js/blob/27920d84831e323275b38f0b5186644b7936e1a2/src/tokenizers.js#L1138
         return text.ranges(of: re).map { String(text[$0]) }
     }
@@ -200,7 +199,7 @@ class DigitsPreTokenizer: PreTokenizer {
         re = "[^\\d]+|\\d\(individualDigits ? "" : "+")"
     }
 
-    func preTokenize(text: String) -> [String] {
+    func preTokenize(text: String, firstSection: Bool = true) -> [String] {
         return text.ranges(of: re).map { String(text[$0]) }
     }
 }
@@ -214,7 +213,7 @@ class SplitPreTokenizer: PreTokenizer {
         invert = config.invert?.boolValue ?? false
     }
 
-    func preTokenize(text: String) -> [String] {
+    func preTokenize(text: String, firstSection: Bool = true) -> [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) -> [String] {
+    func preTokenize(_ text: String, firstSection: Bool) -> [String] {
         guard let preTokenizer = preTokenizer else { return [text] }
-        return preTokenizer(text: text)
+        return preTokenizer(text: text, firstSection: firstSection)
     }
 
     func normalize(_ text: String) -> String {
@@ -229,9 +229,9 @@ public class PreTrainedTokenizer: Tokenizer {
         } else {
             sections = [text]
         }
-        return sections.map { x in
+        return sections.enumerated().map { section, x in
             if addedTokens.contains(x) { return [x] }
-            return preTokenize(normalize(x)).flatMap { model($0) }
+            return preTokenize(normalize(x), firstSection: section == 0).flatMap { model($0) }
         }.flatMap { $0 }
     }