Add repo revision option (#329)

* Add repo revision option
* Update ModelConfiguration inits
* Add tokenizer revision option

Author: petrukha-ivan

Libraries/MLXLMCommon/Load.swift

@@ -24,13 +24,16 @@ public func downloadModel(
 ) async throws -> URL {
     do {
         switch configuration.id {
-        case .id(let id):
+        case .id(let id, let revision):
             // download the model weights
             let repo = Hub.Repo(id: id)
             let modelFiles = ["*.safetensors", "*.json"]
             return try await hub.snapshot(
-                from: repo, matching: modelFiles, progressHandler: progressHandler)
-
+                from: repo,
+                revision: revision,
+                matching: modelFiles,
+                progressHandler: progressHandler
+            )
         case .directory(let directory):
             return directory
         }

Libraries/MLXLMCommon/ModelConfiguration.swift

@@ -9,16 +9,16 @@ import Hub
 public struct ModelConfiguration: Sendable {
 
     public enum Identifier: Sendable {
-        case id(String)
+        case id(String, revision: String = "main")
         case directory(URL)
     }
 
     public var id: Identifier
 
     public var name: String {
         switch id {
-        case .id(let string):
-            string
+        case .id(let id, _):
+            id
         case .directory(let url):
             url.deletingLastPathComponent().lastPathComponent + "/" + url.lastPathComponent
         }
@@ -37,20 +37,22 @@ public struct ModelConfiguration: Sendable {
     public var extraEOSTokens: Set<String>
 
     public init(
-        id: String, tokenizerId: String? = nil, overrideTokenizer: String? = nil,
+        id: String, revision: String = "main",
+        tokenizerId: String? = nil, overrideTokenizer: String? = nil,
         defaultPrompt: String = "hello",
         extraEOSTokens: Set<String> = [],
         preparePrompt: (@Sendable (String) -> String)? = nil
     ) {
-        self.id = .id(id)
+        self.id = .id(id, revision: revision)
         self.tokenizerId = tokenizerId
         self.overrideTokenizer = overrideTokenizer
         self.defaultPrompt = defaultPrompt
         self.extraEOSTokens = extraEOSTokens
     }
 
     public init(
-        directory: URL, tokenizerId: String? = nil, overrideTokenizer: String? = nil,
+        directory: URL,
+        tokenizerId: String? = nil, overrideTokenizer: String? = nil,
         defaultPrompt: String = "hello",
         extraEOSTokens: Set<String> = []
     ) {
@@ -63,7 +65,7 @@ public struct ModelConfiguration: Sendable {
 
     public func modelDirectory(hub: HubApi = HubApi()) -> URL {
         switch id {
-        case .id(let id):
+        case .id(let id, _):
             // download the model weights and config
             let repo = Hub.Repo(id: id)
             return hub.localRepoLocation(repo)
@@ -84,8 +86,8 @@ extension ModelConfiguration.Identifier: Equatable {
         -> Bool
     {
         switch (lhs, rhs) {
-        case (.id(let lhsID), .id(let rhsID)):
-            lhsID == rhsID
+        case (.id(let lhsID, let lhsRevision), .id(let rhsID, let rhsRevision)):
+            lhsID == rhsID && lhsRevision == rhsRevision
         case (.directory(let lhsURL), .directory(let rhsURL)):
             lhsURL == rhsURL
         default:

Libraries/MLXLMCommon/Tokenizer.swift

@@ -24,11 +24,11 @@ public func loadTokenizerConfig(configuration: ModelConfiguration, hub: HubApi)
     let config: LanguageModelConfigurationFromHub
 
     switch configuration.id {
-    case .id(let id):
+    case .id(let id, let revision):
         do {
             // the load can fail (async when we try to use it)
             let loaded = LanguageModelConfigurationFromHub(
-                modelName: configuration.tokenizerId ?? id, hubApi: hub)
+                modelName: configuration.tokenizerId ?? id, revision: revision, hubApi: hub)
             _ = try await loaded.tokenizerConfig
             config = loaded
         } catch {