diff --git a/Package.resolved b/Package.resolved index aa170b49..50a30905 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,5 +1,14 @@ { "pins" : [ + { + "identity" : "fluidaudio", + "kind" : "remoteSourceControl", + "location" : "https://github.com/FluidInference/FluidAudio.git", + "state" : { + "revision" : "19600a485baa4998812e4654b70d2bab8f2c9949", + "version" : "0.15.5" + } + }, { "identity" : "swift-argument-parser", "kind" : "remoteSourceControl", diff --git a/Package.swift b/Package.swift index d709e114..7ce06138 100644 --- a/Package.swift +++ b/Package.swift @@ -7,6 +7,7 @@ let package = Package( dependencies: [ .package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.3.0"), .package(url: "https://github.com/argmaxinc/WhisperKit.git", from: "0.9.0"), + .package(url: "https://github.com/FluidInference/FluidAudio.git", from: "0.15.0"), ], targets: [ .executableTarget( @@ -14,6 +15,7 @@ let package = Package( dependencies: [ .product(name: "ArgumentParser", package: "swift-argument-parser"), .product(name: "WhisperKit", package: "WhisperKit"), + .product(name: "FluidAudio", package: "FluidAudio"), ] ), ] diff --git a/Sources/parrot/Models/ModelRegistry.swift b/Sources/parrot/Models/ModelRegistry.swift index 457bdfcb..04ffc7e4 100644 --- a/Sources/parrot/Models/ModelRegistry.swift +++ b/Sources/parrot/Models/ModelRegistry.swift @@ -11,29 +11,51 @@ enum ModelRegistry { id: "whisper-base.en", displayName: "Whisper Base (English)", engine: .whisperKit, - whisperKitID: "openai_whisper-base.en", + engineID: "openai_whisper-base.en", sizeMB: 145, languages: ["en"], - recommended: true + recommended: false ), TranscriptionModel( id: "whisper-large-v3-turbo", displayName: "Whisper Large v3 Turbo", engine: .whisperKit, - whisperKitID: "openai_whisper-large-v3-v20240930_turbo", + engineID: "openai_whisper-large-v3-v20240930_turbo", sizeMB: 1620, languages: ["multi"], recommended: false ), + TranscriptionModel( + id: "whisper-large-v3-turbo-compressed", + displayName: "Whisper Large v3 Turbo (compressed)", + engine: .whisperKit, + engineID: "openai_whisper-large-v3-v20240930_turbo_632MB", + sizeMB: 632, + languages: ["multi"], + recommended: false + ), TranscriptionModel( id: "whisper-small.en", displayName: "Whisper Small (English)", engine: .whisperKit, - whisperKitID: "openai_whisper-small.en", + engineID: "openai_whisper-small.en", sizeMB: 488, languages: ["en"], recommended: false ), + TranscriptionModel( + id: "parakeet-tdt-v3", + displayName: "Parakeet TDT v3", + engine: .parakeet, + engineID: "parakeet-tdt-0.6b-v3", + sizeMB: 461, + languages: ["multi"], + // Measured over 32 recordings against every other model here: best + // accuracy on the speaker's own language, faster than the rest by an + // order of magnitude, and ahead of whisper-base.en on English too + // while covering 24 more languages. + recommended: true + ), ] static func find(_ id: String) -> TranscriptionModel? { diff --git a/Sources/parrot/Models/ModelStore.swift b/Sources/parrot/Models/ModelStore.swift new file mode 100644 index 00000000..d5cd745b --- /dev/null +++ b/Sources/parrot/Models/ModelStore.swift @@ -0,0 +1,54 @@ +import Foundation + +/// Which model parrot transcribes with: chosen in the menu bar, remembered +/// across restarts. +/// +/// `--model` still wins for the launch that passes it, so a one-off run can try +/// a model without changing what the daemon comes back as. +final class ModelStore { + private static let selectionKey = "modelID" + private static let previousKey = "previousModelID" + + private let defaults: UserDefaults + + init(defaults: UserDefaults = .standard) { + self.defaults = defaults + } + + var selectedID: String? { + get { defaults.string(forKey: Self.selectionKey) } + set { + if let newValue { + defaults.set(newValue, forKey: Self.selectionKey) + } else { + defaults.removeObject(forKey: Self.selectionKey) + } + } + } + + /// The model switched away from, kept on disk so switching back does not + /// mean downloading it again. + var previousID: String? { + get { defaults.string(forKey: Self.previousKey) } + set { + if let newValue { + defaults.set(newValue, forKey: Self.previousKey) + } else { + defaults.removeObject(forKey: Self.previousKey) + } + } + } + + var previous: TranscriptionModel? { + previousID.flatMap(ModelRegistry.find) + } + + /// Model to start with: the flag, else the remembered choice, else the + /// registry's recommendation. A remembered model that no longer exists in + /// the registry falls through instead of failing the launch. + func resolved(flag: String?) -> TranscriptionModel? { + if let flag { return ModelRegistry.find(flag) } + if let selectedID, let model = ModelRegistry.find(selectedID) { return model } + return ModelRegistry.recommended() + } +} diff --git a/Sources/parrot/Models/ModelWeights.swift b/Sources/parrot/Models/ModelWeights.swift new file mode 100644 index 00000000..80aab881 --- /dev/null +++ b/Sources/parrot/Models/ModelWeights.swift @@ -0,0 +1,87 @@ +import Foundation + +/// Where model weights live on disk, and how to get the space back. +/// +/// Every engine downloads under one root so switching models can reclaim the +/// previous one without knowing which engine wrote it. +enum ModelWeights { + /// Not `Documents`: iCloud replicates that folder and evicts large files to + /// dataless stubs, and CoreML's mmap of an evicted weight file blocks forever. + static let root = URL.applicationSupportDirectory.appending(path: "parrot") + + /// WhisperKit appends `models//` to whatever base it is given. + static let whisperKitBase = root.appending(path: "huggingface") + + static func directory(of model: TranscriptionModel) -> URL? { + guard let engineID = model.engineID else { return nil } + switch model.engine { + case .whisperKit: + return whisperKitBase.appending(path: "models/argmaxinc/whisperkit-coreml/\(engineID)") + case .parakeet: + return root.appending(path: "parakeet/\(engineID)") + } + } + + static func isInstalled(_ model: TranscriptionModel) -> Bool { + guard let directory = directory(of: model) else { return false } + return FileManager.default.fileExists(atPath: directory.path(percentEncoded: false)) + } + + /// Deletes the weights of every registered model except the ones named. + /// Called after a switch succeeds, so a failed download never costs the + /// user the model they had. The model switched away from is kept too: + /// otherwise going back and forth between two engines re-downloads a + /// gigabyte each way. + @discardableResult + static func purge(keeping wanted: [TranscriptionModel]) -> Int64 { + let keepIDs = Set(wanted.map(\.id)) + var reclaimed: Int64 = 0 + for model in ModelRegistry.shared where !keepIDs.contains(model.id) { + guard let directory = directory(of: model), + FileManager.default.fileExists(atPath: directory.path(percentEncoded: false)) + else { continue } + reclaimed += bytes(at: directory) + try? FileManager.default.removeItem(at: directory) + } + return reclaimed + sweep(keeping: wanted) + } + + /// WhisperKit keeps tokenizers and download caches beside the weights, under + /// folders the per-model delete above never names. A folder survives only + /// when a kept model's engine id contains its name; guessing wrong costs a + /// few MB of re-download, never a broken model. + private static func sweep(keeping wanted: [TranscriptionModel]) -> Int64 { + let roots = ["models/openai", "models/argmaxinc/whisperkit-coreml"] + let keep = wanted.filter { $0.engine == .whisperKit }.compactMap(\.engineID) + var reclaimed: Int64 = 0 + for root in roots { + let directory = whisperKitBase.appending(path: root) + let children = (try? FileManager.default.contentsOfDirectory( + at: directory, includingPropertiesForKeys: nil)) ?? [] + for child in children + where !keep.contains(where: { $0.contains(child.lastPathComponent) }) { + reclaimed += bytes(at: child) + try? FileManager.default.removeItem(at: child) + } + } + return reclaimed + } + + static func bytes(at directory: URL) -> Int64 { + guard let walker = FileManager.default.enumerator( + at: directory, + includingPropertiesForKeys: [.totalFileAllocatedSizeKey, .fileAllocatedSizeKey]) + else { return 0 } + var total: Int64 = 0 + for case let url as URL in walker { + let size = try? url.resourceValues(forKeys: [.totalFileAllocatedSizeKey, .fileAllocatedSizeKey]) + total += Int64(size?.totalFileAllocatedSize ?? size?.fileAllocatedSize ?? 0) + } + return total + } + + static func describe(bytes: Int64) -> String { + let mb = Double(bytes) / 1_000_000 + return mb >= 1000 ? String(format: "%.1f GB", mb / 1000) : String(format: "%.0f MB", mb) + } +} diff --git a/Sources/parrot/Models/TranscriptionModel.swift b/Sources/parrot/Models/TranscriptionModel.swift index 8f1d53f5..6c50f7b4 100644 --- a/Sources/parrot/Models/TranscriptionModel.swift +++ b/Sources/parrot/Models/TranscriptionModel.swift @@ -9,8 +9,8 @@ struct TranscriptionModel: Codable { let id: String let displayName: String let engine: Engine - /// Engine-specific identifier (e.g. "openai_whisper-base.en" for WhisperKit). - let whisperKitID: String? + /// The id the engine itself knows the model by. + let engineID: String? let sizeMB: Int let languages: [String] let recommended: Bool diff --git a/Sources/parrot/Parrot.swift b/Sources/parrot/Parrot.swift index 05a69ebe..037d3b30 100644 --- a/Sources/parrot/Parrot.swift +++ b/Sources/parrot/Parrot.swift @@ -34,6 +34,14 @@ struct Run: ParsableCommand { @Option(name: .long, help: "Model id to use. Defaults to the recommended model.") var model: String? + /// Both purge sites report the same way; the number is the only reason the + /// user would care that a purge happened at all. + private func reclaim(_ bytes: Int64) { + guard bytes > 0 else { return } + FileHandle.standardError.write(Data( + "freed \(ModelWeights.describe(bytes: bytes)) of unused models\n".utf8)) + } + func run() throws { if !skipDoctor { let checks = DoctorReport.run() @@ -45,23 +53,18 @@ struct Run: ParsableCommand { } } - let chosenModel: TranscriptionModel - if let id = model { - guard let m = ModelRegistry.find(id) else { + let models = ModelStore() + guard let chosenModel = models.resolved(flag: model) else { + if let id = model { FileHandle.standardError.write(Data("unknown model: \(id)\n".utf8)) FileHandle.standardError.write(Data("run `parrot models list` to see options.\n".utf8)) throw ExitCode(1) } - chosenModel = m - } else { - guard let m = ModelRegistry.recommended() else { - FileHandle.standardError.write(Data("no models registered\n".utf8)) - throw ExitCode(1) - } - chosenModel = m + FileHandle.standardError.write(Data("no models registered\n".utf8)) + throw ExitCode(1) } - let transcriber = WhisperKitTranscriber(model: chosenModel) + let transcriber = ActiveTranscriber(model: chosenModel) let warmupSemaphore = DispatchSemaphore(value: 0) var warmupError: Error? Task.detached { @@ -78,6 +81,11 @@ struct Run: ParsableCommand { throw ExitCode(1) } + // Whatever an earlier run left behind is dead weight now that this model + // is loaded — except the one switched away from, which the user is one + // menu click from wanting again. + reclaim(ModelWeights.purge(keeping: [chosenModel, models.previous].compactMap { $0 })) + let app = NSApplication.shared app.setActivationPolicy(.accessory) @@ -88,7 +96,27 @@ struct Run: ParsableCommand { if let overlay { capture.onLevel = { level in overlay.pushLevel(level) } } - let menuBar = MainActor.assumeIsolated { MenuBarController(modelID: chosenModel.id) } + let menuBar = MainActor.assumeIsolated { MenuBarController(model: chosenModel) } + MainActor.assumeIsolated { + menuBar.onModel = { [weak menuBar] next in + Task { + do { + let outgoing = try await transcriber.use(next) { fraction in + Task { @MainActor in menuBar?.setSwitchProgress(fraction) } + } + // Only now: a model that failed to load is not what the + // daemon should come back as after a restart. + models.selectedID = next.id + models.previousID = outgoing?.id + reclaim(ModelWeights.purge(keeping: [next, outgoing].compactMap { $0 })) + await MainActor.run { menuBar?.setModel(next) } + } catch { + FileHandle.standardError.write(Data("switch to \(next.id) failed: \(error)\n".utf8)) + await MainActor.run { menuBar?.setModelFailed(next) } + } + } + } + } do { try monitor.start { event in diff --git a/Sources/parrot/Transcription/ActiveTranscriber.swift b/Sources/parrot/Transcription/ActiveTranscriber.swift new file mode 100644 index 00000000..01f7acb0 --- /dev/null +++ b/Sources/parrot/Transcription/ActiveTranscriber.swift @@ -0,0 +1,49 @@ +import Foundation + +/// The engine dictation currently speaks to, and the only thing that knows how +/// to replace it while the daemon keeps running. +/// +/// Switching is serialised against transcription by the actor itself: a press +/// that lands mid-switch waits for the new model rather than racing a +/// half-unloaded one. +actor ActiveTranscriber { + private(set) var model: TranscriptionModel + private var transcriber: Transcriber + + init(model: TranscriptionModel) { + self.model = model + self.transcriber = Self.make(model) + } + + private static func make(_ model: TranscriptionModel) -> Transcriber { + switch model.engine { + case .whisperKit: WhisperKitTranscriber(model: model) + case .parakeet: ParakeetTranscriber(model: model) + } + } + + func warmUp(onProgress: (@Sendable (Double) -> Void)? = nil) async throws { + try await transcriber.warmUp(onProgress: onProgress) + } + + func transcribe(_ audio: [Float]) async throws -> String { + try await transcriber.transcribe(audio) + } + + /// Loads `next` before dropping the current engine, so a download that fails + /// leaves the user dictating with what they already had. Returns the model + /// that was replaced, which is the one worth keeping on disk. + @discardableResult + func use(_ next: TranscriptionModel, + onProgress: (@Sendable (Double) -> Void)? = nil) async throws -> TranscriptionModel? { + guard next.id != model.id else { return nil } + let incoming = Self.make(next) + try await incoming.warmUp(onProgress: onProgress) + + let outgoing = model + await transcriber.unload() + transcriber = incoming + model = next + return outgoing + } +} diff --git a/Sources/parrot/Transcription/ParakeetTranscriber.swift b/Sources/parrot/Transcription/ParakeetTranscriber.swift new file mode 100644 index 00000000..7fc060c6 --- /dev/null +++ b/Sources/parrot/Transcription/ParakeetTranscriber.swift @@ -0,0 +1,56 @@ +import FluidAudio +import Foundation + +/// Parakeet TDT (CoreML) via FluidAudio. +/// +/// Measured against Whisper large-v3-turbo on the same corpus: an order of +/// magnitude faster, better on short utterances carrying English technical +/// terms, worse on long sentences. Offered as a choice rather than a default +/// because which side of that trade matters depends on how you dictate. +actor ParakeetTranscriber: Transcriber { + let modelID: String + private let model: TranscriptionModel + private var manager: AsrManager? + + init(model: TranscriptionModel) { + self.modelID = model.id + self.model = model + } + + func warmUp(onProgress: (@Sendable (Double) -> Void)? = nil) async throws { + if manager != nil { return } + guard model.engineID != nil, let directory = ModelWeights.directory(of: model) else { + throw TranscriberError.missingEngineID + } + FileHandle.standardError.write(Data("loading \(model.id)...\n".utf8)) + let weights = try await AsrModels.downloadAndLoad( + to: directory, version: .v3, + progressHandler: { progress in onProgress?(progress.fractionCompleted) }) + let manager = AsrManager(config: .default) + try await manager.loadModels(weights) + self.manager = manager + FileHandle.standardError.write(Data("✓ \(model.id) ready\n".utf8)) + } + + func transcribe(_ audio: [Float]) async throws -> String { + if manager == nil { try await warmUp() } + guard let manager else { throw TranscriberError.notLoaded } + + // Shorter than the engine's own minimum it throws instead of returning + // nothing, and a tap on the hotkey is silence rather than a failure. + let minimum = ASRConstants.minimumRequiredSamples( + forSampleRate: Int(AudioCapture.targetSampleRate)) + guard audio.count >= minimum else { return "" } + + // Dictation hands over one finished utterance at a time, so the decoder + // starts clean rather than continuing the previous phrase's state. + var state = try TdtDecoderState() + let result = try await manager.transcribe(audio, decoderState: &state) + return result.text.trimmingCharacters(in: .whitespacesAndNewlines) + } + + func unload() async { + await manager?.cleanup() + manager = nil + } +} diff --git a/Sources/parrot/Transcription/Transcriber.swift b/Sources/parrot/Transcription/Transcriber.swift index d57857f9..e83c50a7 100644 --- a/Sources/parrot/Transcription/Transcriber.swift +++ b/Sources/parrot/Transcription/Transcriber.swift @@ -2,5 +2,13 @@ import Foundation protocol Transcriber { var modelID: String { get } + /// Downloads if needed and loads into memory, so the first hotkey press is + /// not the thing that waits. `onProgress` reports the download as a fraction + /// so a switch that takes minutes can say so; loading afterwards is silent + /// because neither engine reports it. + func warmUp(onProgress: (@Sendable (Double) -> Void)?) async throws func transcribe(_ audio: [Float]) async throws -> String + /// Releases the loaded weights. Switching models calls this on the outgoing + /// engine so two models never sit in memory at once. + func unload() async } diff --git a/Sources/parrot/Transcription/WhisperKitTranscriber.swift b/Sources/parrot/Transcription/WhisperKitTranscriber.swift index 8003194b..0b965f48 100644 --- a/Sources/parrot/Transcription/WhisperKitTranscriber.swift +++ b/Sources/parrot/Transcription/WhisperKitTranscriber.swift @@ -2,6 +2,11 @@ import Foundation import WhisperKit actor WhisperKitTranscriber: Transcriber { + /// Earlier versions downloaded into WhisperKit's `Documents` default. Say so + /// once rather than silently re-downloading gigabytes behind the user's back. + static let legacyModelStore = URL.documentsDirectory.appending(path: "huggingface") + + let modelID: String private let model: TranscriptionModel private var pipeline: WhisperKit? @@ -14,17 +19,49 @@ actor WhisperKitTranscriber: Transcriber { /// Loads the model into memory; downloads first if not already on disk. /// Call once at startup so the first hotkey press isn't blocked on model /// download/load. - func warmUp() async throws { + func warmUp(onProgress: (@Sendable (Double) -> Void)? = nil) async throws { if pipeline != nil { return } - guard let whisperKitID = model.whisperKitID else { + guard let engineID = model.engineID else { throw TranscriberError.missingEngineID } + Self.noteLegacyStore() FileHandle.standardError.write(Data("loading \(model.id)...\n".utf8)) - let config = WhisperKitConfig(model: whisperKitID, verbose: false, prewarm: true, load: true) + // Downloading separately from loading is what makes progress visible: + // WhisperKit reports it on the static download and not on init. + let folder = try await WhisperKit.download( + variant: engineID, + downloadBase: ModelWeights.whisperKitBase, + progressCallback: { progress in onProgress?(progress.fractionCompleted) }) + // `downloadBase` is not redundant beside an explicit `modelFolder`: the + // tokenizer is fetched separately and lands under the base. Drop it and + // tokenizers go back to `Documents`, which is the eviction this avoids. + let config = WhisperKitConfig( + model: engineID, + downloadBase: ModelWeights.whisperKitBase, + modelFolder: folder.path(percentEncoded: false), + verbose: false, + prewarm: true, + load: true, + download: false + ) pipeline = try await WhisperKit(config) FileHandle.standardError.write(Data("✓ \(model.id) ready\n".utf8)) } + /// Point at leftover models from the old `Documents` location so the disk + /// space is reclaimable instead of just abandoned. + private static func noteLegacyStore() { + let path = legacyModelStore.path(percentEncoded: false) + guard FileManager.default.fileExists(atPath: path) else { return } + FileHandle.standardError.write(Data( + "models now live in \(ModelWeights.whisperKitBase.path(percentEncoded: false)) — \(path) is unused, delete it to reclaim space\n".utf8 + )) + } + + func unload() async { + pipeline = nil + } + func transcribe(_ audio: [Float]) async throws -> String { if pipeline == nil { try await warmUp() } guard let pipeline else { throw TranscriberError.notLoaded } diff --git a/Sources/parrot/UI/MenuBarController.swift b/Sources/parrot/UI/MenuBarController.swift index 366ab060..0921f236 100644 --- a/Sources/parrot/UI/MenuBarController.swift +++ b/Sources/parrot/UI/MenuBarController.swift @@ -4,14 +4,21 @@ import AppKit /// a glance and provides the only persistent control surface for the daemon /// (since we run as `.accessory` — no dock icon, no main window). @MainActor -final class MenuBarController { +final class MenuBarController: NSObject, NSMenuDelegate { private let statusItem: NSStatusItem private let modelLabel: NSMenuItem private let stateLabel: NSMenuItem - private let modelID: String + private let modelItem: NSMenuItem + private var model: TranscriptionModel + /// The model being switched to. Set on click so the menu answers the click + /// rather than the download, which can take minutes. + private var pending: TranscriptionModel? + /// Set after construction: switching needs the daemon, and the daemon needs + /// the controller it reports back to. + var onModel: ((TranscriptionModel) -> Void)? - init(modelID: String) { - self.modelID = modelID + init(model: TranscriptionModel) { + self.model = model self.statusItem = NSStatusBar.system.statusItem(withLength: NSStatusItem.variableLength) let menu = NSMenu() @@ -21,10 +28,16 @@ final class MenuBarController { stateLabel.isEnabled = false menu.addItem(stateLabel) - modelLabel = NSMenuItem(title: "model: \(modelID)", action: nil, keyEquivalent: "") + modelLabel = NSMenuItem(title: "model: \(model.id)", action: nil, keyEquivalent: "") modelLabel.isEnabled = false menu.addItem(modelLabel) + modelItem = NSMenuItem(title: "Model", action: nil, keyEquivalent: "") + let modelMenu = NSMenu() + modelMenu.autoenablesItems = false + modelItem.submenu = modelMenu + menu.addItem(modelItem) + menu.addItem(.separator()) let quit = NSMenuItem( @@ -32,13 +45,72 @@ final class MenuBarController { action: #selector(quitClicked), keyEquivalent: "q" ) + super.init() + quit.target = self menu.addItem(quit) + // Rebuild the model list on open so a switch in flight is reflected + // without the menu holding its own copy of the state. + menu.delegate = self + statusItem.menu = menu configureButton(recording: false) } + func menuWillOpen(_ menu: NSMenu) { + guard let modelSubmenu = modelItem.submenu else { return } + modelSubmenu.removeAllItems() + for candidate in ModelRegistry.shared { + let active = candidate.id == model.id + let arriving = candidate.id == pending?.id + let item = NSMenuItem( + title: "\(candidate.displayName) · \(candidate.sizeMB) MB", + action: #selector(modelSelected), + keyEquivalent: "") + item.target = self + item.representedObject = candidate.id + // A dash rather than a tick while it arrives: the choice is taken, + // the model is not ready, and pretending otherwise is what made the + // menu look like it ignored the click. + item.state = arriving ? .mixed : (active ? .on : .off) + // One switch at a time, and the model in use is already here. + item.isEnabled = pending == nil && !active + modelSubmenu.addItem(item) + } + } + + /// Hands the choice to the daemon, which loads it before dropping the model + /// in use. The label follows on `setModel` once that succeeds. + @objc private func modelSelected(_ sender: NSMenuItem) { + guard let id = sender.representedObject as? String, + let chosen = ModelRegistry.find(id), chosen.id != model.id + else { return } + pending = chosen + modelLabel.title = "loading \(chosen.id)…" + onModel?(chosen) + } + + /// Download fraction, which is the only part either engine reports. Loading + /// afterwards shows as the same line without a number. + func setSwitchProgress(_ fraction: Double) { + guard let pending else { return } + modelLabel.title = fraction < 1 + ? String(format: "downloading %@… %.0f%%", pending.id, fraction * 100) + : "loading \(pending.id)…" + } + + func setModel(_ model: TranscriptionModel) { + self.model = model + self.pending = nil + modelLabel.title = "model: \(model.id)" + } + + func setModelFailed(_ attempted: TranscriptionModel) { + pending = nil + modelLabel.title = "model: \(model.id) · \(attempted.id) failed" + } + func setRecording(_ recording: Bool) { stateLabel.title = recording ? "● recording" : "idle · hold fn to dictate" }