feat(mobile): on-device dictation models, voice filter and correction learning on iOS

apps/mobile/app.config.ts

@@ -490,6 +490,7 @@ const config: ExpoConfig = {
       },
     ],
     "./plugins/withIosCocoaPodsUuidCache.cjs",
+    "./plugins/withDictationPods.cjs",
     // Compile out the unused pedometer to avoid an iOS motion purpose string.
     ["expo-sensors", { motionPermission: false }],
     // Mods run in reverse order; expo-widgets must create and wipe the target before assets are added.

apps/mobile/modules/pseudo-dictation/.gitignore

@@ -0,0 +1,2 @@
+.build/
+.swiftpm/

apps/mobile/modules/pseudo-dictation/Package.swift

@@ -0,0 +1,29 @@
+// swift-tools-version: 6.0
+import PackageDescription
+
+// The app builds this module through CocoaPods. This manifest exists so the core can be compiled
+// and tested on a Mac without an iOS build. Keep the FluidAudio revision equal to the one in
+// `apps/mobile/plugins/withDictationPods.cjs`.
+let package = Package(
+  name: "PseudoDictation",
+  platforms: [.macOS(.v14), .iOS(.v18)],
+  dependencies: [
+    .package(
+      url: "https://github.com/FluidInference/FluidAudio.git",
+      revision: "c8b41f98eae668794dd0083bbe028fe7e6585051")
+  ],
+  targets: [
+    .target(
+      name: "PseudoDictationCore",
+      dependencies: [.product(name: "FluidAudio", package: "FluidAudio")],
+      path: "ios/Core",
+      swiftSettings: [.swiftLanguageMode(.v5)]
+    ),
+    .testTarget(
+      name: "PseudoDictationCoreTests",
+      dependencies: ["PseudoDictationCore"],
+      path: "Tests/PseudoDictationCoreTests",
+      swiftSettings: [.swiftLanguageMode(.v5)]
+    ),
+  ]
+)

apps/mobile/modules/pseudo-dictation/THIRD_PARTY_NOTICES.md

@@ -0,0 +1,32 @@
+# Third-party notices
+
+On-device dictation in PseudoCode is built on the work below.
+
+## Shipped in the app
+
+**FluidAudio** by Fluid Inference, Apache License 2.0.
+https://github.com/FluidInference/FluidAudio (commit `c8b41f98eae668794dd0083bbe028fe7e6585051`, v0.18.3).
+Its licence text and the notices for the code it bundles are in that repository's `LICENSE` and
+`ThirdPartyLicenses/`.
+
+**`ios/Resources/dictation-sample.wav`** is FluidAudio's test fixture
+`Tests/FluidAudioTests/ASR/Parakeet/SlidingWindow/Fixtures/03-diff-explanation-16.9s.wav`, unchanged,
+under the same Apache License 2.0.
+
+## Downloaded on the device when the user asks for them
+
+These models are not part of the app. They are fetched from Hugging Face. All are Core ML
+conversions by Fluid Inference, modified from the originals named here.
+
+| Model in the app | Repository                                                                                    | Original                                                                                                                  | Licence   |
+| ---------------- | --------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------- | --------- |
+| Fast             | `FluidInference/parakeet-tdt-ctc-110m-coreml`                                                 | NVIDIA `parakeet-tdt_ctc-110m`                                                                                            | CC BY 4.0 |
+| Balanced         | `FluidInference/parakeet-tdt-0.6b-v3-coreml`                                                  | NVIDIA `parakeet-tdt-0.6b-v3`                                                                                             | CC BY 4.0 |
+| Most accurate    | `FluidInference/parakeet-ultra-coreml`                                                        | moondream `parakeet-ultra`, trained from NVIDIA Parakeet v3                                                               | CC BY 4.0 |
+| Voice filter     | `FluidInference/speaker-diarization-coreml` (`Segmentation`, `FBank`, `Embedding`, `PldaRho`) | `pyannote/speaker-diarization-community-1` by pyannote, with WeSpeaker embeddings and PLDA parameters from BUT Speech@FIT | CC BY 4.0 |
+
+CC BY 4.0: https://creativecommons.org/licenses/by/4.0/
+
+The voice filter repository holds files under more than one licence. Its `NOTICE.md` says which
+licence covers which file; the 4 files above are in its CC BY 4.0 set:
+https://huggingface.co/FluidInference/speaker-diarization-coreml/blob/main/NOTICE.md

apps/mobile/modules/pseudo-dictation/Tests/PseudoDictationCoreTests/DictationEngineTests.swift

@@ -0,0 +1,174 @@
+import FluidAudio
+import XCTest
+
+@testable import PseudoDictationCore
+
+private final class Recorder: @unchecked Sendable {
+  private let lock = NSLock()
+  private var stored: [DownloadUpdate] = []
+
+  var updates: [DownloadUpdate] {
+    lock.lock()
+    defer { lock.unlock() }
+    return stored
+  }
+
+  func add(_ update: DownloadUpdate) {
+    lock.lock()
+    stored.append(update)
+    lock.unlock()
+  }
+}
+
+private func environment(freeDisk: UInt64? = nil, availableMemory: UInt64 = 0) -> DictationEnvironment {
+  DictationEnvironment(
+    appBuild: "1.0(1)", osBuild: "test", isSimulator: true,
+    availableMemory: { availableMemory }, freeDisk: { _ in freeDisk })
+}
+
+final class ProgressRelayTests: XCTestCase {
+  func testLibraryProgressMapsDownloadToTheWholeBar() {
+    XCTAssertEqual(
+      DownloadUpdate(DownloadProgress(fractionCompleted: 0, phase: .listing)),
+      DownloadUpdate(phase: .downloading, fraction: 0))
+    XCTAssertEqual(
+      DownloadUpdate(
+        DownloadProgress(
+          fractionCompleted: 0.25, phase: .downloading(completedFiles: 1, totalFiles: 2))),
+      DownloadUpdate(phase: .downloading, fraction: 0.5))
+    XCTAssertEqual(
+      DownloadUpdate(
+        DownloadProgress(
+          fractionCompleted: 0.5, phase: .downloading(completedFiles: 0, totalFiles: 0))),
+      DownloadUpdate(phase: .downloading, fraction: 1))
+    XCTAssertEqual(
+      DownloadUpdate(
+        DownloadProgress(fractionCompleted: 0.75, phase: .compiling(modelName: "Encoder"))),
+      DownloadUpdate(phase: .preparing, fraction: 1))
+  }
+
+  func testProgressNeverMovesBackwardsAndSkipsSmallSteps() {
+    let recorder = Recorder()
+    let relay = ProgressRelay(emit: { recorder.add($0) })
+    for fraction in [0, 0.004, 0.2, 0.1, 0.204, 0.21, 0.996, 1] {
+      relay.report(DownloadUpdate(phase: .downloading, fraction: fraction))
+    }
+    XCTAssertEqual(recorder.updates.map(\.fraction), [0, 0.2, 0.21, 0.996])
+  }
+
+  func testACachedLaterComponentDoesNotReopenTheDownloadPhase() {
+    let recorder = Recorder()
+    let relay = ProgressRelay(emit: { recorder.add($0) })
+    relay.report(DownloadUpdate(phase: .downloading, fraction: 0.9))
+    relay.report(DownloadUpdate(phase: .preparing, fraction: 1))
+    relay.report(DownloadUpdate(phase: .downloading, fraction: 1))
+    relay.report(DownloadUpdate(phase: .preparing, fraction: 1))
+    XCTAssertEqual(recorder.updates.map(\.phase), [.downloading, .preparing])
+  }
+}
+
+final class DictationEngineTests: XCTestCase {
+  private var storage: DictationStorage!
+
+  override func setUpWithError() throws {
+    storage = try makeTemporaryStorage()
+  }
+
+  override func tearDownWithError() throws {
+    try? FileManager.default.removeItem(at: storage.root)
+  }
+
+  private func assertThrows<T>(
+    _ expected: DictationError, _ body: () async throws -> T, file: StaticString = #filePath,
+    line: UInt = #line
+  ) async {
+    do {
+      _ = try await body()
+      XCTFail("Expected \(expected)", file: file, line: line)
+    } catch {
+      XCTAssertEqual(error as? DictationError, expected, file: file, line: line)
+    }
+  }
+
+  func testDownloadIsRefusedBeforeAnyNetworkUseWhenStorageIsShort() async {
+    let engine = DictationEngine(storage: storage, environment: environment(freeDisk: 10 << 20))
+    await assertThrows(
+      .notEnoughDisk(requiredBytes: DictationModel.fast.downloadDiskBytes, availableBytes: 10 << 20)
+    ) {
+      try await engine.download(.fast, processing: .auto, onProgress: { _ in })
+    }
+    XCTAssertFalse(FileManager.default.fileExists(atPath: storage.modelDirectory(.fast).path))
+  }
+
+  func testPrepareAndTranscribeRefuseAModelThatIsNotDownloaded() async {
+    let engine = DictationEngine(storage: storage, environment: environment())
+    await assertThrows(.notInstalled(.fast)) {
+      try await engine.prepare(.fast, processing: .auto)
+    }
+    await assertThrows(.notInstalled(.balanced)) {
+      try await engine.transcribe(
+        requestId: "1", url: URL(fileURLWithPath: "/dev/null"), model: .balanced,
+        processing: .auto, onlyMyVoice: false, language: "en-US")
+    }
+  }
+
+  func testACancelThatArrivesFirstStopsItsTranscriptionAndNoOther() async {
+    let engine = DictationEngine(storage: storage, environment: environment())
+    let transcribe = { (requestId: String) in
+      try await engine.transcribe(
+        requestId: requestId, url: URL(fileURLWithPath: "/dev/null"), model: .fast,
+        processing: .auto, onlyMyVoice: false, language: "en-US")
+    }
+    await engine.cancelTranscription(requestId: "early")
+    await assertThrows(.notInstalled(.fast)) { try await transcribe("other") }
+    await assertThrows(.cancelled) { try await transcribe("early") }
+    await assertThrows(.notInstalled(.fast)) { try await transcribe("early") }
+  }
+
+  func testTheVoiceFilterIsNotASpeechModel() async {
+    let engine = DictationEngine(storage: storage, environment: environment())
+    await assertThrows(.unknownModel("voice-filter")) {
+      try await engine.prepare(.voiceFilter, processing: .auto)
+    }
+  }
+
+  func testStatusOfAModelThatIsNotDownloaded() async {
+    let engine = DictationEngine(storage: storage, environment: environment())
+    let status = await engine.status(.fast, processing: .auto)
+    XCTAssertEqual(
+      status, ModelStatus(installed: false, warm: false, downloading: false, sizeOnDiskBytes: 0))
+  }
+
+  func testAWarmMarkerDoesNotOutliveTheModelFiles() async throws {
+    try storage.markWarm(
+      .fast, key: WarmKey(appBuild: "1.0(1)", osBuild: "test", processing: "0"))
+    let engine = DictationEngine(storage: storage, environment: environment())
+    let status = await engine.status(.fast, processing: .auto)
+    XCTAssertFalse(status.warm)
+  }
+
+  func testSavingAProfileNeedsEveryClip() async {
+    let engine = DictationEngine(storage: storage, environment: environment())
+    await assertThrows(.enrollmentRejected(.tooFewClips)) {
+      try await engine.saveVoiceProfile([URL(fileURLWithPath: "/dev/null")])
+    }
+    let status = await engine.voiceProfileStatus()
+    XCTAssertFalse(status.enrolled)
+  }
+
+  func testLanguageHintComesFromTheLocalePrefix() {
+    XCTAssertEqual(DictationEngine.languageHint("en-US"), .english)
+    XCTAssertEqual(DictationEngine.languageHint("de_DE"), .german)
+    XCTAssertEqual(DictationEngine.languageHint("FR"), .french)
+    XCTAssertNil(DictationEngine.languageHint("ja-JP"))
+    XCTAssertNil(DictationEngine.languageHint(nil))
+    XCTAssertNil(DictationEngine.languageHint(""))
+  }
+
+  func testForeignErrorsAreDescribedAsCodedDictationErrors() {
+    XCTAssertEqual(DictationEngine.describe(CancellationError()), .cancelled)
+    XCTAssertEqual(DictationEngine.describe(DictationError.noVoiceProfile), .noVoiceProfile)
+    XCTAssertEqual(DictationEngine.describe(DictationError.notInstalled(.fast)).code, "not-installed")
+    XCTAssertEqual(DictationEngine.describe(URLError(.timedOut)).code, "failed")
+  }
+}

apps/mobile/modules/pseudo-dictation/Tests/PseudoDictationCoreTests/DictationStorageTests.swift

@@ -0,0 +1,131 @@
+import FluidAudio
+import XCTest
+
+@testable import PseudoDictationCore
+
+func makeTemporaryStorage() throws -> DictationStorage {
+  let root = FileManager.default.temporaryDirectory
+    .appendingPathComponent("pseudo-dictation-tests", isDirectory: true)
+    .appendingPathComponent(UUID().uuidString, isDirectory: true)
+  let storage = DictationStorage(root: root)
+  try storage.prepare()
+  return storage
+}
+
+/// Writes the files a finished voice filter download leaves behind, without real model weights.
+func writeVoiceFilterFiles(_ storage: DictationStorage) throws {
+  let manager = FileManager.default
+  let directory = storage.modelDirectory(.voiceFilter)
+  for name in ModelNames.OfflineDiarizer.requiredModels {
+    let url = directory.appendingPathComponent(name)
+    if name.hasSuffix(".mlmodelc") {
+      try manager.createDirectory(at: url, withIntermediateDirectories: true)
+      try Data([1]).write(to: url.appendingPathComponent("coremldata.bin"))
+    } else {
+      try manager.createDirectory(at: directory, withIntermediateDirectories: true)
+      try Data([1]).write(to: url)
+    }
+  }
+}
+
+final class DictationStorageTests: XCTestCase {
+  private var storage: DictationStorage!
+
+  override func setUpWithError() throws {
+    storage = try makeTemporaryStorage()
+  }
+
+  override func tearDownWithError() throws {
+    try? FileManager.default.removeItem(at: storage.root)
+  }
+
+  func testModelsStayOutOfBackups() throws {
+    let values = try storage.root.resourceValues(forKeys: [.isExcludedFromBackupKey])
+    XCTAssertEqual(values.isExcludedFromBackup, true)
+  }
+
+  func testFilesWithoutTheMarkerAreNotInstalled() throws {
+    try writeVoiceFilterFiles(storage)
+    XCTAssertFalse(storage.isInstalled(.voiceFilter))
+    try storage.markInstalled(.voiceFilter)
+    XCTAssertTrue(storage.isInstalled(.voiceFilter))
+  }
+
+  func testTheMarkerWithoutFilesIsNotInstalled() throws {
+    try storage.markInstalled(.voiceFilter)
+    try storage.markInstalled(.fast)
+    XCTAssertFalse(storage.isInstalled(.voiceFilter))
+    XCTAssertFalse(storage.isInstalled(.fast))
+  }
+
+  func testAnInterruptedFileIsNotInstalled() throws {
+    try writeVoiceFilterFiles(storage)
+    try storage.markInstalled(.voiceFilter)
+    let partial = storage.modelDirectory(.voiceFilter)
+      .appendingPathComponent("Embedding.mlmodelc/weights/weight.bin.partial")
+    try FileManager.default.createDirectory(
+      at: partial.deletingLastPathComponent(), withIntermediateDirectories: true)
+    try Data([1]).write(to: partial)
+    XCTAssertFalse(storage.isInstalled(.voiceFilter))
+  }
+
+  func testAModelMissingItsCompiledDataIsNotInstalled() throws {
+    try writeVoiceFilterFiles(storage)
+    try storage.markInstalled(.voiceFilter)
+    try FileManager.default.removeItem(
+      at: storage.modelDirectory(.voiceFilter).appendingPathComponent("FBank.mlmodelc/coremldata.bin"))
+    XCTAssertFalse(storage.isInstalled(.voiceFilter))
+  }
+
+  func testDeleteRemovesFilesAndBothMarkers() throws {
+    let key = WarmKey(appBuild: "1(1)", osBuild: "os", processing: "2")
+    try writeVoiceFilterFiles(storage)
+    try storage.markInstalled(.voiceFilter)
+    try storage.markWarm(.voiceFilter, key: key)
+    XCTAssertGreaterThan(storage.sizeOnDisk(.voiceFilter), 0)
+
+    try storage.delete(.voiceFilter)
+
+    XCTAssertFalse(FileManager.default.fileExists(atPath: storage.modelDirectory(.voiceFilter).path))
+    XCTAssertEqual(storage.sizeOnDisk(.voiceFilter), 0)
+    XCTAssertFalse(storage.isWarm(.voiceFilter, key: key))
+    try writeVoiceFilterFiles(storage)
+    XCTAssertFalse(storage.isInstalled(.voiceFilter))
+  }
+
+  func testWarmOnlyForTheBuildOsAndProcessingItWasLoadedUnder() throws {
+    let key = WarmKey(appBuild: "1.0(7)", osBuild: "Version 18.0", processing: "2")
+    XCTAssertFalse(storage.isWarm(.fast, key: key))
+    try storage.markWarm(.fast, key: key)
+    XCTAssertTrue(storage.isWarm(.fast, key: key))
+    XCTAssertFalse(storage.isWarm(.balanced, key: key))
+    XCTAssertFalse(storage.isWarm(.fast, key: WarmKey(appBuild: "1.0(8)", osBuild: "Version 18.0", processing: "2")))
+    XCTAssertFalse(storage.isWarm(.fast, key: WarmKey(appBuild: "1.0(7)", osBuild: "Version 18.1", processing: "2")))
+    XCTAssertFalse(storage.isWarm(.fast, key: WarmKey(appBuild: "1.0(7)", osBuild: "Version 18.0", processing: "0")))
+  }
+
+  func testVoiceProfileRoundTrip() throws {
+    let profile = VoiceProfile(
+      extractor: VoiceProfile.currentExtractor, embedding: [0.6, 0.8], clipCount: 3,
+      createdAt: Date(timeIntervalSince1970: 1_700_000_000))
+    XCTAssertNil(storage.loadVoiceProfile())
+    try storage.saveVoiceProfile(profile)
+    XCTAssertEqual(storage.loadVoiceProfile(), profile)
+    try storage.deleteVoiceProfile()
+    XCTAssertNil(storage.loadVoiceProfile())
+    XCTAssertNoThrow(try storage.deleteVoiceProfile())
+  }
+
+  func testAProfileThisBuildCannotCompareReadsAsAbsent() throws {
+    let url = storage.root.appendingPathComponent("voice-profile.json")
+    try storage.saveVoiceProfile(
+      VoiceProfile(extractor: "another-extractor", embedding: [1, 0], clipCount: 3, createdAt: Date()))
+    XCTAssertNil(storage.loadVoiceProfile())
+    try storage.saveVoiceProfile(
+      VoiceProfile(
+        extractor: VoiceProfile.currentExtractor, embedding: [], clipCount: 3, createdAt: Date()))
+    XCTAssertNil(storage.loadVoiceProfile())
+    try Data("{ not json".utf8).write(to: url)
+    XCTAssertNil(storage.loadVoiceProfile())
+  }
+}