feat: add ModelStorage and ModelCatalog; transcriber uses ModelStorage

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

Sources/SstWhisper/ModelStorage.swift

@@ -0,0 +1,73 @@
+import Foundation
+
+/// A speech model the user can pick in Settings.
+public struct ModelOption: Identifiable, Equatable, Sendable {
+    public let id: String          // WhisperKit variant name
+    public let title: String
+    public let detail: String
+    public let approxSizeMB: Int
+
+    public init(id: String, title: String, detail: String, approxSizeMB: Int) {
+        self.id = id
+        self.title = title
+        self.detail = detail
+        self.approxSizeMB = approxSizeMB
+    }
+}
+
+public enum ModelCatalog {
+    public static let defaultVariant = "large-v3-v20240930_turbo"
+
+    /// Sizes measured from the Hugging Face repository `argmaxinc/whisperkit-coreml`.
+    public static let all: [ModelOption] = [
+        ModelOption(id: defaultVariant, title: "Whisper Large v3 Turbo", detail: "Aanbevolen · nauwkeurigst", approxSizeMB: 1638),
+        ModelOption(id: "large-v3-v20240930_turbo_632MB", title: "Whisper Large v3 Turbo (compact)", detail: "Kleiner, iets minder nauwkeurig", approxSizeMB: 646),
+        ModelOption(id: "small", title: "Whisper Small", detail: "Snel en licht, minder nauwkeurig", approxSizeMB: 486),
+    ]
+}
+
+/// Where models live on disk: `<base>/models/argmaxinc/whisperkit-coreml/openai_whisper-<variant>`.
+public struct ModelStorage: Sendable {
+    public static let repo = "argmaxinc/whisperkit-coreml"
+    public static let requiredModels = ["MelSpectrogram", "AudioEncoder", "TextDecoder"]
+
+    public let base: URL
+
+    public init(base: URL = ModelStorage.defaultBase) {
+        self.base = base
+    }
+
+    /// Application Support, so models are neither in ~/Documents (privacy prompt, iCloud) nor re-fetched.
+    public static var defaultBase: URL {
+        FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0]
+            .appending(path: "Sst/models", directoryHint: .isDirectory)
+    }
+
+    var repoDirectory: URL {
+        base.appending(path: "models/\(Self.repo)", directoryHint: .isDirectory)
+    }
+
+    private func folders(for variant: String) -> [URL] {
+        let names = (try? FileManager.default.contentsOfDirectory(atPath: repoDirectory.path)) ?? []
+        return names
+            .filter { $0 == variant || $0 == "openai_whisper-\(variant)" }
+            .map { repoDirectory.appending(path: $0, directoryHint: .isDirectory) }
+    }
+
+    /// The folder of a completely downloaded variant, or nil.
+    public func folder(for variant: String) -> URL? {
+        folders(for: variant).first { folder in
+            Self.requiredModels.allSatisfy {
+                FileManager.default.fileExists(atPath: folder.appending(path: "\($0).mlmodelc/coremldata.bin").path)
+            }
+        }
+    }
+
+    public func isInstalled(_ variant: String) -> Bool {
+        folder(for: variant) != nil
+    }
+
+    public func delete(_ variant: String) throws {
+        for folder in folders(for: variant) { try FileManager.default.removeItem(at: folder) }
+    }
+}

Sources/SstWhisper/WhisperKitTranscriber.swift

@@ -10,44 +10,24 @@ public final class WhisperKitTranscriber: Transcriber {
     }
 
     private let modelName: String
+    private let storage: ModelStorage
     private var pipeline: WhisperKit?
 
     /// Download progress (0...1) during the one-time model download. May be called from any thread.
     public var onProgress: ((Double) -> Void)?
 
-    public init(model: String = "large-v3-v20240930_turbo") {
+    public init(model: String = ModelCatalog.defaultVariant, storage: ModelStorage = ModelStorage()) {
         self.modelName = model
+        self.storage = storage
     }
 
-    /// Hugging Face hub root; WhisperKit stores models in `<root>/models/<repo>/<variant>`.
-    /// Kept in Application Support so it is neither in ~/Documents (TCC prompt, iCloud) nor re-fetched.
-    private static var storageBase: URL {
-        FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0]
-            .appending(path: "Sst/models", directoryHint: .isDirectory)
-    }
+    public var variant: String { modelName }
 
-    private static let repo = "argmaxinc/whisperkit-coreml"
-
-    /// A previously downloaded variant folder, if it contains the compiled Core ML models.
-    private func localModelFolder(base: URL) -> URL? {
-        let repoDir = base.appending(path: "models/\(Self.repo)", directoryHint: .isDirectory)
-        let folders = (try? FileManager.default.contentsOfDirectory(atPath: repoDir.path)) ?? []
-        return folders
-            .filter { $0.hasSuffix(modelName) }
-            .map { repoDir.appending(path: $0, directoryHint: .isDirectory) }
-            .first { folder in
-                ["MelSpectrogram", "AudioEncoder", "TextDecoder"].allSatisfy {
-                    FileManager.default.fileExists(
-                        atPath: folder.appending(path: "\($0).mlmodelc/coremldata.bin").path)
-                }
-            }
-    }
-
-    private func loadPipeline(from folder: URL, base: URL) async throws -> WhisperKit {
+    private func loadPipeline(from folder: URL) async throws -> WhisperKit {
         try await WhisperKit(WhisperKitConfig(
-            downloadBase: base,
+            downloadBase: storage.base,
             modelFolder: folder.path,
-            tokenizerFolder: base,
+            tokenizerFolder: storage.base,
             prewarm: true,
             load: true,
             download: false
@@ -58,22 +38,21 @@ public final class WhisperKitTranscriber: Transcriber {
     /// once downloaded, launching needs no network.
     public func load() async throws {
         guard pipeline == nil else { return }
-        let base = Self.storageBase
-        try FileManager.default.createDirectory(at: base, withIntermediateDirectories: true)
+        try FileManager.default.createDirectory(at: storage.base, withIntermediateDirectories: true)
 
-        if let folder = localModelFolder(base: base),
-           let loaded = try? await loadPipeline(from: folder, base: base) {
+        if let folder = storage.folder(for: modelName),
+           let loaded = try? await loadPipeline(from: folder) {
             pipeline = loaded
             return
         }
         let folder = try await WhisperKit.download(
             variant: modelName,
-            downloadBase: base,
-            from: Self.repo
+            downloadBase: storage.base,
+            from: ModelStorage.repo
         ) { [onProgress] progress in
             onProgress?(progress.fractionCompleted)
         }
-        pipeline = try await loadPipeline(from: folder, base: base)
+        pipeline = try await loadPipeline(from: folder)
     }
 
     public func transcribe(_ samples: [Float]) async throws -> String {

Tests/SstWhisperTests/ModelStorageTests.swift

@@ -0,0 +1,63 @@
+import Foundation
+import Testing
+@testable import SstWhisper
+
+struct ModelStorageTests {
+    func makeStorage() -> ModelStorage {
+        ModelStorage(base: FileManager.default.temporaryDirectory.appending(path: "sst-models-\(UUID().uuidString)"))
+    }
+
+    /// Creates `<repo>/<folder>/<model>.mlmodelc/coremldata.bin` for the given models.
+    func install(_ storage: ModelStorage, folder: String, models: [String] = ModelStorage.requiredModels) throws {
+        for model in models {
+            let dir = storage.repoDirectory.appending(path: "\(folder)/\(model).mlmodelc", directoryHint: .isDirectory)
+            try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
+            try Data([1]).write(to: dir.appending(path: "coremldata.bin"))
+        }
+    }
+
+    @Test func emptyStorageHasNothingInstalled() {
+        let storage = makeStorage()
+        #expect(storage.folder(for: "small") == nil)
+        #expect(!storage.isInstalled("small"))
+    }
+
+    @Test func completeInstallIsFound() throws {
+        let storage = makeStorage()
+        try install(storage, folder: "openai_whisper-large-v3-v20240930_turbo")
+        #expect(storage.isInstalled("large-v3-v20240930_turbo"))
+        #expect(storage.folder(for: "large-v3-v20240930_turbo")?.lastPathComponent == "openai_whisper-large-v3-v20240930_turbo")
+    }
+
+    @Test func incompleteInstallDoesNotCount() throws {
+        let storage = makeStorage()
+        try install(storage, folder: "openai_whisper-small", models: ["MelSpectrogram", "AudioEncoder"])
+        #expect(!storage.isInstalled("small"))
+    }
+
+    @Test func variantsAreMatchedExactly() throws {
+        let storage = makeStorage()
+        try install(storage, folder: "openai_whisper-large-v3-v20240930_turbo_632MB")
+        #expect(!storage.isInstalled("large-v3-v20240930_turbo"))
+        #expect(storage.isInstalled("large-v3-v20240930_turbo_632MB"))
+    }
+
+    @Test func deleteRemovesOnlyThatVariant() throws {
+        let storage = makeStorage()
+        try install(storage, folder: "openai_whisper-small")
+        try install(storage, folder: "openai_whisper-base")
+        try storage.delete("small")
+        #expect(!storage.isInstalled("small"))
+        #expect(storage.isInstalled("base"))
+    }
+
+    @Test func deletingSomethingNotInstalledIsHarmless() throws {
+        try makeStorage().delete("small")
+    }
+
+    @Test func catalogStartsWithTheDefaultAndHasUniqueIds() {
+        #expect(ModelCatalog.all.first?.id == ModelCatalog.defaultVariant)
+        #expect(Set(ModelCatalog.all.map(\.id)).count == ModelCatalog.all.count)
+        #expect(ModelCatalog.all.allSatisfy { $0.approxSizeMB > 0 })
+    }
+}