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 }) + } +}