Added Tokens (#93)

pcuenca · web-flow · commit 0a606f50af21 · 2024-04-28T19:31:11.000+02:00
* Split by regexp with capture groups

The other split helpers we have don't work for capture groups.
We had to resort to raw `NSRegularExpression`s

* Build added tokens split regexp, shortcut before pre-tokenization

* Update PreTokenizers so Metaspace can conditionally act

* Create LlamaPreTrainedTokenizer subclass

We need some custom behaviour that's not in the config :(

* Rename test

* Replace with enum for future extensibility
diff --git a/Sources/Tokenizers/PreTokenizer.swift b/Sources/Tokenizers/PreTokenizer.swift
@@ -8,28 +8,33 @@
 import Foundation
 import Hub
 
+public enum PreTokenizerOption: String {
+    case firstSection
+}
+
+public typealias PreTokenizerOptions = Set<PreTokenizerOption>
+
 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, 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]) -> [String] {
-        texts.flatMap { preTokenize(text: $0) }
+    func preTokenize(texts: [String], options: PreTokenizerOptions = [.firstSection]) -> [String] {
+        texts.flatMap { preTokenize(text: $0, options: options) }
     }
 
-    func callAsFunction(texts: [String]) -> [String] {
-        return preTokenize(texts: texts)
+    func callAsFunction(texts: [String], options: PreTokenizerOptions = [.firstSection]) -> [String] {
+        return preTokenize(texts: texts, options: options)
     }
     
-    func callAsFunction(text: String) -> [String] {
-        return preTokenize(text: text)
+    func callAsFunction(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
+        return preTokenize(text: text, options: options)
     }
-    
 }
 
 enum PreTokenizerType: String {
@@ -71,9 +76,9 @@ class PreTokenizerSequence: PreTokenizer {
         preTokenizers = configs.compactMap { PreTokenizerFactory.fromConfig(config: $0) }
     }
     
-    func preTokenize(text: String) -> [String] {
+    func preTokenize(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
         preTokenizers.reduce([text]) { current, preTokenizer in
-            preTokenizer(texts: current)
+            preTokenizer(texts: current, options: options)
         }
     }
 }
@@ -85,7 +90,7 @@ class WhitespacePreTokenizer: PreTokenizer {
         re = #"\S+"#
     }
 
-    func preTokenize(text: String) -> [String] {
+    func preTokenize(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
         return text.ranges(of: re).map { String(text[$0]) }
     }
 }
@@ -125,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) -> [String] {
+    func preTokenize(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
         let normalized = text.replacingOccurrences(of: " ", with: stringReplacement)
         
         // We add a prefix space if:
@@ -141,7 +146,7 @@ class MetaspacePreTokenizer: PreTokenizer {
             if prependScheme == .always {
                 prepend = stringReplacement
             }
-            if prependScheme == .first /* && first_section */ {
+            if prependScheme == .first && options.contains(.firstSection) {
                 prepend = stringReplacement
             }
         }
@@ -164,7 +169,7 @@ class ByteLevelPreTokenizer: PreTokenizer {
         useRegex = config.useRegex?.boolValue ?? true
     }
     
-    func preTokenize(text: String) -> [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
@@ -186,7 +191,7 @@ class PunctuationPreTokenizer: PreTokenizer {
         re = "[^\(PUNCTUATION_REGEX)]+|[\(PUNCTUATION_REGEX)]+"
     }
 
-    func preTokenize(text: String) -> [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]) }
     }
@@ -200,7 +205,7 @@ class DigitsPreTokenizer: PreTokenizer {
         re = "[^\\d]+|\\d\(individualDigits ? "" : "+")"
     }
 
-    func preTokenize(text: String) -> [String] {
+    func preTokenize(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
         return text.ranges(of: re).map { String(text[$0]) }
     }
 }
@@ -214,7 +219,7 @@ class SplitPreTokenizer: PreTokenizer {
         invert = config.invert?.boolValue ?? false
     }
 
-    func preTokenize(text: String) -> [String] {
+    func preTokenize(text: String, options: PreTokenizerOptions = [.firstSection]) -> [String] {
         guard let pattern = pattern else { return [text] }
         return pattern.split(text, invert: invert)
     }
@@ -248,7 +253,7 @@ extension StringSplitPattern {
     }
 }
 
-extension String {
+public extension String {
     func ranges(of string: String, options: CompareOptions = .regularExpression) -> [Range<Index>] {
         var result: [Range<Index>] = []
         var start = startIndex
@@ -277,6 +282,42 @@ extension String {
         return result
     }
 
+    /// This version supports capture groups, wheres the one above doesn't
+    func split(by captureRegex: NSRegularExpression) -> [String] {
+        // Find the matching capture groups
+        let selfRange = NSRange(startIndex..<endIndex, in: self)
+        let matches = captureRegex.matches(in: self, options: [], range: selfRange)
+
+        if matches.first == nil { return [self] }
+
+        var result: [String] = []
+        var start = startIndex
+        for match in matches {
+            // Append prefix before matched separator
+            let prefixEnd = index(startIndex, offsetBy: match.range.lowerBound)
+            if start < prefixEnd {
+                result.append(String(self[start..<prefixEnd]))
+            }
+            start = index(startIndex, offsetBy: match.range.upperBound)
+
+            // Append separator, supporting capture groups
+            for r in (0..<match.numberOfRanges).reversed() {
+                let matchRange = match.range(at: r)
+                if let sepRange = Range(matchRange, in:self) {
+                    result.append(String(self[sepRange]))
+                    break
+                }
+            }
+        }
+
+        // Append remaining suffix
+        let beginningOfEnd = index(startIndex, offsetBy: matches.last!.range.upperBound)
+        if beginningOfEnd < endIndex {
+            result.append(String(self[beginningOfEnd...]))
+        }
+
+        return result
+    }
 }
 
 public enum SplitDelimiterBehavior {
diff --git a/Sources/Tokenizers/Tokenizer.swift b/Sources/Tokenizers/Tokenizer.swift
@@ -141,6 +141,7 @@ public class PreTrainedTokenizer: Tokenizer {
 
     private let addedTokens: Set<String>
     private let specialTokens: [String: Int]
+    private let addedTokensRegex: NSRegularExpression?
 
     private let preTokenizer: PreTokenizer?
     private let normalizer: Normalizer?
@@ -161,6 +162,16 @@ public class PreTrainedTokenizer: Tokenizer {
                 specialTokens[content] = id
             }
         }
+
+        let addedTokensRegexString = (tokenizerData.addedTokens?.arrayValue ?? []).compactMap { addedToken in
+               guard let content = addedToken.content?.stringValue else { return nil }
+               let prefix = (addedToken.lstrip?.boolValue ?? false ? #"\s*"# : "")
+               let suffix = (addedToken.rstrip?.boolValue ?? false ? #"\s*"# : "")
+               let token = NSRegularExpression.escapedPattern(for: content)
+               return "\(prefix)(\(token))\(suffix)"
+        }.joined(separator: "|")
+        addedTokensRegex = try? NSRegularExpression(pattern: addedTokensRegexString, options: [])
+
         // TODO: specialTokens are stored but never used
         self.specialTokens = specialTokens
         self.addedTokens = Set(addedTokens.keys)
@@ -174,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, options: PreTokenizerOptions) -> [String] {
         guard let preTokenizer = preTokenizer else { return [text] }
-        return preTokenizer(text: text)
+        return preTokenizer(text: text, options: options)
     }
 
     func normalize(_ text: String) -> String {
@@ -211,7 +222,17 @@ public class PreTrainedTokenizer: Tokenizer {
     }
 
     public func tokenize(text: String) -> [String] {
-        preTokenize(normalize(text)).flatMap { model($0) }
+        // Take care of special tokens first
+        let sections: [String]
+        if let regex = self.addedTokensRegex {
+            sections = text.split(by: regex)
+        } else {
+            sections = [text]
+        }
+        return sections.enumerated().map { section, x in
+            if addedTokens.contains(x) { return [x] }
+            return preTokenize(normalize(x), options: section == 0 ? [.firstSection] : []).flatMap { model($0) }
+        }.flatMap { $0 }
     }
 
     /// Main entry point
@@ -241,9 +262,32 @@ public class PreTrainedTokenizer: Tokenizer {
 
 public struct AutoTokenizer {}
 
+struct PreTrainedTokenizerClasses {
+    /// Class overrides for custom behaviour
+    /// Not to be confused with the TokenizerModel classes defined in TokenizerModel
+    static let tokenizerClasses: [String : PreTrainedTokenizer.Type] = [
+        "LlamaTokenizer": LlamaPreTrainedTokenizer.self
+    ]
+}
+
 extension AutoTokenizer {
+    static func tokenizerClass(for tokenizerConfig: Config) -> PreTrainedTokenizer.Type {
+        guard let tokenizerClassName = tokenizerConfig.tokenizerClass?.stringValue else {
+            return PreTrainedTokenizer.self
+        }
+
+        // Some tokenizer_class entries use a Fast suffix
+        let tokenizerName = tokenizerClassName.replacingOccurrences(of: "Fast", with: "")
+        if let tokenizerClass = PreTrainedTokenizerClasses.tokenizerClasses[tokenizerName] {
+            return tokenizerClass
+        }
+
+        return PreTrainedTokenizer.self
+    }
+
     public static func from(tokenizerConfig: Config, tokenizerData: Config) throws -> Tokenizer {
-        return try PreTrainedTokenizer(tokenizerConfig: tokenizerConfig, tokenizerData: tokenizerData)
+        let tokenizerClass = tokenizerClass(for: tokenizerConfig)
+        return try tokenizerClass.init(tokenizerConfig: tokenizerConfig, tokenizerData: tokenizerData)
     }
 
     public static func from(
@@ -254,7 +298,7 @@ extension AutoTokenizer {
         guard let tokenizerConfig = try await config.tokenizerConfig else { throw TokenizerError.missingConfig }
         let tokenizerData = try await config.tokenizerData
 
-        return try PreTrainedTokenizer(tokenizerConfig: tokenizerConfig, tokenizerData: tokenizerData)
+        return try AutoTokenizer.from(tokenizerConfig: tokenizerConfig, tokenizerData: tokenizerData)
     }
     
     public static func from(
@@ -281,3 +325,25 @@ class CodeLlamaTokenizer: BPETokenizer {}
 class CohereTokenizer   : BPETokenizer {}
 
 class T5Tokenizer       : UnigramTokenizer {}
+
+
+// MARK: - PreTrainedTokenizer classes
+
+let sentencePieceUnderline = "▁"
+
+// See https://github.com/xenova/transformers.js/blob/1a9964fb09b8f54fcbeac46dc6aae8d76795809d/src/tokenizers.js#L3203 for these exceptions
+class LlamaPreTrainedTokenizer: PreTrainedTokenizer {
+    let isLegacy: Bool
+
+    required init(tokenizerConfig: Config, tokenizerData: Config) throws {
+        isLegacy = tokenizerConfig.legacy?.boolValue ?? true
+        var configDictionary = tokenizerData.dictionary
+        if !isLegacy {
+            configDictionary.removeValue(forKey: "normalizer")
+            configDictionary["pre_tokenizer"] = ["type": "Metaspace", "replacement": sentencePieceUnderline, "add_prefix_space": true, "prepend_scheme": "first"]
+        }
+        let updatedData = Config(configDictionary)
+
+        try super.init(tokenizerConfig: tokenizerConfig, tokenizerData: updatedData)
+    }
+}
diff --git a/Tests/TokenizersTests/AddedTokensTests.swift b/Tests/TokenizersTests/AddedTokensTests.swift
@@ -0,0 +1,67 @@
+//
+//  AddedTokensTests.swift
+//
+//
+//  Created by Pedro Cuenca on 20240426.
+//
+
+import XCTest
+import Tokenizers
+import Hub
+
+class AddedTokensTests: XCTestCase {
+    func testPhiAddedTokens() async throws {
+        let tokenizer = try await AutoTokenizer.from(pretrained: "mlx-community/Phi-3-mini-128k-instruct-4bit")
+        let inputIds = tokenizer("This is the <|end|>. My only friend, the <|end|>")
+        XCTAssertEqual(inputIds, [1, 910, 338, 278, 29871, 32007, 29889, 1619, 871, 5121, 29892, 278, 29871, 32007])
+
+        let decoded = tokenizer.decode(tokens: inputIds)
+        XCTAssertEqual(decoded, "<s> This is the <|end|>. My only friend, the <|end|>")
+    }
+
+    func testSplitWithCaptureGroups() {
+        let addedTokensRegexp = #"(<\|end\|>)\s*|(<\|raw\|>)\s*"#
+        let captureRegex = try! NSRegularExpression(pattern: addedTokensRegexp, options: [])
+
+        XCTAssertEqual(
+            "eating <|raw|> meat <|end|> That's all".split(by: captureRegex),
+            ["eating ", "<|raw|>", "meat ", "<|end|>", "That's all"]
+        )
+
+        XCTAssertEqual(
+            "<|raw|>".split(by: captureRegex),
+            ["<|raw|>"]
+        )
+
+        XCTAssertEqual(
+            "This string doesn't have those separators".split(by: captureRegex),
+            ["This string doesn't have those separators"]
+        )
+
+        XCTAssertEqual(
+            "start <|end|>".split(by: captureRegex),
+            ["start ", "<|end|>"]
+        )
+
+        XCTAssertEqual(
+            "start <|end|> ".split(by: captureRegex),
+            ["start ", "<|end|>"]
+        )
+
+        XCTAssertEqual(
+            "start <|end|>       ".split(by: captureRegex),
+            ["start ", "<|end|>"]
+        )
+
+        XCTAssertEqual(
+            "start <|end|>       for real".split(by: captureRegex),
+            ["start ", "<|end|>", "for real"]
+        )
+
+        XCTAssertEqual(
+            "<|raw|><|end|>".split(by: captureRegex),
+            ["<|raw|>", "<|end|>"]
+        )
+
+    }
+}
