diff --git a/Sources/SstWhisper/ModelStorage.swift b/Sources/SstWhisper/ModelStorage.swift
new file mode 100644
index 0000000..a352f84
--- /dev/null
+++ b/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: `/models/argmaxinc/whisperkit-coreml/openai_whisper-`.
+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) }
+ }
+}
diff --git a/Sources/SstWhisper/WhisperKitTranscriber.swift b/Sources/SstWhisper/WhisperKitTranscriber.swift
index b64781f..e0e81bf 100644
--- a/Sources/SstWhisper/WhisperKitTranscriber.swift
+++ b/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 `/models//`.
- /// 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 {
diff --git a/Tests/SstWhisperTests/ModelStorageTests.swift b/Tests/SstWhisperTests/ModelStorageTests.swift
new file mode 100644
index 0000000..40465ab
--- /dev/null
+++ b/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 `//.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 })
+ }
+}