Refactored tokenizer config, added more supported models

Author: jkrukowski

README.md

@@ -62,6 +62,8 @@ Some of the supported models on `Hugging Face`:
 - [minishlab/potion-base-2M](https://huggingface.co/minishlab/potion-base-2M)
 - [minishlab/potion-base-4M](https://huggingface.co/minishlab/potion-base-4M)
 - [minishlab/potion-base-8M](https://huggingface.co/minishlab/potion-base-8M)
+- [minishlab/potion-retrieval-32M](https://huggingface.co/minishlab/potion-retrieval-32M)
+- [minishlab/potion-base-32M](https://huggingface.co/minishlab/potion-base-32M)
 - [minishlab/M2V_base_output](https://huggingface.co/minishlab/M2V_base_output)
 
 ### Static Embeddings

Sources/Embeddings/EmbeddingsUtils.swift

@@ -47,16 +47,21 @@ extension String {
     }
 }
 
+public enum TokenizerConfigType {
+    case filePath(String)
+    case data([String: Any])
+}
+
 public struct TokenizerConfig {
-    public let dataFileName: String
-    public let tokenizerClass: String
+    public let data: TokenizerConfigType
+    public let config: TokenizerConfigType
 
     public init(
-        dataFileName: String = "tokenizer.json",
-        tokenizerClass: String = "BertTokenizer"
+        data: TokenizerConfigType = .filePath("tokenizer.json"),
+        config: TokenizerConfigType = .filePath("tokenizer_config.json")
     ) {
-        self.dataFileName = dataFileName
-        self.tokenizerClass = tokenizerClass
+        self.data = data
+        self.config = config
     }
 }
 
@@ -104,9 +109,11 @@ extension LoadConfig {
             modelConfig: ModelConfig(
                 weightsFileName: "0_StaticEmbedding/model.safetensors"
             ),
+            // In case of `StaticEmbeddings` tokenizer `data` is loaded from `0_StaticEmbedding/tokenizer.json` file
+            // and tokenizer `config` is a dictionary with a single key `tokenizerClass` and value `BertTokenizer`.
             tokenizerConfig: TokenizerConfig(
-                dataFileName: "0_StaticEmbedding/tokenizer.json",
-                tokenizerClass: "BertTokenizer"
+                data: .filePath("0_StaticEmbedding/tokenizer.json"),
+                config: .data(["tokenizerClass": "BertTokenizer"])
             )
         )
     }

Sources/Embeddings/StaticEmbeddings/StaticEmbeddingsUtils.swift

@@ -31,9 +31,9 @@ extension StaticEmbeddings {
         let tokenizer =
             if let tokenizerConfig = loadConfig.tokenizerConfig {
                 try AutoTokenizer.from(
-                    tokenizerDataFile: modelFolder.appendingPathComponent(
-                        tokenizerConfig.dataFileName),
-                    tokenizerClass: tokenizerConfig.tokenizerClass
+                    modelFolder: modelFolder,
+                    tokenizerData: tokenizerConfig.data,
+                    tokenizerConfig: tokenizerConfig.config
                 )
             } else {
                 try await AutoTokenizer.from(modelFolder: modelFolder)
@@ -65,17 +65,35 @@ extension StaticEmbeddings {
 
 extension AutoTokenizer {
     static func from(
-        tokenizerDataFile: URL,
-        tokenizerClass: String
+        modelFolder: URL,
+        tokenizerData: TokenizerConfigType,
+        tokenizerConfig: TokenizerConfigType
     ) throws -> any Tokenizer {
-        let data = try Data(contentsOf: tokenizerDataFile)
-        let parsedData = try JSONSerialization.jsonObject(with: data, options: [])
-        guard let tokenizerData = parsedData as? [NSString: Any] else {
-            throw EmbeddingsError.invalidFile
-        }
+        let tokenizerConfig = try resolveConfig(tokenizerConfig, in: modelFolder)
+        let tokenizerData = try resolveConfig(tokenizerData, in: modelFolder)
         return try AutoTokenizer.from(
-            tokenizerConfig: Config(["tokenizerClass": tokenizerClass]),
-            tokenizerData: Config(tokenizerData)
+            tokenizerConfig: tokenizerConfig,
+            tokenizerData: tokenizerData
         )
     }
 }
+
+func resolveConfig(_ tokenizerConfig: TokenizerConfigType, in modelFolder: URL) throws -> Config {
+    switch tokenizerConfig {
+    case .filePath(let filePath):
+        let fileURL = modelFolder.appendingPathComponent(filePath)
+        let data = try loadJSONConfig(at: fileURL)
+        return Config(data as [NSString: Any])
+    case .data(let data):
+        return Config(data as [NSString: Any])
+    }
+}
+
+func loadJSONConfig(at filePath: URL) throws -> [String: Any] {
+    let data = try Data(contentsOf: filePath)
+    let parsedData = try JSONSerialization.jsonObject(with: data, options: [])
+    guard let config = parsedData as? [String: Any] else {
+        throw EmbeddingsError.invalidFile
+    }
+    return config
+}