menu bar: switch models, and keep only the one in use on disk

Package.swift

@@ -7,13 +7,15 @@ 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(
             name: "parrot",
             dependencies: [
                 .product(name: "ArgumentParser", package: "swift-argument-parser"),
                 .product(name: "WhisperKit", package: "WhisperKit"),
+                .product(name: "FluidAudio", package: "FluidAudio"),
             ]
         ),
     ]

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? {

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

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/<repo>/<id>` 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)
+    }
+}

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

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

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