Add repo revision option (#202)

* Add repo revision option

* Add tokenizer revision option

* Add revision option to CLI

* Fix error type

* Add repo revision unit tests

* Update CLI option description

Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

* Use correct revision when retrieving filenames from API

* Use revision in getFileMetadata method

* Apply formatting

* Change error type

---------

Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

Author: petrukha-ivan

Sources/Hub/Hub.swift

@@ -155,10 +155,11 @@ public class LanguageModelConfigurationFromHub {
 
     public init(
         modelName: String,
+        revision: String = "main",
         hubApi: HubApi = .shared
     ) {
         configPromise = Task.init {
-            try await self.loadConfig(modelName: modelName, hubApi: hubApi)
+            try await self.loadConfig(modelName: modelName, revision: revision, hubApi: hubApi)
         }
     }
 
@@ -216,13 +217,14 @@ public class LanguageModelConfigurationFromHub {
 
     func loadConfig(
         modelName: String,
+        revision: String,
         hubApi: HubApi = .shared
     ) async throws -> Configurations {
         let filesToDownload = ["config.json", "tokenizer_config.json", "chat_template.json", "tokenizer.json"]
         let repo = Hub.Repo(id: modelName)
 
         do {
-            let downloadedModelFolder = try await hubApi.snapshot(from: repo, matching: filesToDownload)
+            let downloadedModelFolder = try await hubApi.snapshot(from: repo, revision: revision, matching: filesToDownload)
             return try await loadConfig(modelFolder: downloadedModelFolder, hubApi: hubApi)
         } catch {
             // Convert generic errors to more specific ones

Sources/Hub/HubApi.swift

@@ -134,16 +134,17 @@ public extension HubApi {
 
         switch response.statusCode {
         case 200..<400: break // Allow redirects to pass through to the redirect delegate
-        case 400..<500: throw Hub.HubClientError.authorizationRequired
+        case 401, 403: throw Hub.HubClientError.authorizationRequired
+        case 404: throw Hub.HubClientError.fileNotFound(url.lastPathComponent)
         default: throw Hub.HubClientError.httpStatusCode(response.statusCode)
         }
 
         return (data, response)
     }
 
-    func getFilenames(from repo: Repo, matching globs: [String] = []) async throws -> [String] {
+    func getFilenames(from repo: Repo, revision: String = "main", matching globs: [String] = []) async throws -> [String] {
         // Read repo info and only parse "siblings"
-        let url = URL(string: "\(endpoint)/api/\(repo.type)/\(repo.id)")!
+        let url = URL(string: "\(endpoint)/api/\(repo.type)/\(repo.id)/revision/\(revision)")!
         let (data, _) = try await httpGet(for: url)
         let response = try JSONDecoder().decode(SiblingsResponse.self, from: data)
         let filenames = response.siblings.map { $0.rfilename }
@@ -335,6 +336,7 @@ public extension HubApi {
     struct HubFileDownloader {
         let hub: HubApi
         let repo: Repo
+        let revision: String
         let repoDestination: URL
         let repoMetadataDestination: URL
         let relativeFilename: String
@@ -349,7 +351,7 @@ public extension HubApi {
                 url = url.appending(component: repo.type.rawValue)
             }
             url = url.appending(path: repo.id)
-            url = url.appending(path: "resolve/main") // TODO: revisions
+            url = url.appending(path: "resolve/\(revision)")
             url = url.appending(path: relativeFilename)
             return url
         }
@@ -462,7 +464,7 @@ public extension HubApi {
     }
 
     @discardableResult
-    func snapshot(from repo: Repo, matching globs: [String] = [], progressHandler: @escaping (Progress) -> Void = { _ in }) async throws -> URL {
+    func snapshot(from repo: Repo, revision: String = "main", matching globs: [String] = [], progressHandler: @escaping (Progress) -> Void = { _ in }) async throws -> URL {
         let repoDestination = localRepoLocation(repo)
         let repoMetadataDestination = repoDestination
             .appendingPathComponent(".cache")
@@ -504,13 +506,14 @@ public extension HubApi {
             return repoDestination
         }
 
-        let filenames = try await getFilenames(from: repo, matching: globs)
+        let filenames = try await getFilenames(from: repo, revision: revision, matching: globs)
         let progress = Progress(totalUnitCount: Int64(filenames.count))
         for filename in filenames {
             let fileProgress = Progress(totalUnitCount: 100, parent: progress, pendingUnitCount: 1)
             let downloader = HubFileDownloader(
                 hub: self,
                 repo: repo,
+                revision: revision,
                 repoDestination: repoDestination,
                 repoMetadataDestination: repoMetadataDestination,
                 relativeFilename: filename,
@@ -529,18 +532,18 @@ public extension HubApi {
     }
 
     @discardableResult
-    func snapshot(from repoId: String, matching globs: [String] = [], progressHandler: @escaping (Progress) -> Void = { _ in }) async throws -> URL {
-        try await snapshot(from: Repo(id: repoId), matching: globs, progressHandler: progressHandler)
+    func snapshot(from repoId: String, revision: String = "main", matching globs: [String] = [], progressHandler: @escaping (Progress) -> Void = { _ in }) async throws -> URL {
+        try await snapshot(from: Repo(id: repoId), revision: revision, matching: globs, progressHandler: progressHandler)
     }
 
     @discardableResult
-    func snapshot(from repo: Repo, matching glob: String, progressHandler: @escaping (Progress) -> Void = { _ in }) async throws -> URL {
-        try await snapshot(from: repo, matching: [glob], progressHandler: progressHandler)
+    func snapshot(from repo: Repo, revision: String = "main", matching glob: String, progressHandler: @escaping (Progress) -> Void = { _ in }) async throws -> URL {
+        try await snapshot(from: repo, revision: revision, matching: [glob], progressHandler: progressHandler)
     }
 
     @discardableResult
-    func snapshot(from repoId: String, matching glob: String, progressHandler: @escaping (Progress) -> Void = { _ in }) async throws -> URL {
-        try await snapshot(from: Repo(id: repoId), matching: [glob], progressHandler: progressHandler)
+    func snapshot(from repoId: String, revision: String = "main", matching glob: String, progressHandler: @escaping (Progress) -> Void = { _ in }) async throws -> URL {
+        try await snapshot(from: Repo(id: repoId), revision: revision, matching: [glob], progressHandler: progressHandler)
     }
 }
 
@@ -596,9 +599,9 @@ public extension HubApi {
         )
     }
 
-    func getFileMetadata(from repo: Repo, matching globs: [String] = []) async throws -> [FileMetadata] {
+    func getFileMetadata(from repo: Repo, revision: String = "main", matching globs: [String] = []) async throws -> [FileMetadata] {
         let files = try await getFilenames(from: repo, matching: globs)
-        let url = URL(string: "\(endpoint)/\(repo.id)/resolve/main")! // TODO: revisions
+        let url = URL(string: "\(endpoint)/\(repo.id)/resolve/\(revision)")!
         var selectedMetadata: [FileMetadata] = []
         for file in files {
             let fileURL = url.appending(path: file)
@@ -607,16 +610,16 @@ public extension HubApi {
         return selectedMetadata
     }
 
-    func getFileMetadata(from repoId: String, matching globs: [String] = []) async throws -> [FileMetadata] {
-        try await getFileMetadata(from: Repo(id: repoId), matching: globs)
+    func getFileMetadata(from repoId: String, revision: String = "main", matching globs: [String] = []) async throws -> [FileMetadata] {
+        try await getFileMetadata(from: Repo(id: repoId), revision: revision, matching: globs)
     }
 
-    func getFileMetadata(from repo: Repo, matching glob: String) async throws -> [FileMetadata] {
-        try await getFileMetadata(from: repo, matching: [glob])
+    func getFileMetadata(from repo: Repo, revision: String = "main", matching glob: String) async throws -> [FileMetadata] {
+        try await getFileMetadata(from: repo, revision: revision, matching: [glob])
     }
 
-    func getFileMetadata(from repoId: String, matching glob: String) async throws -> [FileMetadata] {
-        try await getFileMetadata(from: Repo(id: repoId), matching: [glob])
+    func getFileMetadata(from repoId: String, revision: String = "main", matching glob: String) async throws -> [FileMetadata] {
+        try await getFileMetadata(from: Repo(id: repoId), revision: revision, matching: [glob])
     }
 }
 

Sources/HubCLI/HubCLI.swift

@@ -48,6 +48,9 @@ struct Download: AsyncParsableCommand, SubcommandWithToken {
     @Option(help: "Repo type")
     var repoType: RepoType = .model
 
+    @Option(help: "Specific revision (e.g. branch, commit hash or tag)")
+    var revision: String = "main"
+
     @Option(help: "Glob patterns for files to include")
     var include: [String] = []
 
@@ -57,7 +60,7 @@ struct Download: AsyncParsableCommand, SubcommandWithToken {
     func run() async throws {
         let hubApi = HubApi(hfToken: hfToken)
         let repo = Hub.Repo(id: repo, type: repoType.asHubApiRepoType)
-        let downloadedTo = try await hubApi.snapshot(from: repo, matching: include) { progress in
+        let downloadedTo = try await hubApi.snapshot(from: repo, revision: revision, matching: include) { progress in
             DispatchQueue.main.async {
                 let totalPercent = 100 * progress.fractionCompleted
                 print("\(progress.completedUnitCount)/\(progress.totalUnitCount) \(totalPercent.formatted("%.02f"))%", terminator: "\r")

Tests/HubTests/HubApiTests.swift

@@ -1081,4 +1081,45 @@ class SnapshotDownloadTests: XCTestCase {
         XCTAssertTrue(FileManager.default.fileExists(atPath: filePath.path),
                       "Downloaded file should exist at \(filePath.path)")
     }
+
+    func testDownloadWithRevision() async throws {
+        let hubApi = HubApi(downloadBase: downloadDestination)
+        var lastProgress: Progress? = nil
+
+        let commitHash = "eaf97358a37d03fd48e5a87d15aff2e8423c1afb"
+        let downloadedTo = try await hubApi.snapshot(from: repo, revision: commitHash, matching: "*.json") { progress in
+            print("Total Progress: \(progress.fractionCompleted)")
+            print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)")
+            lastProgress = progress
+        }
+
+        let downloadedFilenames = getRelativeFiles(url: downloadDestination, repo: repo)
+        XCTAssertEqual(lastProgress?.fractionCompleted, 1)
+        XCTAssertEqual(lastProgress?.completedUnitCount, 6)
+        XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)"))
+        XCTAssertEqual(
+            Set(downloadedFilenames),
+            Set([
+                "config.json", "tokenizer.json", "tokenizer_config.json",
+                "llama-2-7b-chat.mlpackage/Manifest.json",
+                "llama-2-7b-chat.mlpackage/Data/com.apple.CoreML/FeatureDescriptions.json",
+                "llama-2-7b-chat.mlpackage/Data/com.apple.CoreML/Metadata.json",
+            ])
+        )
+
+        do {
+            let revision = "nonexistent-revision"
+            try await hubApi.snapshot(from: repo, revision: revision, matching: "*.json")
+            XCTFail("Expected an error to be thrown")
+        } catch let error as Hub.HubClientError {
+            switch error {
+            case .fileNotFound:
+                break // Error type is correct
+            default:
+                XCTFail("Wrong error type: \(error)")
+            }
+        } catch {
+            XCTFail("Unexpected error: \(error)")
+        }
+    }
 }