Create LlamaPreTrainedTokenizer subclass

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

Author: pcuenca

Sources/Tokenizers/Tokenizer.swift

@@ -262,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(
@@ -275,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(
@@ -302,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)
+    }
+}