feat: Support custom endpoint (#57)

Author: johnmai-dev

Sources/Hub/HubApi.swift

@@ -8,21 +8,22 @@
 import Foundation
 
 public struct HubApi {
-    let endpoint = "https://huggingface.co/api"
     var downloadBase: URL
-    var hfToken: String? = nil
+    var hfToken: String?
+    var endpoint: String
 
     public typealias RepoType = Hub.RepoType
     public typealias Repo = Hub.Repo
 
-    public init(downloadBase: URL? = nil, hfToken: String? = nil) {
+    public init(downloadBase: URL? = nil, hfToken: String? = nil, endpoint: String = "https://huggingface.co") {
         if downloadBase == nil {
             let documents = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first!
             self.downloadBase = documents.appending(component: "huggingface")
         } else {
             self.downloadBase = downloadBase!
         }
         self.hfToken = hfToken
+        self.endpoint = endpoint
     }
 
     static let shared = HubApi()
@@ -51,15 +52,15 @@ public extension HubApi {
         switch response.statusCode {
         case 200..<300: break
         case 400..<500: throw Hub.HubClientError.authorizationRequired
-        default       : throw Hub.HubClientError.httpStatusCode(response.statusCode)
+        default: throw Hub.HubClientError.httpStatusCode(response.statusCode)
         }
 
         return (data, response)
     }
 
     func getFilenames(from repo: Repo, matching globs: [String] = []) async throws -> [String] {
         // Read repo info and only parse "siblings"
-        let url = URL(string: "\(endpoint)/\(repo.type)/\(repo.id)")!
+        let url = URL(string: "\(endpoint)/api/\(repo.type)/\(repo.id)")!
         let (data, _) = try await httpGet(for: url)
         let response = try JSONDecoder().decode(SiblingsResponse.self, from: data)
         let filenames = response.siblings.map { $0.rfilename }
@@ -103,7 +104,7 @@ public extension HubApi {
     func whoami() async throws -> Config {
         guard hfToken != nil else { throw Hub.HubClientError.authorizationRequired }
 
-        let url = URL(string: "\(endpoint)/whoami-v2")!
+        let url = URL(string: "\(endpoint)/api/whoami-v2")!
         let (data, _) = try await httpGet(for: url)
 
         let parsed = try JSONSerialization.jsonObject(with: data, options: [])
@@ -123,15 +124,16 @@ public extension HubApi {
         let repoDestination: URL
         let relativeFilename: String
         let hfToken: String?
+        let endpoint: String?
 
         var source: URL {
             // https://huggingface.co/coreml-projects/Llama-2-7b-chat-coreml/resolve/main/tokenizer.json?download=true
-            var url = URL(string: "https://huggingface.co")!
+            var url = URL(string: endpoint ?? "https://huggingface.co")!
             if repo.type != .models {
                 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/main") // TODO: revisions
             url = url.appending(path: relativeFilename)
             return url
         }
@@ -177,7 +179,7 @@ public extension HubApi {
         let repoDestination = localRepoLocation(repo)
         for filename in filenames {
             let fileProgress = Progress(totalUnitCount: 100, parent: progress, pendingUnitCount: 1)
-            let downloader = HubFileDownloader(repo: repo, repoDestination: repoDestination, relativeFilename: filename, hfToken: hfToken)
+            let downloader = HubFileDownloader(repo: repo, repoDestination: repoDestination, relativeFilename: filename, hfToken: hfToken, endpoint: endpoint)
             try await downloader.download { fractionDownloaded in
                 fileProgress.completedUnitCount = Int64(100 * fractionDownloaded)
                 progressHandler(progress)
@@ -243,9 +245,8 @@ public extension Hub {
     }
 }
 
-public extension Array<String> {
+public extension [String] {
     func matching(glob: String) -> [String] {
-        self.filter { fnmatch(glob, $0, 0) == 0 }
+        filter { fnmatch(glob, $0, 0) == 0 }
     }
 }
-

Tests/HubTests/HubApiTests.swift

@@ -4,9 +4,8 @@
 //  Created by Pedro Cuenca on 20231230.
 //
 
-import XCTest
 @testable import Hub
-
+import XCTest
 
 class HubApiTests: XCTestCase {
     override func setUp() {
@@ -18,7 +17,7 @@ class HubApiTests: XCTestCase {
     }
 
     // MARK: use a specific revision for these tests
-    
+
     func testFilenameRetrieval() async {
         do {
             let filenames = try await Hub.getFilenames(from: "coreml-projects/Llama-2-7b-chat-coreml")
@@ -27,7 +26,7 @@ class HubApiTests: XCTestCase {
             XCTFail("\(error)")
         }
     }
-    
+
     func testFilenameRetrievalWithGlob() async {
         do {
             try await {
@@ -75,7 +74,7 @@ class HubApiTests: XCTestCase {
             XCTFail("\(error)")
         }
     }
-    
+
     func testFilenameRetrievalWithMultiplePatterns() async {
         do {
             let patterns = ["config.json", "tokenizer.json", "tokenizer_*.json"]
@@ -96,9 +95,8 @@ class SnapshotDownloadTests: XCTestCase {
         let base = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first!
         return base.appending(component: "huggingface-tests")
     }()
-    
-    override func setUp() {
-    }
+
+    override func setUp() {}
 
     override func tearDown() {
         do {
@@ -126,7 +124,7 @@ class SnapshotDownloadTests: XCTestCase {
         }
         return filenames
     }
-    
+
     func testDownload() async throws {
         let hubApi = HubApi(downloadBase: downloadDestination)
         var lastProgress: Progress? = nil
@@ -138,7 +136,31 @@ class SnapshotDownloadTests: XCTestCase {
         XCTAssertEqual(lastProgress?.fractionCompleted, 1)
         XCTAssertEqual(lastProgress?.completedUnitCount, 6)
         XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)"))
-        
+
+        let downloadedFilenames = getRelativeFiles(url: downloadDestination)
+        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",
+            ])
+        )
+    }
+
+    func testCustomEndpointDownload() async throws {
+        let hubApi = HubApi(downloadBase: downloadDestination, endpoint: "https://hf-mirror.com")
+        var lastProgress: Progress? = nil
+        let downloadedTo = try await hubApi.snapshot(from: repo, matching: "*.json") { progress in
+            print("Total Progress: \(progress.fractionCompleted)")
+            print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)")
+            lastProgress = progress
+        }
+        XCTAssertEqual(lastProgress?.fractionCompleted, 1)
+        XCTAssertEqual(lastProgress?.completedUnitCount, 6)
+        XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)"))
+
         let downloadedFilenames = getRelativeFiles(url: downloadDestination)
         XCTAssertEqual(
             Set(downloadedFilenames),