diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index f6141dc..73b87c8 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -50,7 +50,7 @@ jobs:
- name: Install XcodeGen
run: brew install xcodegen
- name: Generate project
- run: xcodegen generate
+ run: ./Scripts/generate-xcodeproj.sh
- name: Build
run: |
set -o pipefail
@@ -76,7 +76,7 @@ jobs:
- name: Install XcodeGen
run: brew install xcodegen
- name: Generate project
- run: xcodegen generate
+ run: ./Scripts/generate-xcodeproj.sh
- name: Run tests
run: |
set -o pipefail
diff --git a/CHANGELOG.md b/CHANGELOG.md
index dc4a73e..36a0ef4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,13 +7,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [Unreleased]
-### Changed
-- **iPhone only**: Set `TARGETED_DEVICE_FAMILY` to `"1"` for both `OSGKeyboard` and `OSGKeyboardExt` targets. Removed `UIRequiresFullScreen` and trimmed `UISupportedInterfaceOrientations` to Portrait only (iPad is no longer a supported device).
-- **Remove top divider line**: Deleted the 0.5 pt `palette.divider` overlay from `KeyboardRootView` — the subtle highlight gradient is retained; the hard separator line is gone.
-- **Keyboard preview always dark**: `KeyboardPreviewSheet` now injects `.environment(\.themePalette, Palette.dark)` alongside `.environment(\.colorScheme, .dark)` on `KeyboardPreviewStub`, so the preview palette is always the dark variant regardless of the app's active theme.
-- **Docs consistency**: README/README.zh now consistently describe the currently implemented capability set (`iOS 26+`, on-device `SpeechAnalyzer` + `DictationTranscriber`) with no deferred-ASR wording.
+## [0.2.0] - 2026-06-22
-## [0.1.2] - In Progress
+### Added
+- **Local engine with on-device Qwen models**: Optional Qwen3-ASR 0.6B speech recognition and Qwen3.5-0.8B text polish, fully offline after download.
+- **On-device model management**: Download, progress, delete, and readiness status in Settings; mirror auto-selection between ModelScope and Hugging Face with fallback.
+- **Engine picker**: Choose between local (on-device ASR + polish) and cloud (ASR + user-configured LLM polish).
+- **Flow session dictation**: TypeWhisper-style continuous capture in the host app with keyboard handoff via App Group.
+- **Open-source licenses** screen for bundled third-party components.
+
+### Changed
+- **Settings simplified**: Merged language and model sections; cloud mode always enables polish (removed off/transcribe mode picker).
+- **Keyboard UI**: Local/cloud engine badges replace the mode menu; shows model-not-downloaded guidance when the local stack is incomplete.
+- **iPhone only**: Set `TARGETED_DEVICE_FAMILY` to `"1"` for both targets; portrait-only orientations.
+- **Keyboard preview always dark**: Preview injects the dark palette regardless of app theme.
+- **Docs consistency**: README/README.zh aligned to the iOS 26+ capability set.
+
+### Fixed
+- **ModelScope download progress**: Progress now tracks byte counts instead of jumping to 50% after the first file.
+- **Light/Dark mode consistency** for shared button/card modifiers via `@Environment(\.themePalette)`.
+
+## [0.1.2] - 2026-06-20
### Fixed
- **Light/Dark mode consistency**: `cardSurface()`, `primaryButton()`, `secondaryButton()`, and `pillChip()` view modifiers in `Theme.swift` now use `ViewModifier` structs that read from `@Environment(\.themePalette)`. Previously they used hardcoded dark `Palette` constants, causing cards and buttons to always render in dark mode even when the main App was in light mode.
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index d2a7a9f..2f6bb68 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -29,7 +29,7 @@ Open an issue using the **Feature request** template. Briefly describe:
2. **Generate the project locally:**
```bash
brew install xcodegen swiftlint
- xcodegen generate
+ xcodegen generate # or: ./Scripts/generate-xcodeproj.sh
```
3. **Code style.** SwiftLint config lives in `.swiftlint.yml` — keep it green. We use Swift 6 strict concurrency, no `Sendable` shims where avoidable.
4. **Tests.** Add XCTest coverage in `OSGKeyboardTests/` for any non-trivial logic.
diff --git a/OSGKeyboard/OSGKeyboardApp.swift b/OSGKeyboard/OSGKeyboardApp.swift
index 181df28..6e610f3 100644
--- a/OSGKeyboard/OSGKeyboardApp.swift
+++ b/OSGKeyboard/OSGKeyboardApp.swift
@@ -14,6 +14,14 @@ struct OSGKeyboardApp: App {
init() {
MaterialIconsFont.registerIfNeeded()
+ // Register backend-specific ASR providers. The shared
+ // framework ships a built-in SpeechAnalyzer provider; we
+ // install the Qwen3-ASR provider here because linking
+ // `Qwen3ASR` pulls in mlx-swift, which the keyboard
+ // extension's `APPLICATION_EXTENSION_API_ONLY` build would
+ // refuse. Doing it in the host app's `init` keeps the heavy
+ // dependency localised.
+ ASRServiceFactory.providers[.qwen3ASR] = Qwen3ASRServiceProvider()
}
var body: some Scene {
@@ -29,8 +37,10 @@ struct OSGKeyboardApp: App {
AppGroupErrorView()
}
}
+ .environment(\.locale, config.uiLanguage.swiftUILocale)
.environmentObject(flowManager)
.onAppear {
+ FlowAppLifecycle.shared.setForeground(scenePhase == .active)
flowManager.setAppForeground(scenePhase == .active)
}
.onOpenURL { url in
@@ -51,11 +61,19 @@ struct OSGKeyboardApp: App {
)
}
.onChange(of: config.hasCompletedOnboarding) { _, done in
- if done { flowManager.autoStartIfNeeded() }
+ if done {
+ flowManager.autoStartIfNeeded()
+ if config.isLocalEngine {
+ OnDeviceModelWarmup.shared.warmUpIfNeeded()
+ }
+ }
}
.onChange(of: scenePhase) { _, phase in
- flowManager.setAppForeground(phase == .active)
+ flowManager.handleScenePhase(phase)
guard phase == .active, AppGroup.isAvailable, config.hasCompletedOnboarding else { return }
+ if config.isLocalEngine {
+ OnDeviceModelWarmup.shared.ensureReadyAfterBackground()
+ }
if flowManager.isActive {
flowManager.extendSession()
} else {
diff --git a/OSGKeyboard/Resources/PrivacyPolicy.html b/OSGKeyboard/Resources/PrivacyPolicy.html
new file mode 100644
index 0000000..d5629d3
--- /dev/null
+++ b/OSGKeyboard/Resources/PrivacyPolicy.html
@@ -0,0 +1,98 @@
+
+
+
+
+
+ OSGKeyboard Privacy Policy
+
+
+
+ 中文 · English
+
+
OSGKeyboard Privacy Policy
+
Last updated: June 19, 2026
+
OSGKeyboard is a custom iOS keyboard that turns your voice into text. This policy explains what data the app processes and how it is used.
+
+
What we collect
+
+ - Voice audio — captured only while you actively record. On-device mode transcribes locally with Apple’s speech APIs; raw audio is not uploaded by OSGKeyboard.
+ - Transcribed text — in Cloud polish mode, the final text (not audio) may be sent to the LLM provider you configure (e.g. OpenAI) for punctuation and formatting.
+ - API credentials — stored in the iOS Keychain on your device and shared only between the main app and keyboard extension via an App Group.
+ - App preferences — engine mode, language, and keyboard settings stored in App Group UserDefaults on your device.
+
+
+
What we do not collect
+
+ - We do not log or upload ordinary keystrokes you type with the keyboard.
+ - We do not operate analytics or advertising SDKs.
+ - We do not sell personal data.
+
+
+
Permissions
+
+ - Microphone — required for voice input and background voice sessions.
+ - Speech recognition — required for on-device transcription.
+ - Full Access — required so the keyboard can reach the microphone, read your API key, and communicate with the main app. Full Access does not grant us access to everything you type; we do not exfiltrate keystrokes.
+
+
+
Third parties
+
When you choose Cloud polish mode, transcribed text is sent to the API endpoint you configure. That provider’s privacy policy applies to those requests.
+
+
Data retention
+
Settings and API keys remain on your device until you delete the app or reset settings. Transcription results are passed to the host app you are typing in and are not stored long-term by OSGKeyboard.
+
Voice history — successful transcripts may be saved locally in the main app’s History tab for your convenience. This history stays on your device only, is never uploaded, and can be cleared at any time from History or by resetting settings.
+
+
Contact
+
Questions: open an issue at github.com/hkgood/OSGKeyboard.
+
+
+
+ OSGKeyboard 隐私政策
+ 更新日期:2026 年 6 月 19 日
+ OSGKeyboard 是一款 iOS 自定义键盘,可将语音转为文字。本政策说明应用处理哪些数据及用途。
+
+ 我们处理的数据
+
+ - 语音音频 — 仅在你主动录音时采集。本地模式在设备端通过 Apple 语音识别转写,OSGKeyboard 不会上传原始录音。
+ - 转写文字 — 云端润色模式下,最终文字(非音频)可能发送到你配置的 LLM 服务商以整理标点和格式。
+ - API 凭证 — 保存在设备 Keychain,仅通过 App Group 在主 App 与键盘扩展间共享。
+ - 应用偏好 — 引擎、语言等设置保存在设备 App Group 中。
+
+
+ 我们不收集的内容
+
+ - 我们不会记录或上传你平时在键盘上的击键内容。
+ - 我们不会集成广告或第三方分析 SDK。
+ - 我们不会出售个人数据。
+
+
+ 权限说明
+
+ - 麦克风 — 语音输入与后台语音会话所需。
+ - 语音识别 — 端侧转写所需。
+ - 完全访问 — 使键盘能使用麦克风、读取 API Key 并与主 App 通信。完全访问不代表我们会收集全部击键内容。
+
+
+ 第三方
+ 选择云端润色时,转写文字会发往你配置的 API,该服务商的隐私政策适用于相关请求。
+
+ 数据保留
+ 设置与 API Key 保留在设备上,直至卸载或重置。识别结果写入你正在使用的宿主 App,OSGKeyboard 不会长期存储。
+ 语音历史 — 成功的转写可能保存在主 App「历史」页,仅供本机查看,不会上传,可随时在历史页清空或通过重置设置清除。
+
+ 联系
+ 问题反馈:github.com/hkgood/OSGKeyboard
+
+
diff --git a/OSGKeyboard/Services/AppPermissions.swift b/OSGKeyboard/Services/AppPermissions.swift
index cf46743..8aa5f2d 100644
--- a/OSGKeyboard/Services/AppPermissions.swift
+++ b/OSGKeyboard/Services/AppPermissions.swift
@@ -72,6 +72,7 @@ enum AppPermissions {
}
}
+ @MainActor
static func openSystemSettings() {
guard let url = URL(string: UIApplication.openSettingsURLString) else { return }
UIApplication.shared.open(url)
@@ -82,12 +83,12 @@ enum AppPermissions {
let micMissing = micStatus != .granted
let speechMissing = speechStatus != .granted
if micMissing && speechMissing {
- return NSLocalizedString("home.setup.permission.both", comment: "")
+ return AppL10n.string("home.setup.permission.both")
}
if micMissing {
- return NSLocalizedString("home.setup.permission.mic", comment: "")
+ return AppL10n.string("home.setup.permission.mic")
}
- return NSLocalizedString("home.setup.permission.speech", comment: "")
+ return AppL10n.string("home.setup.permission.speech")
}
/// True when at least one permission can still be requested in-app.
diff --git a/OSGKeyboard/Services/FlowDiagnostics.swift b/OSGKeyboard/Services/FlowDiagnostics.swift
new file mode 100644
index 0000000..3861257
--- /dev/null
+++ b/OSGKeyboard/Services/FlowDiagnostics.swift
@@ -0,0 +1,22 @@
+// FlowDiagnostics.swift
+// OSGKeyboard · Main App
+//
+// Structured logging for the Flow dictation pipeline. Visible in Xcode
+// console (DEBUG) and Console.app via `subsystem: com.osgkeyboard.ios`.
+
+import Foundation
+import os
+
+enum FlowDiagnostics {
+ private static let logger = Logger(
+ subsystem: "com.osgkeyboard.ios",
+ category: "Flow"
+ )
+
+ static func log(_ message: String) {
+ logger.info("\(message, privacy: .public)")
+ #if DEBUG
+ print("🌊[OSGFlow] \(message)")
+ #endif
+ }
+}
diff --git a/OSGKeyboard/Services/FlowSessionManager.swift b/OSGKeyboard/Services/FlowSessionManager.swift
index 6a278a2..3880147 100644
--- a/OSGKeyboard/Services/FlowSessionManager.swift
+++ b/OSGKeyboard/Services/FlowSessionManager.swift
@@ -3,12 +3,14 @@
//
// Session Owner for TypeWhisper-style Flow dictation: continuous
// `.playAndRecord` capture for the whole session, utterance gating for
-// ASR, optional LLM polish, and App Group result delivery.
+// ASR and cloud LLM polish, with App Group result delivery.
import Foundation
import AVFoundation
import Speech
import OSGKeyboardShared
+import UIKit
+import SwiftUI
@MainActor
final class FlowSessionManager: ObservableObject {
@@ -19,9 +21,22 @@ final class FlowSessionManager: ObservableObject {
@Published private(set) var sessionWarning: String?
private let capture = FlowContinuousCapture()
- private let asr: ASRService = ASRServiceFactory.make()
- private let polisher = PolishingService()
private let store = AppGroupStore()
+ /// Cloud-engine polish only; local engine delivers raw ASR text.
+ private var polisher: PolishingService {
+ PolishingService()
+ }
+ /// Cached ASR instance shared with `OnDeviceModelWarmup`.
+ private var sessionASR: ASRService?
+ private var asr: ASRService {
+ if let sessionASR { return sessionASR }
+ let service = OnDeviceModelWarmup.shared.asrService(
+ engineMode: store.engineMode,
+ localBackend: store.localASRBackend
+ )
+ sessionASR = service
+ return service
+ }
private var pollingTask: Task?
private var heartbeatTask: Task?
@@ -33,10 +48,13 @@ final class FlowSessionManager: ObservableObject {
private var isUtteranceProcessing = false
private var finalizeTask: Task?
private var asrTask: Task?
+ private var chunkedPipeline: ChunkedUtterancePipeline?
private var currentPartial = ""
private var lastFinal = ""
+ private var chunkWarnings: [String] = []
/// True while the host app scene is `.active` — drives foreground renewal.
private var isAppForeground = false
+ private var backgroundTaskID: UIBackgroundTaskIdentifier = .invalid
init() {
Task { @MainActor [weak self] in
@@ -126,6 +144,9 @@ final class FlowSessionManager: ObservableObject {
startLevelPublishing()
scheduleExpiry(after: remaining)
+ OnDeviceModelWarmup.shared.warmUpIfNeeded()
+ bindSessionASR()
+
debug("Flow session restored (\(Int(remaining))s remaining)")
}
@@ -149,13 +170,17 @@ final class FlowSessionManager: ObservableObject {
if isUtteranceRecording || isUtteranceProcessing {
capture.cancelUtterance()
asrTask?.cancel()
+ Task { await chunkedPipeline?.cancel() }
asr.cancel()
}
asrTask = nil
+ chunkedPipeline = nil
isUtteranceRecording = false
isUtteranceProcessing = false
capture.stop()
+ endBackgroundKeepAlive()
+ sessionASR = nil
FlowSessionBridge.markSessionInactive()
FlowSessionDarwin.postSessionChanged()
isActive = false
@@ -179,6 +204,81 @@ final class FlowSessionManager: ObservableObject {
}
}
+ /// Full scene lifecycle — keeps Flow + ASR alive across app switches.
+ func handleScenePhase(_ phase: ScenePhase) {
+ switch phase {
+ case .active:
+ FlowAppLifecycle.shared.setForeground(true)
+ setAppForeground(true)
+ resumeAfterForeground()
+ case .inactive:
+ writeHeartbeatIfActive()
+ case .background:
+ FlowAppLifecycle.shared.setForeground(false)
+ setAppForeground(false)
+ beginBackgroundKeepAlive()
+ @unknown default:
+ break
+ }
+ }
+
+ private func writeHeartbeatIfActive() {
+ guard isActive else { return }
+ FlowSessionBridge.writeHeartbeat()
+ }
+
+ private func beginBackgroundKeepAlive() {
+ guard isActive else { return }
+ FlowSessionBridge.writeHeartbeat()
+
+ guard backgroundTaskID == .invalid else { return }
+ backgroundTaskID = UIApplication.shared.beginBackgroundTask { [weak self] in
+ self?.endBackgroundKeepAlive()
+ }
+ debug("background keep-alive started")
+ }
+
+ private func endBackgroundKeepAlive() {
+ guard backgroundTaskID != .invalid else { return }
+ UIApplication.shared.endBackgroundTask(backgroundTaskID)
+ backgroundTaskID = .invalid
+ debug("background keep-alive ended")
+ }
+
+ private func resumeAfterForeground() {
+ guard isActive else {
+ endBackgroundKeepAlive()
+ return
+ }
+
+ FlowSessionBridge.writeHeartbeat()
+ endBackgroundKeepAlive()
+
+ Task { @MainActor [weak self] in
+ await self?.reactivateCaptureIfNeeded()
+ OnDeviceModelWarmup.shared.ensureReadyAfterBackground()
+ self?.bindSessionASR()
+ }
+ }
+
+ private func reactivateCaptureIfNeeded() async {
+ guard isActive else { return }
+
+ if capture.running {
+ capture.reassertIfRunning()
+ return
+ }
+
+ do {
+ try capture.start()
+ debug("capture restarted after foreground")
+ } catch {
+ let message = (error as? LocalizedError)?.errorDescription ?? error.localizedDescription
+ sessionWarning = message
+ debug("capture restart failed: \(message)")
+ }
+ }
+
/// Extend the session before it expires while the host app stays in foreground.
private func renewSessionIfNeededWhileForeground() {
guard isActive, isAppForeground else { return }
@@ -220,14 +320,24 @@ final class FlowSessionManager: ObservableObject {
startLevelPublishing()
scheduleExpiry(after: duration)
+ OnDeviceModelWarmup.shared.warmUpIfNeeded()
+ bindSessionASR()
+
debug("Flow session started (\(Int(duration))s), continuous capture running")
}
+ private func bindSessionASR() {
+ sessionASR = OnDeviceModelWarmup.shared.asrService(
+ engineMode: store.engineMode,
+ localBackend: store.localASRBackend
+ )
+ }
+
private func permissionWarningMessage() -> String {
if AppPermissions.micStatus != .granted {
- return NSLocalizedString("flow.error.micRequired", comment: "")
+ return AppL10n.string("flow.error.micRequired")
}
- return NSLocalizedString("flow.error.speechRequired", comment: "")
+ return AppL10n.string("flow.error.speechRequired")
}
// MARK: - Polling
@@ -259,19 +369,20 @@ final class FlowSessionManager: ObservableObject {
private func beginUtterance() {
guard capture.running else {
- failUtterance(message: NSLocalizedString("flow.error.audioUnavailable", comment: ""))
+ failUtterance(message: AppL10n.string("flow.error.audioUnavailable"))
return
}
- // Mirror `LiveDictationController.start`: only begin when the previous
- // utterance fully finished. Never cancel an in-flight analyzer here —
- // that was the source of intermittent CancellationError / noSpeech.
guard !isUtteranceProcessing else {
debug("beginUtterance ignored — previous utterance still processing")
return
}
+ // Honor engine / ASR backend changes without restarting the session.
+ bindSessionASR()
+
currentPartial = ""
lastFinal = ""
+ chunkWarnings = []
let localeId = store.localeId
FlowSessionBridge.setTranscriptionLanguage(localeId)
@@ -279,28 +390,42 @@ final class FlowSessionManager: ObservableObject {
let locale = SpeechLocaleResolver.resolve(localeId)
let stream = capture.beginUtterance()
- let events = asr.transcribe(stream: stream, locale: locale)
+ let pipeline = ChunkedUtterancePipeline(asr: asr, locale: locale)
+ chunkedPipeline = pipeline
isUtteranceRecording = true
+ FlowDiagnostics.log(
+ "beginUtterance engine=\(store.engineMode) asr=\(store.localASRBackend.rawValue) " +
+ "modelsInMemory=\(OnDeviceModelStatus.modelsLoadedInMemory()) " +
+ "asrType=\(type(of: asr)) pipelined=true max=\(Int(FlowSessionKeys.maxUtteranceDuration))s"
+ )
- asrTask = Task { @MainActor [weak self] in
- guard let self else { return }
- for await event in events {
- switch event {
- case .capability:
- break
- case .partial(let text):
- self.currentPartial = text
- case .final(let text):
- self.lastFinal = text.trimmingCharacters(in: .whitespacesAndNewlines)
- self.currentPartial = ""
- case .error(let message):
- self.debug("asr error: \(message)")
- if self.isUtteranceRecording {
- self.failUtterance(message: message)
- } else if self.isUtteranceProcessing {
- self.finishProcessing(withError: message)
+ asrTask = Task.detached(priority: .userInitiated) { [weak manager = self] in
+ let outcome = await pipeline.transcribe(stream: stream) { partial in
+ Task { @MainActor in
+ manager?.currentPartial = partial
+ }
+ }
+ await MainActor.run {
+ guard let manager else { return }
+ FlowDiagnostics.log(
+ "chunkedASR finished partialLen=\(manager.currentPartial.count) " +
+ "finalPending=\(manager.lastFinal.isEmpty)"
+ )
+ switch outcome {
+ case .success(let success):
+ manager.lastFinal = success.text
+ manager.chunkWarnings = success.chunkWarnings
+ manager.currentPartial = ""
+ case .failure(let message):
+ manager.debug("asr error: \(message)")
+ if manager.isUtteranceRecording {
+ manager.failUtterance(message: message)
+ } else if manager.isUtteranceProcessing {
+ manager.finishProcessing(withError: message)
}
+ case .cancelled:
+ break
}
}
}
@@ -335,10 +460,13 @@ final class FlowSessionManager: ObservableObject {
finalizeTask?.cancel()
finalizeTask = nil
asrTask?.cancel()
+ Task { await chunkedPipeline?.cancel() }
+ chunkedPipeline = nil
asr.cancel()
capture.cancelUtterance()
currentPartial = ""
lastFinal = ""
+ chunkWarnings = []
FlowSessionBridge.setRecordingState(.idle)
debug("utterance aborted")
}
@@ -349,10 +477,13 @@ final class FlowSessionManager: ObservableObject {
finalizeTask?.cancel()
finalizeTask = nil
asrTask?.cancel()
+ Task { await chunkedPipeline?.cancel() }
+ chunkedPipeline = nil
asr.cancel()
capture.cancelUtterance()
currentPartial = ""
lastFinal = ""
+ chunkWarnings = []
FlowSessionBridge.storeTranscriptionError(message)
FlowSessionBridge.setRecordingState(.idle)
debug("utterance failed: \(message)")
@@ -362,26 +493,43 @@ final class FlowSessionManager: ObservableObject {
isUtteranceProcessing = false
finalizeTask?.cancel()
finalizeTask = nil
+ chunkedPipeline = nil
currentPartial = ""
lastFinal = ""
+ chunkWarnings = []
FlowSessionBridge.storeTranscriptionError(message)
FlowSessionBridge.setRecordingState(.idle)
debug("utterance processing failed: \(message)")
}
private func finalizeUtterance() async {
+ let pipelineStarted = Date()
defer {
isUtteranceProcessing = false
FlowSessionBridge.setRecordingState(.idle)
}
- let deadline = Date().addingTimeInterval(30)
- while Date() < deadline {
+ let asrWait = asrWaitTimeout()
+ FlowDiagnostics.log(
+ "finalize start asrWait=\(Int(asrWait))s engine=\(store.engineMode) " +
+ "backend=\(store.localASRBackend.rawValue)"
+ )
+
+ let asrDeadline = Date().addingTimeInterval(asrWait)
+ while Date() < asrDeadline {
if !lastFinal.isEmpty { break }
if asrTask?.isCancelled == true { break }
try? await Task.sleep(nanoseconds: 100_000_000)
}
+ if lastFinal.isEmpty, let asrTask {
+ FlowDiagnostics.log("ASR wait elapsed — awaiting asrTask completion")
+ _ = await asrTask.value
+ }
+
+ let asrElapsed = Date().timeIntervalSince(pipelineStarted)
+ FlowDiagnostics.log("ASR phase done in \(String(format: "%.1f", asrElapsed))s finalLen=\(lastFinal.count)")
+
var text = lastFinal.trimmingCharacters(in: .whitespacesAndNewlines)
if text.isEmpty {
text = currentPartial.trimmingCharacters(in: .whitespacesAndNewlines)
@@ -390,36 +538,72 @@ final class FlowSessionManager: ObservableObject {
let key = (asrTask?.isCancelled == true)
? "flow.error.recognitionInterrupted"
: "flow.error.noSpeech"
+ FlowDiagnostics.log("finalize failed: empty transcript after \(String(format: "%.1f", asrElapsed))s")
FlowSessionBridge.storeTranscriptionError(
- NSLocalizedString(key, comment: "")
+ AppL10n.string(key)
)
return
}
let engineMode = store.engineMode
- let modeId = store.modeId
- let shouldPolish = engineMode != "local" && modeId == "polish"
+
+ if engineMode == "local" {
+ let warning = Self.chunkWarningMessage(chunkWarnings)
+ FlowSessionBridge.storeTranscriptionResult(text, polishWarning: warning)
+ FlowDiagnostics.log(
+ "finalize ASR-only total=\(String(format: "%.1f", Date().timeIntervalSince(pipelineStarted)))s " +
+ "len=\(text.count)"
+ )
+ SpeechHistoryStore.shared.append(text: text, engineMode: engineMode)
+ currentPartial = ""
+ lastFinal = ""
+ chunkWarnings = []
+ debug("utterance finalized length=\(text.count)")
+ return
+ }
var delivered = text
- if shouldPolish {
- do {
- let polished = try await polisher.polish(text)
- delivered = polished
- FlowSessionBridge.storeTranscriptionResult(polished)
- } catch {
- FlowSessionBridge.storeTranscriptionResult(text)
- }
- } else {
- FlowSessionBridge.storeTranscriptionResult(text)
+ let chunkNote = Self.chunkWarningMessage(chunkWarnings)
+ let polishStarted = Date()
+ do {
+ let polished = try await polisher.polish(text)
+ delivered = polished
+ FlowSessionBridge.storeTranscriptionResult(polished, polishWarning: chunkNote)
+ FlowDiagnostics.log(
+ "polish done in \(String(format: "%.1f", Date().timeIntervalSince(polishStarted)))s " +
+ "total=\(String(format: "%.1f", Date().timeIntervalSince(pipelineStarted)))s"
+ )
+ } catch {
+ FlowDiagnostics.log(
+ "polish failed after \(String(format: "%.1f", Date().timeIntervalSince(polishStarted)))s: " +
+ "\(error.localizedDescription)"
+ )
+ FlowSessionBridge.storeTranscriptionResult(text, polishWarning: chunkNote)
}
SpeechHistoryStore.shared.append(text: delivered, engineMode: engineMode)
currentPartial = ""
lastFinal = ""
+ chunkWarnings = []
+ chunkedPipeline = nil
debug("utterance finalized length=\(text.count)")
}
+ private static func chunkWarningMessage(_ warnings: [String]) -> String? {
+ guard !warnings.isEmpty else { return nil }
+ return warnings.joined(separator: "\n")
+ }
+
+ private func asrWaitTimeout() -> TimeInterval {
+ if store.engineMode == "local" {
+ return store.localASRBackend == .qwen3ASR
+ ? FlowSessionKeys.localQwen3ASRWaitTimeout
+ : FlowSessionKeys.localASRWaitTimeout
+ }
+ return FlowSessionKeys.cloudASRWaitTimeout
+ }
+
// MARK: - Level publishing (main thread only)
private func startLevelPublishing() {
@@ -461,8 +645,6 @@ final class FlowSessionManager: ObservableObject {
}
private func debug(_ message: String) {
- #if DEBUG
- print("🌊[FlowSession] \(message)")
- #endif
+ FlowDiagnostics.log(message)
}
}
diff --git a/OSGKeyboard/Services/LegalLinks.swift b/OSGKeyboard/Services/LegalLinks.swift
index 0c9136e..2f751cb 100644
--- a/OSGKeyboard/Services/LegalLinks.swift
+++ b/OSGKeyboard/Services/LegalLinks.swift
@@ -4,7 +4,9 @@
import Foundation
enum LegalLinks {
- /// Public privacy policy (GitHub Pages).
+ static let repositoryURL = URL(string: "https://github.com/hkgood/OSGKeyboard")!
+
+ /// Public privacy policy (GitHub Pages). Also bundled in-app as PrivacyPolicy.html.
static var privacyPolicyURL: URL? {
URL(string: "https://hkgood.github.io/OSGKeyboard/privacy/")
}
diff --git a/OSGKeyboard/Services/ModelDownloadSourcePicker.swift b/OSGKeyboard/Services/ModelDownloadSourcePicker.swift
new file mode 100644
index 0000000..9cecae5
--- /dev/null
+++ b/OSGKeyboard/Services/ModelDownloadSourcePicker.swift
@@ -0,0 +1,126 @@
+// ModelDownloadSourcePicker.swift
+// OSGKeyboard · Main App
+//
+// Picks ModelScope vs Hugging Face by probing both mirrors on the
+// user's current network. Result is cached briefly so consecutive
+// downloads in one session don't re-probe.
+
+import Foundation
+import os
+
+enum ModelDownloadSourcePicker {
+
+ private struct CacheState {
+ var source: ModelDownloadSource?
+ var expiresAt: Date?
+ }
+
+ private static let lock = OSAllocatedUnfairLock(initialState: CacheState())
+ private static let cacheTTL: TimeInterval = 300
+
+ /// Resolves the fastest reachable mirror for the current network.
+ static func resolve() async -> ModelDownloadSource {
+ if let cached = cachedValue() { return cached }
+
+ let winner = await probeFastest() ?? defaultHeuristic()
+ storeCache(winner)
+ return winner
+ }
+
+ /// Alternate mirror — used when the first download attempt fails.
+ static func alternate(to source: ModelDownloadSource) -> ModelDownloadSource {
+ switch source {
+ case .modelScope: return .huggingface
+ case .huggingface: return .modelScope
+ }
+ }
+
+ // MARK: - Probe
+
+ private static func probeFastest() async -> ModelDownloadSource? {
+ await withTaskGroup(of: (ModelDownloadSource, TimeInterval)?.self) { group in
+ for source in ModelDownloadSource.allCases {
+ group.addTask {
+ guard let latency = await probeLatency(for: source) else { return nil }
+ return (source, latency)
+ }
+ }
+
+ var best: (ModelDownloadSource, TimeInterval)?
+ for await candidate in group {
+ guard let candidate else { continue }
+ if best == nil || candidate.1 < best!.1 {
+ best = candidate
+ }
+ }
+ return best?.0
+ }
+ }
+
+ private static func probeLatency(for source: ModelDownloadSource) async -> TimeInterval? {
+ guard let url = probeURL(for: source) else { return nil }
+
+ var request = URLRequest(url: url)
+ request.httpMethod = "HEAD"
+ request.timeoutInterval = 4
+ request.cachePolicy = .reloadIgnoringLocalCacheData
+
+ let started = CFAbsoluteTimeGetCurrent()
+ do {
+ let (_, response) = try await URLSession.shared.data(for: request)
+ guard let http = response as? HTTPURLResponse else { return nil }
+ guard (200...399).contains(http.statusCode) else { return nil }
+ return CFAbsoluteTimeGetCurrent() - started
+ } catch {
+ // Some hosts reject HEAD — retry with a tiny GET.
+ var get = URLRequest(url: url)
+ get.httpMethod = "GET"
+ get.timeoutInterval = 4
+ get.cachePolicy = .reloadIgnoringLocalCacheData
+ do {
+ let (_, response) = try await URLSession.shared.data(for: get)
+ guard let http = response as? HTTPURLResponse else { return nil }
+ guard (200...399).contains(http.statusCode) else { return nil }
+ return CFAbsoluteTimeGetCurrent() - started
+ } catch {
+ return nil
+ }
+ }
+ }
+
+ private static func probeURL(for source: ModelDownloadSource) -> URL? {
+ switch source {
+ case .modelScope:
+ return URL(string: "https://modelscope.cn")
+ case .huggingface:
+ return URL(string: "https://huggingface.co")
+ }
+ }
+
+ /// When both probes fail (offline, captive portal, etc.).
+ private static func defaultHeuristic() -> ModelDownloadSource {
+ if Locale.current.region?.identifier == "CN" { return .modelScope }
+ if TimeZone.current.identifier.hasPrefix("Asia/Shanghai") { return .modelScope }
+ return .huggingface
+ }
+
+ // MARK: - Cache
+
+ private static func cachedValue() -> ModelDownloadSource? {
+ lock.withLock { state in
+ guard let source = state.source,
+ let expiresAt = state.expiresAt,
+ expiresAt > Date() else {
+ return nil
+ }
+ return source
+ }
+ }
+
+ private static func storeCache(_ source: ModelDownloadSource) {
+ lock.withLock { state in
+ state.source = source
+ state.expiresAt = Date().addingTimeInterval(cacheTTL)
+ }
+ }
+}
diff --git a/OSGKeyboard/Services/ModelManager.swift b/OSGKeyboard/Services/ModelManager.swift
new file mode 100644
index 0000000..aa891c0
--- /dev/null
+++ b/OSGKeyboard/Services/ModelManager.swift
@@ -0,0 +1,492 @@
+// ModelManager.swift
+// OSGKeyboard · Main App
+//
+// Owns the lifecycle of on-device ML models that back the local
+// ASR backend (Qwen3-ASR-0.6B CoreML, ~1.6 GB).
+//
+// `runDownload` fetches CoreML bundles + tokenizer files — it does
+// not load models into memory (warm-up happens in `OnDeviceModelWarmup`).
+// Weights land under `~/Library/Caches/qwen3-speech/` using the Hub
+// layout from `HuggingFaceDownloader`.
+//
+// Why this lives in the host app: the ASR model is loaded via
+// soniqo/speech-swift, which is only linked into the main App
+// target (Qwen3Speech pulls mlx-swift as a transitive dependency).
+//
+// Mirror selection: resolved automatically at download time via
+// `ModelDownloadSourcePicker` (latency probe + locale fallback).
+
+import Foundation
+import SwiftUI
+import OSGKeyboardShared
+import Qwen3ASR
+
+private enum Qwen3CoreMLDownloadArtifacts {
+ static let coreMLBundleGlobs = [
+ "encoder.mlmodelc/**",
+ "embedding.mlmodelc/**",
+ "decoder_part1.mlmodelc/**",
+ "decoder_part2.mlmodelc/**",
+ "config.json",
+ ]
+
+ static let tokenizerFiles = [
+ "vocab.json",
+ "merges.txt",
+ "tokenizer_config.json",
+ ]
+}
+
+/// Where on-device model weights are downloaded from.
+enum ModelDownloadSource: String, CaseIterable, Identifiable, Sendable {
+ case huggingface
+ case modelScope
+
+ var id: String { rawValue }
+
+ var registry: ModelRegistry {
+ switch self {
+ case .huggingface: return .huggingFace()
+ case .modelScope: return .modelScope()
+ }
+ }
+
+ /// Host shown in error messages.
+ var hostLabel: String {
+ switch self {
+ case .huggingface: return "huggingface.co"
+ case .modelScope: return "modelscope.cn"
+ }
+ }
+}
+
+enum ModelDownloadState: Equatable, Sendable {
+ case notDownloaded
+ case downloading(progress: Double)
+ case downloaded
+ case failed(String)
+
+ var isTerminal: Bool {
+ switch self {
+ case .downloaded, .failed: return true
+ case .notDownloaded, .downloading: return false
+ }
+ }
+
+ var downloadProgress: Double? {
+ if case .downloading(let progress) = self { return progress }
+ return nil
+ }
+}
+
+/// Per-model state tracked by `ModelManager`. The manager keeps a
+/// dictionary of these and re-emits it on the main actor whenever
+/// any field changes.
+struct ModelState: Equatable, Sendable {
+ var download: ModelDownloadState
+ var lastError: String?
+}
+
+/// Observable holder that the Settings UI binds to. All mutating
+/// methods dispatch onto the main actor so SwiftUI views can
+/// observe without ceremony.
+@MainActor
+final class ModelManager: ObservableObject {
+
+ static let shared = ModelManager()
+
+ @Published private(set) var states: [OnDeviceModel: ModelState] = [:]
+ @Published private(set) var activeDownloads: Set = []
+
+ private var downloadTasks: [OnDeviceModel: Task] = [:]
+
+ init() {
+ for model in OnDeviceModel.allCases {
+ states[model] = ModelState(download: .notDownloaded, lastError: nil)
+ }
+ refreshAll()
+ }
+
+ // MARK: - Queries
+
+ /// Synchronous check on whether the model is already on disk.
+ /// Used by the UI to decide whether to show "Download" or
+ /// "Delete". Doesn't touch the network.
+ func isDownloaded(_ model: OnDeviceModel) -> Bool {
+ Self.weightsOnDisk(for: model)
+ }
+
+ /// Disk-only probe safe to call from background ASR tasks.
+ /// Returns `false` until the user downloads via Settings.
+ nonisolated static func weightsOnDisk(for model: OnDeviceModel) -> Bool {
+ existingCacheDirectory(for: model) != nil
+ }
+
+ /// Approximate on-disk bytes used by the model directory. Used
+ /// by the Settings "Storage" badge.
+ func onDiskBytes(_ model: OnDeviceModel) -> Int64 {
+ guard let dir = Self.existingCacheDirectory(for: model) else { return 0 }
+ guard let enumerator = FileManager.default.enumerator(
+ at: dir,
+ includingPropertiesForKeys: [.totalFileAllocatedSizeKey, .isRegularFileKey]
+ ) else { return 0 }
+ var total: Int64 = 0
+ for case let url as URL in enumerator {
+ let values = try? url.resourceValues(forKeys: [.totalFileAllocatedSizeKey, .isRegularFileKey])
+ if values?.isRegularFile == true {
+ total += Int64(values?.totalFileAllocatedSize ?? 0)
+ }
+ }
+ return total
+ }
+
+ // MARK: - Mutations
+
+ /// Triggers a background download of the model. The call returns
+ /// immediately; observe `states[model].download` for progress.
+ /// Calling this while a download is in progress is a no-op.
+ func startDownload(_ model: OnDeviceModel) {
+ if activeDownloads.contains(model) { return }
+ if isDownloaded(model) {
+ states[model]?.download = .downloaded
+ return
+ }
+ activeDownloads.insert(model)
+ states[model] = ModelState(download: .downloading(progress: 0), lastError: nil)
+ publishStatusToAppGroup()
+
+ let task = Task.detached(priority: .userInitiated) { [weak self] in
+ guard let self else { return }
+ let primary = await ModelDownloadSourcePicker.resolve()
+ do {
+ try await self.runDownload(model, registry: primary.registry)
+ } catch is CancellationError {
+ await self.finishDownloadCancelled(model)
+ } catch {
+ let fallback = ModelDownloadSourcePicker.alternate(to: primary)
+ await self.reportDownloadProgress(model, fraction: 0, monotonic: false)
+ do {
+ try await self.runDownload(model, registry: fallback.registry)
+ } catch is CancellationError {
+ await self.finishDownloadCancelled(model)
+ } catch {
+ await self.finishDownloadFailed(model, error: error)
+ }
+ }
+ }
+ downloadTasks[model] = task
+ }
+
+ func cancelDownload(_ model: OnDeviceModel) {
+ downloadTasks[model]?.cancel()
+ downloadTasks[model] = nil
+ activeDownloads.remove(model)
+ states[model] = ModelState(download: .notDownloaded, lastError: nil)
+ publishStatusToAppGroup()
+ }
+
+ func deleteModel(_ model: OnDeviceModel) {
+ for dir in Self.candidateCacheDirectories(for: model) {
+ try? FileManager.default.removeItem(at: dir)
+ }
+ states[model] = ModelState(download: .notDownloaded, lastError: nil)
+ publishStatusToAppGroup()
+ OnDeviceModelWarmup.shared.invalidate()
+ }
+
+ /// Cheap refresh that re-reads the on-disk state for every
+ /// tracked model. Called from `init` and after a successful
+ /// download so the Settings row updates from "Downloading…" to
+ /// "Downloaded · 1.4 GB" without needing a separate notifier.
+ func refreshAll() {
+ for model in OnDeviceModel.allCases {
+ if activeDownloads.contains(model) { continue }
+ if isDownloaded(model) {
+ states[model] = ModelState(download: .downloaded, lastError: nil)
+ } else if case .failed = states[model]?.download {
+ // Preserve any existing failure message so the UI
+ // can show "Download failed: " instead of
+ // resetting it back to "Not downloaded" every
+ // refresh.
+ continue
+ } else {
+ states[model] = ModelState(download: .notDownloaded, lastError: states[model]?.lastError)
+ }
+ }
+ publishStatusToAppGroup()
+ }
+
+ /// Mirror disk/download state into the App Group for the keyboard
+ /// extension, which cannot probe the host app's Caches folder.
+ private func publishStatusToAppGroup() {
+ for model in OnDeviceModel.allCases {
+ let downloaded: Bool
+ let progress: Double?
+ switch states[model]?.download {
+ case .downloaded:
+ downloaded = true
+ progress = nil
+ case .downloading(let fraction):
+ downloaded = false
+ progress = fraction
+ case .failed, .notDownloaded, .none:
+ downloaded = isDownloaded(model)
+ progress = nil
+ }
+ OnDeviceModelStatus.setDownloaded(downloaded, for: model)
+ OnDeviceModelStatus.setProgress(progress, for: model)
+ }
+ scheduleWarmupIfNeeded()
+ }
+
+ private func scheduleWarmupIfNeeded() {
+ let config = ProviderConfig.shared
+ guard config.isLocalEngine else {
+ OnDeviceModelWarmup.shared.invalidate()
+ return
+ }
+ if OnDeviceModelStatus.isLocalStackReady(asrBackend: config.localASRBackend) {
+ // Do not force-restart an in-flight warm-up — `publishStatusToAppGroup`
+ // runs on download progress ticks and would otherwise cancel load
+ // mid-flight, leaving the UI stuck on "warming".
+ OnDeviceModelWarmup.shared.warmUpIfNeeded()
+ } else {
+ OnDeviceModelWarmup.shared.invalidate()
+ }
+ }
+
+ // MARK: - Internals
+
+ /// Runs off the main actor; updates `@Published` state via `MainActor.run`.
+ /// Throws on failure so `startDownload` can fall back to the alternate mirror.
+ nonisolated private func runDownload(_ model: OnDeviceModel, registry: ModelRegistry) async throws {
+ switch model {
+ case .qwen3ASR:
+ try await Self.downloadQwen3CoreMLWeights(
+ model: model,
+ registry: registry,
+ progressHandler: { @Sendable [weak self] fraction, _ in
+ Task { @MainActor [weak self] in
+ guard let self else { return }
+ self.reportDownloadProgress(model, fraction: fraction)
+ }
+ }
+ )
+ }
+
+ await MainActor.run { [weak self] in
+ guard let self else { return }
+ self.activeDownloads.remove(model)
+ self.downloadTasks[model] = nil
+ self.states[model] = ModelState(download: .downloaded, lastError: nil)
+ self.publishStatusToAppGroup()
+ let config = ProviderConfig.shared
+ if config.isLocalEngine,
+ OnDeviceModelStatus.isLocalStackReady(asrBackend: config.localASRBackend) {
+ // Retry warm-up after a prior load failure once weights land on disk.
+ OnDeviceModelWarmup.shared.warmUpIfNeeded(force: true)
+ }
+ }
+ }
+
+ nonisolated private func finishDownloadCancelled(_ model: OnDeviceModel) async {
+ await MainActor.run { [weak self] in
+ guard let self else { return }
+ self.activeDownloads.remove(model)
+ self.downloadTasks[model] = nil
+ self.states[model] = ModelState(download: .notDownloaded, lastError: nil)
+ self.publishStatusToAppGroup()
+ }
+ }
+
+ nonisolated private func finishDownloadFailed(_ model: OnDeviceModel, error: Error) async {
+ let message = Self.userFacingDownloadError(error)
+ await MainActor.run { [weak self] in
+ guard let self else { return }
+ self.activeDownloads.remove(model)
+ self.downloadTasks[model] = nil
+ self.states[model] = ModelState(download: .failed(message), lastError: message)
+ self.publishStatusToAppGroup()
+ }
+ }
+
+ /// Updates UI progress. By default keeps the bar monotonic so brief
+ /// per-file jumps inside the downloader never move backwards.
+ private func reportDownloadProgress(
+ _ model: OnDeviceModel,
+ fraction: Double,
+ monotonic: Bool = true
+ ) {
+ let clamped = min(max(fraction, 0), 1)
+ let previous = states[model]?.download.downloadProgress ?? 0
+ let value = monotonic ? max(previous, clamped) : clamped
+ states[model] = ModelState(download: .downloading(progress: value), lastError: nil)
+ publishStatusToAppGroup()
+ }
+
+ /// Short, user-readable download failure text for Settings UI.
+ nonisolated private static func userFacingDownloadError(_ error: Error) -> String {
+ let raw = (error as? LocalizedError)?.errorDescription ?? error.localizedDescription
+ if raw.localizedCaseInsensitiveContains("metadata") {
+ return AppL10n.string("settings.models.error.metadata")
+ }
+ if raw.localizedCaseInsensitiveContains("offline mode") {
+ return AppL10n.string("settings.models.error.offline")
+ }
+ if raw.count > 280 {
+ return String(raw.prefix(277)) + "…"
+ }
+ return raw
+ }
+
+ /// Resolve on-disk cache directories for a model. Matches the layout
+ /// `HuggingFaceDownloader.getCacheDirectory(for:)` uses in Qwen3Speech.
+ nonisolated static func candidateCacheDirectories(for model: OnDeviceModel) -> [URL] {
+ let base = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first!
+ .appendingPathComponent("qwen3-speech", isDirectory: true)
+ let repoId = model.repoId
+ var candidates: [URL] = []
+
+ // Hub-style path (current default).
+ let parts = repoId.split(separator: "/", omittingEmptySubsequences: true)
+ if parts.count == 2 {
+ candidates.append(
+ base
+ .appendingPathComponent("models/\(parts[0])/\(parts[1])", isDirectory: true)
+ )
+ }
+
+ // Legacy flat path kept by HuggingFaceDownloader for backward compat.
+ let sanitized = repoId.replacingOccurrences(of: "/", with: "_")
+ candidates.append(base.appendingPathComponent(sanitized, isDirectory: true))
+
+ // Older OSGKeyboard probe paths (pre-alignment); delete still sweeps these.
+ switch model {
+ case .qwen3ASR:
+ candidates.append(base.appendingPathComponent("Qwen3ASR", isDirectory: true))
+ candidates.append(
+ base.appendingPathComponent("models/aufklarer/Qwen3-ASR-0.6B-MLX-4bit", isDirectory: true)
+ )
+ candidates.append(base.appendingPathComponent("aufklarer_Qwen3-ASR-0.6B-MLX-4bit", isDirectory: true))
+ }
+
+ return candidates
+ }
+
+ /// First candidate directory that already contains downloaded weights.
+ nonisolated static func existingCacheDirectory(for model: OnDeviceModel) -> URL? {
+ candidateCacheDirectories(for: model).first { dir in
+ weightsExist(in: dir, model: model)
+ }
+ }
+
+ nonisolated private static func weightsExist(in directory: URL, model: OnDeviceModel) -> Bool {
+ let fm = FileManager.default
+ switch model {
+ case .qwen3ASR:
+ let encoder = directory.appendingPathComponent("encoder.mlmodelc", isDirectory: true)
+ let decoder = directory.appendingPathComponent("decoder_part1.mlmodelc", isDirectory: true)
+ let vocab = directory.appendingPathComponent("vocab.json")
+ return fm.fileExists(atPath: encoder.path)
+ && fm.fileExists(atPath: decoder.path)
+ && fm.fileExists(atPath: vocab.path)
+ }
+ }
+
+ // MARK: - CoreML download
+
+ /// Downloads CoreML encoder/decoder bundles and tokenizer files into one cache dir.
+ nonisolated static func downloadQwen3CoreMLWeights(
+ model: OnDeviceModel,
+ registry: ModelRegistry,
+ progressHandler: @escaping @Sendable (Double, String) -> Void
+ ) async throws {
+ let coreMLId = model.repoId
+ let tokenizerId = model.tokenizerRepoId
+ let dir = try HuggingFaceDownloader.getCacheDirectory(for: coreMLId)
+
+ switch registry {
+ case .huggingFace(let hubEndpoint):
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: coreMLId,
+ to: dir,
+ additionalFiles: Qwen3CoreMLDownloadArtifacts.coreMLBundleGlobs,
+ hubEndpoint: hubEndpoint,
+ progressHandler: { progressHandler($0 * 0.85, "CoreML") }
+ )
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: tokenizerId,
+ to: dir,
+ additionalFiles: Qwen3CoreMLDownloadArtifacts.tokenizerFiles,
+ hubEndpoint: hubEndpoint,
+ progressHandler: { progressHandler(0.85 + $0 * 0.15, "Tokenizer") }
+ )
+ case .modelScope(let baseURL, let revision):
+ try await downloadQwen3CoreMLViaModelScope(
+ coreMLId: coreMLId,
+ tokenizerId: tokenizerId,
+ to: dir,
+ baseURL: baseURL,
+ revision: revision,
+ progressHandler: progressHandler
+ )
+ }
+ progressHandler(1.0, "Ready")
+ }
+
+ nonisolated private static func downloadQwen3CoreMLViaModelScope(
+ coreMLId: String,
+ tokenizerId: String,
+ to directory: URL,
+ baseURL: String,
+ revision: String,
+ progressHandler: @escaping @Sendable (Double, String) -> Void
+ ) async throws {
+ let coreListed = try await ModelScopeDownloader.listAllFiles(
+ modelId: coreMLId,
+ baseURL: baseURL,
+ revision: revision
+ )
+ let corePaths = coreListed.map(\.path).filter { path in
+ path.contains(".mlmodelc/") || path == "config.json"
+ }
+ guard !corePaths.isEmpty else {
+ throw DownloadError.failedToDownload("\(coreMLId): no CoreML files on ModelScope")
+ }
+ let coreSizes = Dictionary(uniqueKeysWithValues: coreListed.map { ($0.path, $0.size) })
+ try await ModelScopeDownloader.downloadFiles(
+ modelId: coreMLId,
+ to: directory,
+ files: corePaths,
+ fileSizes: coreSizes,
+ baseURL: baseURL,
+ revision: revision,
+ progressHandler: { progressHandler($0 * 0.85, "CoreML") }
+ )
+
+ let tokListed = try await ModelScopeDownloader.listAllFiles(
+ modelId: tokenizerId,
+ baseURL: baseURL,
+ revision: revision
+ )
+ let tokPaths = Qwen3CoreMLDownloadArtifacts.tokenizerFiles.filter { name in
+ tokListed.contains { $0.path == name }
+ }
+ let tokSizes = Dictionary(uniqueKeysWithValues: tokListed.map { ($0.path, $0.size) })
+ try await ModelScopeDownloader.downloadFiles(
+ modelId: tokenizerId,
+ to: directory,
+ files: tokPaths,
+ fileSizes: tokSizes,
+ baseURL: baseURL,
+ revision: revision,
+ progressHandler: { progressHandler(0.85 + $0 * 0.15, "Tokenizer") }
+ )
+ }
+
+ /// Preferred cache directory for display / storage badges.
+ nonisolated static func cacheDirectory(for model: OnDeviceModel) -> URL {
+ existingCacheDirectory(for: model)
+ ?? candidateCacheDirectories(for: model).first!
+ }
+}
diff --git a/OSGKeyboard/Services/OnDeviceModelWarmup.swift b/OSGKeyboard/Services/OnDeviceModelWarmup.swift
new file mode 100644
index 0000000..48b93b9
--- /dev/null
+++ b/OSGKeyboard/Services/OnDeviceModelWarmup.swift
@@ -0,0 +1,197 @@
+// OnDeviceModelWarmup.swift
+// OSGKeyboard · Main App
+//
+// Preloads on-device ASR weights for Flow sessions.
+
+import Foundation
+import OSGKeyboardShared
+
+@MainActor
+final class OnDeviceModelWarmup: ObservableObject {
+
+ static let shared = OnDeviceModelWarmup()
+
+ enum Phase: Equatable {
+ case idle
+ case warming
+ case ready
+ case failed(String)
+ case notNeeded
+
+ var isFailed: Bool {
+ if case .failed = self { return true }
+ return false
+ }
+ }
+
+ @Published private(set) var phase: Phase = .idle
+
+ /// Bumped on `invalidate()` and each new warm-up so cancelled tasks
+ /// cannot leave `phase` stuck on `.warming`.
+ private var warmupGeneration = 0
+ private var warmupTask: Task?
+ private var qwenASRService: Qwen3ASRService?
+ private var speechAnalyzerService: ASRService?
+ private var cloudASRService: ASRService?
+
+ private init() {}
+
+ /// Loads ASR into memory when the local stack is ready on disk.
+ func warmUpIfNeeded(force: Bool = false) {
+ let store = AppGroupStore()
+ guard store.engineMode == "local" else {
+ resetInstances()
+ phase = .notNeeded
+ publishMemoryReady(false)
+ return
+ }
+
+ guard store.localASRBackend != .qwen3ASR || OnDeviceMLRuntime.supportsOnDeviceQwen3 else {
+ resetInstances()
+ phase = .notNeeded
+ publishMemoryReady(false)
+ return
+ }
+
+ guard OnDeviceModelStatus.isLocalStackReady(asrBackend: store.localASRBackend) else {
+ resetInstances()
+ phase = .idle
+ publishMemoryReady(false)
+ return
+ }
+
+ var shouldForce = force
+ if phase == .ready, !shouldForce {
+ if needsModelReload() {
+ shouldForce = true
+ } else {
+ publishMemoryReady(true)
+ return
+ }
+ }
+ if phase == .warming { return }
+ if case .failed = phase, !shouldForce { return }
+
+ warmupTask?.cancel()
+ warmupGeneration += 1
+ let generation = warmupGeneration
+ phase = .warming
+ publishMemoryReady(false)
+
+ let asrBackend = store.localASRBackend
+ warmupTask = Task { @MainActor [weak self] in
+ guard let self else { return }
+ do {
+ try await self.performWarmup(asrBackend: asrBackend)
+ guard generation == self.warmupGeneration, !Task.isCancelled else { return }
+ self.phase = .ready
+ self.publishMemoryReady(true)
+ } catch {
+ guard generation == self.warmupGeneration, !Task.isCancelled else { return }
+ let message = (error as? LocalizedError)?.errorDescription
+ ?? error.localizedDescription
+ self.phase = .failed(message)
+ self.publishMemoryReady(false)
+ }
+ }
+ }
+
+ func invalidate() {
+ warmupGeneration += 1
+ warmupTask?.cancel()
+ warmupTask = nil
+ resetInstances()
+ phase = .idle
+ publishMemoryReady(false)
+ }
+
+ /// Called when returning from background — re-verify CoreML weights and
+ /// unstick a warmup that was frozen while the app was suspended.
+ func ensureReadyAfterBackground() {
+ let store = AppGroupStore()
+ guard store.engineMode == "local" else {
+ phase = .notNeeded
+ publishMemoryReady(false)
+ return
+ }
+
+ guard OnDeviceModelStatus.isLocalStackReady(asrBackend: store.localASRBackend) else {
+ phase = .idle
+ publishMemoryReady(false)
+ return
+ }
+
+ switch phase {
+ case .warming, .ready:
+ if needsModelReload() {
+ warmUpIfNeeded(force: true)
+ }
+ case .failed, .idle:
+ warmUpIfNeeded(force: true)
+ case .notNeeded:
+ break
+ }
+ }
+
+ func asrService(engineMode: String, localBackend: LocalASRBackend) -> ASRService {
+ if engineMode != "local" {
+ if cloudASRService == nil {
+ cloudASRService = ASRServiceFactory.make(
+ engineMode: engineMode,
+ localBackend: localBackend
+ )
+ }
+ return cloudASRService!
+ }
+
+ switch localBackend {
+ case .qwen3ASR:
+ if qwenASRService == nil {
+ qwenASRService = Qwen3ASRService()
+ }
+ return qwenASRService!
+ case .speechAnalyzer:
+ if speechAnalyzerService == nil {
+ speechAnalyzerService = ASRServiceFactory.make(
+ engineMode: "local",
+ localBackend: .speechAnalyzer
+ )
+ }
+ return speechAnalyzerService!
+ }
+ }
+
+ // MARK: - Internals
+
+ private func performWarmup(asrBackend: LocalASRBackend) async throws {
+ switch asrBackend {
+ case .qwen3ASR:
+ if qwenASRService == nil {
+ qwenASRService = Qwen3ASRService()
+ }
+ FlowDiagnostics.log("warmup ASR start backend=qwen3ASR")
+ try await qwenASRService!.warmUp()
+ FlowDiagnostics.log("warmup ASR done")
+ case .speechAnalyzer:
+ FlowDiagnostics.log("warmup skipped — speechAnalyzer backend")
+ }
+ }
+
+ private func resetInstances() {
+ qwenASRService = nil
+ speechAnalyzerService = nil
+ cloudASRService = nil
+ }
+
+ private func publishMemoryReady(_ ready: Bool) {
+ OnDeviceModelStatus.setModelsLoadedInMemory(ready)
+ }
+
+ private func needsModelReload() -> Bool {
+ let store = AppGroupStore()
+ guard store.engineMode == "local", store.localASRBackend == .qwen3ASR else {
+ return false
+ }
+ return qwenASRService?.isModelInMemory != true
+ }
+}
diff --git a/OSGKeyboard/Services/OpenSourceLicenseCatalog.swift b/OSGKeyboard/Services/OpenSourceLicenseCatalog.swift
new file mode 100644
index 0000000..33be0ac
--- /dev/null
+++ b/OSGKeyboard/Services/OpenSourceLicenseCatalog.swift
@@ -0,0 +1,115 @@
+// OpenSourceLicenseCatalog.swift
+// OSGKeyboard · Main App
+//
+// Single source of truth for third-party open-source components shipped
+// with or downloaded by OSGKeyboard. Consumed by Settings → About →
+// Third-Party Licenses.
+//
+// Keep this list aligned with `project.yml` package dependencies and
+// the default model ID in `Qwen3ASRService`.
+
+import Foundation
+
+enum OpenSourceLicenseCatalog {
+
+ struct Entry: Identifiable, Hashable {
+ let id: String
+ let name: String
+ let licenseName: String
+ /// One-line explanation shown in the popup and above the full text.
+ let purpose: String
+ let url: URL?
+ /// Verbatim license body for the long-scroll disclosure page.
+ let licenseText: String
+ }
+
+ /// Bundled libraries and runtime model artefacts referenced by the app.
+ static let entries: [Entry] = [
+ .init(
+ id: "speech-swift",
+ name: "soniqo/speech-swift",
+ licenseName: "Apache-2.0",
+ purpose: "On-device ASR runtime. Vendored locally as the Qwen3Speech SPM package (Qwen3ASR CoreML path, AudioCommon, SpeechVAD).",
+ url: URL(string: "https://github.com/soniqo/speech-swift"),
+ licenseText: apache2Text
+ ),
+ .init(
+ id: "swift-transformers",
+ name: "huggingface/swift-transformers",
+ licenseName: "Apache-2.0",
+ purpose: "Hugging Face Hub client and tokenizer bindings. Used to resolve and download model snapshots at runtime.",
+ url: URL(string: "https://github.com/huggingface/swift-transformers"),
+ licenseText: apache2Text
+ ),
+ .init(
+ id: "qwen3-asr-coreml",
+ name: "aufklarer/Qwen3-ASR-CoreML",
+ licenseName: "Apache-2.0",
+ purpose: "CoreML INT8 weights for Qwen3-ASR-0.6B (derived from Alibaba Qwen team). Downloaded on first use (~1.6 GB); not bundled in the app binary.",
+ url: URL(string: "https://huggingface.co/aufklarer/Qwen3-ASR-CoreML"),
+ licenseText: apache2Text
+ ),
+ .init(
+ id: "qwen3-asr-upstream",
+ name: "Qwen/Qwen3-ASR-0.6B",
+ licenseName: "Apache-2.0",
+ purpose: "Original ASR model by Alibaba's Qwen team. CoreML bundle and tokenizer files are derived from these weights.",
+ url: URL(string: "https://huggingface.co/Qwen/Qwen3-ASR-0.6B"),
+ licenseText: apache2Text
+ ),
+ .init(
+ id: "material-icons",
+ name: "Google Material Icons",
+ licenseName: "Apache-2.0",
+ purpose: "MaterialIcons-Regular.ttf bundled for Settings and navigation iconography.",
+ url: URL(string: "https://github.com/google/material-design-icons"),
+ licenseText: apache2Text
+ ),
+ ]
+
+ // MARK: - License bodies
+
+ static let apache2Text = """
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied. See the License for the specific language governing
+ permissions and limitations under the License.
+ """
+
+ static let mitText = """
+ MIT License
+
+ Copyright (c) 2023 Apple Inc.
+
+ Permission is hereby granted, free of charge, to any person
+ obtaining a copy of this software and associated documentation
+ files (the "Software"), to deal in the Software without
+ restriction, including without limitation the rights to use, copy,
+ modify, merge, publish, distribute, sublicense, and/or sell copies
+ of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+
+ The above copyright notice and this permission notice shall be
+ included in all copies or substantial portions of the Software.
+
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
+ OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+ NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
+ HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+ WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+ OTHER DEALINGS IN THE SOFTWARE.
+ """
+}
diff --git a/OSGKeyboard/Services/Qwen3ASRService.swift b/OSGKeyboard/Services/Qwen3ASRService.swift
new file mode 100644
index 0000000..a3f5cfb
--- /dev/null
+++ b/OSGKeyboard/Services/Qwen3ASRService.swift
@@ -0,0 +1,257 @@
+// Qwen3ASRService.swift
+// OSGKeyboard · Main App
+//
+// On-device ASR via Qwen3-ASR-0.6B CoreML (Neural Engine + CPU). Uses the
+// MLX-free `transcribeBackgroundSafe` path so Flow dictation works while the
+// host app is backgrounded (no Metal GPU).
+
+import Foundation
+import AVFoundation
+import os
+import OSGKeyboardShared
+@preconcurrency import Qwen3ASR
+
+struct Qwen3ASRServiceProvider: ASRServiceProvider {
+ let backend: LocalASRBackend = .qwen3ASR
+ func make() -> ASRService { Qwen3ASRService() }
+}
+
+final class Qwen3ASRService: ASRService, @unchecked Sendable {
+
+ private enum TranscribeConstants {
+ static let sampleRate = 16_000
+ /// CoreML encoder is exported for 30 s windows — align with Flow chunking.
+ static let chunkDurationSeconds = Int(FlowUtteranceChunkConfig.flowDefault.maxChunkDurationSeconds)
+ }
+
+ private let lock = OSAllocatedUnfairLock()
+ private var currentTask: Task?
+ private var cancelled = false
+
+ private var model: CoreMLASRModel?
+ private var loadError: Error?
+ private var loadingTask: Task?
+
+ private func resolveModel() async throws -> CoreMLASRModel {
+ if let model = lock.withLock({ self.model }) { return model }
+ if let err = lock.withLock({ self.loadError }) { throw err }
+
+ guard ModelManager.weightsOnDisk(for: .qwen3ASR) else {
+ throw ASRServiceError.modelNotDownloaded
+ }
+
+ guard OnDeviceMLRuntime.supportsOnDeviceQwen3 else {
+ throw ASRServiceError.unsupportedOS
+ }
+
+ let task: Task = lock.withLock {
+ if let existing = loadingTask { return existing }
+ let new = Task { [weak self] in
+ guard let self else { throw ASRServiceError.notReady }
+ let cacheDir = ModelManager.cacheDirectory(for: .qwen3ASR)
+ let loaded = try await CoreMLASRModel.fromPretrained(
+ tokenizerModelId: OnDeviceModel.qwen3ASR.tokenizerRepoId,
+ cacheDir: cacheDir,
+ offlineMode: true,
+ progressHandler: { @Sendable _, _ in }
+ )
+ try loaded.warmUp()
+ self.lock.withLock { self.model = loaded }
+ return loaded
+ }
+ loadingTask = new
+ return new
+ }
+ do {
+ let model = try await task.value
+ return model
+ } catch {
+ lock.withLock { self.loadError = error }
+ throw error
+ }
+ }
+
+ func warmUp() async throws {
+ FlowDiagnostics.log("Qwen3ASR CoreML warmUp start")
+ resetForNewUtterance()
+ _ = try await resolveModel()
+ FlowDiagnostics.log("Qwen3ASR CoreML warmUp done")
+ }
+
+ func resetForNewUtterance() {
+ lock.withLock { cancelled = false }
+ }
+
+ var isModelInMemory: Bool {
+ lock.withLock { model != nil }
+ }
+
+ func transcribe(
+ stream: AsyncStream,
+ locale: Locale
+ ) -> AsyncStream {
+ AsyncStream { continuation in
+ continuation.yield(.capability(onDeviceSupported: true))
+
+ let task = Task { [weak self] in
+ guard let self else { return }
+ defer { self.lock.withLock { self.currentTask = nil } }
+
+ var samples: [Float] = []
+ samples.reserveCapacity(
+ Int(Double(TranscribeConstants.sampleRate) * FlowSessionKeys.maxUtteranceDuration) + 16_000
+ )
+ for await snap in stream {
+ if Task.isCancelled || self.cancelledNow() { break }
+ samples.append(contentsOf: snap.samples)
+ }
+ guard !Task.isCancelled, !self.cancelledNow() else {
+ continuation.finish()
+ return
+ }
+ if samples.isEmpty {
+ continuation.yield(.error(SharedL10n.string("error.asr.noSpeech")))
+ continuation.finish()
+ return
+ }
+
+ do {
+ let model = try await self.resolveModel()
+ let language = Self.languageHint(from: locale)
+ let durationSec = Double(samples.count) / 16_000.0
+ FlowDiagnostics.log(
+ "Qwen3ASR CoreML transcribe start samples=\(samples.count) " +
+ "duration=\(String(format: "%.1f", durationSec))s"
+ )
+ let text = self.transcribeInChunks(
+ model: model,
+ samples: samples,
+ language: language
+ )
+ FlowDiagnostics.log("Qwen3ASR CoreML transcribe done chars=\(text.count)")
+ let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
+ if trimmed.isEmpty {
+ continuation.yield(.error(SharedL10n.string("error.asr.noSpeech")))
+ } else {
+ continuation.yield(.final(trimmed))
+ }
+ continuation.finish()
+ } catch {
+ Self.debug("Qwen3ASR.transcribe failed: \(error.localizedDescription)")
+ continuation.yield(.error(error.localizedDescription))
+ continuation.finish()
+ }
+ }
+ self.lock.withLock { self.currentTask = task }
+
+ continuation.onTermination = { @Sendable [weak self] _ in
+ self?.cancel()
+ }
+ }
+ }
+
+ func cancel() {
+ lock.withLock {
+ cancelled = true
+ currentTask?.cancel()
+ currentTask = nil
+ }
+ }
+
+ func transcribeChunk(samples: [Float], locale: Locale) async -> ASRChunkResult {
+ if cancelledNow() || Task.isCancelled { return .cancelled }
+ guard !samples.isEmpty else { return .success("") }
+
+ do {
+ let model = try await resolveModel()
+ let language = Self.languageHint(from: locale)
+ let text = model.transcribeBackgroundSafe(
+ audio: samples,
+ sampleRate: TranscribeConstants.sampleRate,
+ language: language
+ )
+ .trimmingCharacters(in: .whitespacesAndNewlines)
+ if text.hasPrefix("[CoreML error:") {
+ return .failure(text)
+ }
+ return .success(text)
+ } catch {
+ let message = (error as? LocalizedError)?.errorDescription ?? error.localizedDescription
+ return .failure(message)
+ }
+ }
+
+ private func cancelledNow() -> Bool {
+ lock.withLock { cancelled }
+ }
+
+ private func transcribeInChunks(
+ model: CoreMLASRModel,
+ samples: [Float],
+ language: String?
+ ) -> String {
+ let chunkSize = TranscribeConstants.sampleRate * TranscribeConstants.chunkDurationSeconds
+ guard samples.count > chunkSize else {
+ return model.transcribeBackgroundSafe(
+ audio: samples,
+ sampleRate: TranscribeConstants.sampleRate,
+ language: language
+ )
+ }
+
+ var parts: [String] = []
+ parts.reserveCapacity((samples.count + chunkSize - 1) / chunkSize)
+ var offset = 0
+ var chunkIndex = 0
+ while offset < samples.count {
+ let end = min(offset + chunkSize, samples.count)
+ let chunk = Array(samples[offset.. String? {
+ let id = locale.identifier.lowercased()
+ if id.hasPrefix("zh") { return "zh" }
+ if id.hasPrefix("en") { return "en" }
+ return locale.language.languageCode?.identifier
+ }
+
+ private static func debug(_ message: String) {
+ #if DEBUG
+ print("🎙️[Qwen3ASR] \(message)")
+ #endif
+ }
+}
+
+private enum ASRServiceError: Error, LocalizedError {
+ case notReady
+ case modelNotDownloaded
+ case unsupportedOS
+
+ var errorDescription: String? {
+ switch self {
+ case .notReady:
+ return nil
+ case .modelNotDownloaded:
+ return AppL10n.string("asr.error.modelNotDownloaded")
+ case .unsupportedOS:
+ return AppL10n.string("asr.error.unsupportedOS")
+ }
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Package.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Package.swift
new file mode 100644
index 0000000..74a28c9
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Package.swift
@@ -0,0 +1,99 @@
+// swift-tools-version: 5.10
+import PackageDescription
+
+// Local fork of soniqo/speech-swift that ships ONLY what OSGKeyboard
+// consumes: Qwen3ASR + Qwen3Chat. The original repo's `Package.swift`
+// references a `CSpeechCore` binary target whose URL doesn't match
+// its declared filename (`SpeechCore.xcframework.zip` vs target
+// name `CSpeechCore`), which breaks SwiftPM resolve on a clean
+// checkout. The only thing we need from speech-swift for the
+// OSGKeyboard on-device path is the two Qwen3 modules and the
+// AudioCommon / MLXCommon / SpeechVAD slices they depend on — the
+// AudioServer / AudioCLI / AudioCLILib targets that pulled in
+// SpeechCore aren't part of our build graph.
+//
+// Source provenance: every `.swift` file in `Sources//` is
+// copied from https://github.com/soniqo/speech-swift (commit pinned
+// to v0.0.21 of the upstream tag tree). Original copyright
+// headers are preserved in each file. Apache-2.0 license.
+//
+// Track upstream: when soniqo fixes the binary-target mismatch in
+// their main `Package.swift`, delete this local package and
+// re-enable the upstream dependency in the host project.
+
+let package = Package(
+ name: "Qwen3Speech",
+ platforms: [
+ .iOS("18.0"),
+ .macOS("15.0")
+ ],
+ products: [
+ .library(name: "Qwen3ASR", targets: ["Qwen3ASR"]),
+ .library(name: "Qwen3Chat", targets: ["Qwen3Chat"]),
+ ],
+ dependencies: [
+ // mlx-swift is the Apple MLX array framework bindings; Qwen3
+ // runtime depends on the GPU side, the chat runtime depends
+ // on the linear-attention kernels exposed by MLXNN / MLXFast.
+ //
+ // We pin to a local flattened copy at `~/.local/mlx-swift`
+ // (an exported snapshot of mlx-swift 0.31.4 with its Cmlx /
+ // mlx-c submodules baked in as plain directories) because
+ // SwiftPM can't reliably fetch the upstream's git submodules
+ // on this network — the Cmlx/mlx submodule is ~700 MB of
+ // history and the clone drops mid-fetch. The snapshot is
+ // generated once on a healthy network, kept outside the
+ // project, and re-used on every resolve.
+ .package(path: "/Users/rocky/.local/mlx-swift"),
+ // swift-transformers exposes Hugging Face Hub and tokenizers
+ // — AudioCommon uses Hub to resolve repo → snapshot path.
+ .package(url: "https://github.com/huggingface/swift-transformers", from: "1.1.6"),
+ ],
+ targets: [
+ .target(
+ name: "AudioCommon",
+ dependencies: [
+ .product(name: "Hub", package: "swift-transformers"),
+ ]
+ ),
+ .target(
+ name: "MLXCommon",
+ dependencies: [
+ "AudioCommon",
+ .product(name: "MLX", package: "mlx-swift"),
+ .product(name: "MLXNN", package: "mlx-swift"),
+ .product(name: "MLXFast", package: "mlx-swift"),
+ ]
+ ),
+ .target(
+ name: "SpeechVAD",
+ dependencies: [
+ "AudioCommon",
+ "MLXCommon",
+ .product(name: "MLX", package: "mlx-swift"),
+ .product(name: "MLXNN", package: "mlx-swift"),
+ ]
+ ),
+ .target(
+ name: "Qwen3ASR",
+ dependencies: [
+ "AudioCommon",
+ "MLXCommon",
+ "SpeechVAD",
+ .product(name: "MLX", package: "mlx-swift"),
+ .product(name: "MLXNN", package: "mlx-swift"),
+ .product(name: "MLXFast", package: "mlx-swift"),
+ ]
+ ),
+ .target(
+ name: "Qwen3Chat",
+ dependencies: [
+ "AudioCommon",
+ "MLXCommon",
+ .product(name: "MLX", package: "mlx-swift"),
+ .product(name: "MLXNN", package: "mlx-swift"),
+ .product(name: "MLXFast", package: "mlx-swift"),
+ ]
+ ),
+ ]
+)
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioFileLoader.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioFileLoader.swift
new file mode 100644
index 0000000..42812d6
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioFileLoader.swift
@@ -0,0 +1,382 @@
+import Foundation
+import AVFoundation
+
+/// Sample-rate-conversion quality. Both options fully drain the converter and
+/// produce exact-length output; they differ only in the SRC filter.
+public enum ResampleQuality {
+ /// Framework-default band-limited SRC (`Normal` algorithm). Anti-aliases
+ /// steep downsamples and retains high frequencies well below Nyquist;
+ /// rolls off slightly more near Nyquist than `.mastering`. The right
+ /// default for speech/voice, which is band-limited and usually
+ /// downsampled (e.g. 44.1k→16k for ASR), where mastering-grade filtering
+ /// is wasted cost.
+ case standard
+ /// Mastering algorithm at maximum quality — fullest high-frequency
+ /// retention right up to Nyquist, at higher cost. Use for music (source
+ /// separation) and upsampling/super-resolution, where full-band fidelity
+ /// matters.
+ case mastering
+}
+
+/// Loads audio files and converts to float samples
+public enum AudioFileLoader {
+ /// Load audio file and return samples at target sample rate.
+ /// `quality` selects the SRC filter when resampling (default `.standard`;
+ /// pass `.mastering` for music/upsampling).
+ public static func load(url: URL, targetSampleRate: Int = 24000, quality: ResampleQuality = .standard) throws -> [Float] {
+ let audioFile = try AVAudioFile(forReading: url)
+ let format = audioFile.processingFormat
+ let frameCount = AVAudioFrameCount(audioFile.length)
+
+ guard let buffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: frameCount) else {
+ throw AudioLoadError.bufferCreationFailed
+ }
+
+ try audioFile.read(into: buffer)
+
+ guard let floatData = buffer.floatChannelData else {
+ throw AudioLoadError.noFloatData
+ }
+
+ // Get mono samples (use first channel)
+ let samples = Array(UnsafeBufferPointer(start: floatData[0], count: Int(buffer.frameLength)))
+
+ // Resample if needed
+ let inputSampleRate = Int(format.sampleRate)
+ if inputSampleRate != targetSampleRate {
+ return resample(samples, from: inputSampleRate, to: targetSampleRate, quality: quality)
+ }
+
+ return samples
+ }
+
+ /// Load audio file and return stereo channels at target sample rate.
+ /// Returns `[left, right]` — mono files are duplicated to stereo.
+ /// `quality` selects the SRC filter when resampling (default `.standard`;
+ /// pass `.mastering` for music).
+ public static func loadStereo(url: URL, targetSampleRate: Int = 44100, quality: ResampleQuality = .standard) throws -> [[Float]] {
+ let audioFile = try AVAudioFile(forReading: url)
+ let format = audioFile.processingFormat
+ let frameCount = AVAudioFrameCount(audioFile.length)
+
+ guard let buffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: frameCount) else {
+ throw AudioLoadError.bufferCreationFailed
+ }
+
+ try audioFile.read(into: buffer)
+
+ guard let floatData = buffer.floatChannelData else {
+ throw AudioLoadError.noFloatData
+ }
+
+ let count = Int(buffer.frameLength)
+ let left = Array(UnsafeBufferPointer(start: floatData[0], count: count))
+ let right: [Float]
+ if format.channelCount >= 2 {
+ right = Array(UnsafeBufferPointer(start: floatData[1], count: count))
+ } else {
+ right = left // Mono → duplicate
+ }
+
+ let inputSampleRate = Int(format.sampleRate)
+ if inputSampleRate != targetSampleRate {
+ // Resample both channels in one converter pass so L/R stay
+ // phase-aligned (two independent converters can drift).
+ return resampleStereo([left, right], from: inputSampleRate, to: targetSampleRate, quality: quality)
+ }
+
+ return [left, right]
+ }
+
+ /// Load WAV file directly (for 16-bit PCM)
+ public static func loadWAV(url: URL) throws -> (samples: [Float], sampleRate: Int) {
+ let data = try Data(contentsOf: url)
+
+ // Parse WAV header
+ guard data.count > 44 else {
+ throw AudioLoadError.invalidWAVFile
+ }
+
+ // Check RIFF header
+ let riff = String(data: data[0..<4], encoding: .ascii)
+ guard riff == "RIFF" else {
+ throw AudioLoadError.invalidWAVFile
+ }
+
+ // Check WAVE format
+ let wave = String(data: data[8..<12], encoding: .ascii)
+ guard wave == "WAVE" else {
+ throw AudioLoadError.invalidWAVFile
+ }
+
+ // Parse format chunk (handle unaligned reads)
+ let audioFormat = data[20..<22].withUnsafeBytes { $0.loadUnaligned(as: UInt16.self) }
+ let numChannels = data[22..<24].withUnsafeBytes { $0.loadUnaligned(as: UInt16.self) }
+ let sampleRate = data[24..<28].withUnsafeBytes { $0.loadUnaligned(as: UInt32.self) }
+ let bitsPerSample = data[34..<36].withUnsafeBytes { $0.loadUnaligned(as: UInt16.self) }
+
+ guard audioFormat == 1 else { // PCM
+ throw AudioLoadError.unsupportedFormat("Not PCM format")
+ }
+
+ guard numChannels > 0 else {
+ throw AudioLoadError.invalidWAVFile
+ }
+
+ guard bitsPerSample == 16 else {
+ throw AudioLoadError.unsupportedFormat("Not 16-bit")
+ }
+
+ // Find data chunk
+ var dataOffset = 36
+ var dataChunkSize: UInt32? = nil
+ while dataOffset < data.count - 8 {
+ let chunkId = String(data: data[dataOffset..<(dataOffset+4)], encoding: .ascii)
+ let chunkSize = data[(dataOffset+4)..<(dataOffset+8)].withUnsafeBytes { $0.loadUnaligned(as: UInt32.self) }
+
+ if chunkId == "data" {
+ dataOffset += 8
+ dataChunkSize = chunkSize
+ break
+ }
+
+ // Validate chunk advance to avoid out-of-bounds.
+ let nextOffset = dataOffset + 8 + Int(chunkSize)
+ guard nextOffset >= dataOffset, nextOffset <= data.count else {
+ throw AudioLoadError.invalidWAVFile
+ }
+ dataOffset = nextOffset
+ }
+
+ // Read samples
+ guard let chunkSize = dataChunkSize else {
+ throw AudioLoadError.invalidWAVFile
+ }
+ let chunkSizeInt = Int(chunkSize)
+ guard dataOffset >= 0, dataOffset <= data.count, dataOffset + chunkSizeInt <= data.count else {
+ throw AudioLoadError.invalidWAVFile
+ }
+
+ let sampleData = data[dataOffset..<(dataOffset + chunkSizeInt)]
+ let channels = Int(numChannels)
+ let bytesPerSample = 2
+ let frameSize = bytesPerSample * channels
+ let sampleCount = sampleData.count / frameSize
+
+ var samples = [Float](repeating: 0, count: sampleCount)
+ sampleData.withUnsafeBytes { ptr in
+ let int16Ptr = ptr.bindMemory(to: Int16.self)
+ for i in 0.. [Float] {
+ guard inputRate != outputRate, !samples.isEmpty else { return samples }
+
+ guard let sourceFormat = AVAudioFormat(
+ commonFormat: .pcmFormatFloat32, sampleRate: Double(inputRate),
+ channels: 1, interleaved: false),
+ let targetFormat = AVAudioFormat(
+ commonFormat: .pcmFormatFloat32, sampleRate: Double(outputRate),
+ channels: 1, interleaved: false),
+ let converter = AVAudioConverter(from: sourceFormat, to: targetFormat),
+ let sourceBuffer = AVAudioPCMBuffer(
+ pcmFormat: sourceFormat, frameCapacity: AVAudioFrameCount(samples.count))
+ else {
+ return samples
+ }
+
+ configureSRC(converter, quality: quality)
+ sourceBuffer.frameLength = AVAudioFrameCount(samples.count)
+ samples.withUnsafeBufferPointer { src in
+ sourceBuffer.floatChannelData![0].update(from: src.baseAddress!, count: samples.count)
+ }
+
+ let ratio = Double(outputRate) / Double(inputRate)
+ guard let out = convertDrained(
+ converter: converter, source: sourceBuffer, targetFormat: targetFormat,
+ inputFrames: samples.count, ratio: ratio, channels: 1)
+ else {
+ return samples
+ }
+ return out[0]
+ }
+
+ /// Resample a stereo signal in a single converter pass so the two channels
+ /// stay phase-aligned. `channels[0]` = left, `channels[1]` = right; both
+ /// must have equal length. `quality` selects the SRC filter (default
+ /// `.standard`; pass `.mastering` for music). Falls back to per-channel
+ /// mono resampling for non-stereo input or on converter-setup failure.
+ public static func resampleStereo(_ channels: [[Float]], from inputRate: Int, to outputRate: Int, quality: ResampleQuality = .standard) -> [[Float]] {
+ guard channels.count == 2 else {
+ return channels.map { resample($0, from: inputRate, to: outputRate, quality: quality) }
+ }
+ let n = channels[0].count
+ guard inputRate != outputRate, n > 0, channels[1].count == n else {
+ return channels
+ }
+
+ guard let sourceFormat = AVAudioFormat(
+ commonFormat: .pcmFormatFloat32, sampleRate: Double(inputRate),
+ channels: 2, interleaved: false),
+ let targetFormat = AVAudioFormat(
+ commonFormat: .pcmFormatFloat32, sampleRate: Double(outputRate),
+ channels: 2, interleaved: false),
+ let converter = AVAudioConverter(from: sourceFormat, to: targetFormat),
+ let sourceBuffer = AVAudioPCMBuffer(
+ pcmFormat: sourceFormat, frameCapacity: AVAudioFrameCount(n))
+ else {
+ return channels.map { resample($0, from: inputRate, to: outputRate, quality: quality) }
+ }
+
+ configureSRC(converter, quality: quality)
+ sourceBuffer.frameLength = AVAudioFrameCount(n)
+ channels[0].withUnsafeBufferPointer {
+ sourceBuffer.floatChannelData![0].update(from: $0.baseAddress!, count: n)
+ }
+ channels[1].withUnsafeBufferPointer {
+ sourceBuffer.floatChannelData![1].update(from: $0.baseAddress!, count: n)
+ }
+
+ let ratio = Double(outputRate) / Double(inputRate)
+ guard let out = convertDrained(
+ converter: converter, source: sourceBuffer, targetFormat: targetFormat,
+ inputFrames: n, ratio: ratio, channels: 2)
+ else {
+ return channels.map { resample($0, from: inputRate, to: outputRate, quality: quality) }
+ }
+ return out
+ }
+
+ /// Configure the converter's SRC filter. Must be set before the first
+ /// `convert`. `.standard` leaves the framework default (`Normal`); only
+ /// `.mastering` opts into the slower, full-band Mastering algorithm.
+ private static func configureSRC(_ converter: AVAudioConverter, quality: ResampleQuality) {
+ switch quality {
+ case .standard:
+ break // framework default Normal SRC — already drains + exact length
+ case .mastering:
+ converter.sampleRateConverterAlgorithm = AVSampleRateConverterAlgorithm_Mastering
+ converter.sampleRateConverterQuality = .max
+ }
+ }
+
+ /// Run the converter to completion, draining its internal tail via
+ /// `.endOfStream`, and return one Float array per channel normalized to the
+ /// exact expected frame count.
+ ///
+ /// Returns `nil` unless the converter reaches `.endOfStream` cleanly. Only
+ /// `.endOfStream` is success: `.error` (or a thrown `NSError`) is a hard
+ /// failure, and a no-progress step that isn't end-of-stream means the
+ /// converter is stuck. In every non-success case the partial output is
+ /// discarded rather than returned, so callers can fall back instead of
+ /// silently propagating a truncated buffer (which would desync downstream
+ /// audio/video).
+ private static func convertDrained(
+ converter: AVAudioConverter,
+ source: AVAudioPCMBuffer,
+ targetFormat: AVAudioFormat,
+ inputFrames: Int,
+ ratio: Double,
+ channels: Int
+ ) -> [[Float]]? {
+ // ceil + headroom for the sinc filter's priming/tail latency.
+ let capacity = AVAudioFrameCount(ceil(Double(inputFrames) * ratio)) + 4096
+ guard let target = AVAudioPCMBuffer(pcmFormat: targetFormat, frameCapacity: capacity) else {
+ return nil
+ }
+
+ var out = [[Float]](repeating: [], count: channels)
+ for c in 0.. 0, let chans = target.floatChannelData {
+ for c in 0.. 0, !out[0].isEmpty else { return nil }
+ for c in 0.. expected {
+ out[c].removeLast(out[c].count - expected)
+ } else if out[c].count < expected {
+ out[c].append(contentsOf: repeatElement(0, count: expected - out[c].count))
+ }
+ }
+ return out
+ }
+}
+
+public enum AudioLoadError: Error, LocalizedError {
+ case bufferCreationFailed
+ case noFloatData
+ case invalidWAVFile
+ case unsupportedFormat(String)
+
+ public var errorDescription: String? {
+ switch self {
+ case .bufferCreationFailed:
+ return "Failed to create audio buffer"
+ case .noFloatData:
+ return "No float channel data available"
+ case .invalidWAVFile:
+ return "Invalid WAV file format"
+ case .unsupportedFormat(let reason):
+ return "Unsupported audio format: \(reason)"
+ }
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioIO.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioIO.swift
new file mode 100644
index 0000000..ad7fd3e
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioIO.swift
@@ -0,0 +1,174 @@
+#if canImport(AVFoundation)
+import AVFoundation
+import os
+
+/// Reusable audio I/O manager — handles mic capture, resampling, and playback.
+///
+/// Eliminates AVAudioEngine boilerplate that every demo app reimplements.
+///
+/// ```swift
+/// let audio = AudioIO()
+/// try audio.startMicrophone(targetSampleRate: 16000) { samples in
+/// pipeline.pushAudio(samples)
+/// }
+/// audio.player.scheduleChunk(ttsOutput)
+/// audio.stopMicrophone()
+/// ```
+public final class AudioIO {
+ /// Microphone state.
+ public enum MicrophoneState: Sendable {
+ case stopped, running, error(String)
+ }
+
+ /// Audio player for TTS output. Attached to the engine when mic starts.
+ public let player = StreamingAudioPlayer()
+
+ /// Current microphone state.
+ public private(set) var microphoneState: MicrophoneState = .stopped
+
+ /// RMS audio level (0.0–1.0) for UI meters. Updated on each mic buffer.
+ public private(set) var audioLevel: Float = 0
+
+ /// Whether to enable Voice Processing I/O for echo cancellation.
+ public let enableAEC: Bool
+
+ /// Playback sample rate (for TTS output).
+ public let playbackSampleRate: Double
+
+ private var engine: AVAudioEngine?
+ private static let log = Logger(subsystem: "audio.soniqo", category: "AudioIO")
+
+ public init(enableAEC: Bool = false, playbackSampleRate: Double = 24000) {
+ self.enableAEC = enableAEC
+ self.playbackSampleRate = playbackSampleRate
+ }
+
+ /// Start microphone capture, resampled to targetSampleRate.
+ ///
+ /// Also attaches the player to the engine for simultaneous playback.
+ /// Call `player.scheduleChunk()` to play audio while recording.
+ ///
+ /// - Parameters:
+ /// - targetSampleRate: Output sample rate for onSamples (default 16kHz for VAD/ASR)
+ /// - onSamples: Callback with resampled mono Float32 samples (called on audio thread)
+ public func startMicrophone(
+ targetSampleRate: Int = 16000,
+ onSamples: @escaping ([Float]) -> Void
+ ) throws {
+ stopMicrophone()
+
+ #if os(iOS)
+ let session = AVAudioSession.sharedInstance()
+ if enableAEC {
+ try session.setCategory(.playAndRecord, mode: .voiceChat, options: [.defaultToSpeaker, .allowBluetoothHFP])
+ } else {
+ try session.setCategory(.playAndRecord, mode: .default, options: [.defaultToSpeaker, .allowBluetoothHFP])
+ }
+ try session.setActive(true)
+ #endif
+
+ let engine = AVAudioEngine()
+ let inputNode = engine.inputNode
+ let hwFormat = inputNode.outputFormat(forBus: 0)
+
+ // Mono intermediate at hardware rate
+ guard let monoFormat = AVAudioFormat(
+ commonFormat: .pcmFormatFloat32,
+ sampleRate: hwFormat.sampleRate,
+ channels: 1,
+ interleaved: false
+ ) else {
+ microphoneState = .error("Cannot create mono format")
+ return
+ }
+
+ // Target format for VAD/ASR
+ guard let targetFormat = AVAudioFormat(
+ commonFormat: .pcmFormatFloat32,
+ sampleRate: Double(targetSampleRate),
+ channels: 1,
+ interleaved: false
+ ) else {
+ microphoneState = .error("Cannot create target format")
+ return
+ }
+
+ guard let resampler = AVAudioConverter(from: monoFormat, to: targetFormat) else {
+ microphoneState = .error("Cannot create resampler")
+ return
+ }
+
+ inputNode.installTap(onBus: 0, bufferSize: 1024, format: hwFormat) { [weak self] buffer, _ in
+ guard let self else { return }
+ guard let srcData = buffer.floatChannelData else { return }
+ let frameLen = Int(buffer.frameLength)
+ guard frameLen > 0 else { return }
+
+ // Extract channel 0 into mono buffer
+ guard let monoBuffer = AVAudioPCMBuffer(pcmFormat: monoFormat, frameCapacity: buffer.frameCapacity) else { return }
+ monoBuffer.frameLength = buffer.frameLength
+ memcpy(monoBuffer.floatChannelData![0], srcData[0], frameLen * MemoryLayout.size)
+
+ // Resample
+ let outFrameCount = AVAudioFrameCount(Double(frameLen) * Double(targetSampleRate) / hwFormat.sampleRate)
+ guard outFrameCount > 0,
+ let outBuffer = AVAudioPCMBuffer(pcmFormat: targetFormat, frameCapacity: outFrameCount) else { return }
+
+ var error: NSError?
+ resampler.convert(to: outBuffer, error: &error) { _, outStatus in
+ outStatus.pointee = .haveData
+ return monoBuffer
+ }
+ if error != nil { return }
+
+ guard let outData = outBuffer.floatChannelData else { return }
+ let count = Int(outBuffer.frameLength)
+ guard count > 0 else { return }
+ let samples = Array(UnsafeBufferPointer(start: outData[0], count: count))
+
+ // RMS for audio level
+ var sum: Float = 0
+ for s in samples { sum += s * s }
+ self.audioLevel = sqrt(sum / max(Float(count), 1))
+
+ onSamples(samples)
+ }
+
+ // Attach player for TTS output
+ guard let playerFormat = AVAudioFormat(
+ commonFormat: .pcmFormatFloat32,
+ sampleRate: playbackSampleRate,
+ channels: 1,
+ interleaved: false
+ ) else { return }
+ player.attach(to: engine, format: playerFormat)
+
+ do {
+ try engine.start()
+ player.startPlayback()
+ self.engine = engine
+ microphoneState = .running
+ Self.log.info("Microphone started at \(targetSampleRate)Hz, player at \(self.playbackSampleRate)Hz")
+ } catch {
+ microphoneState = .error(error.localizedDescription)
+ throw error
+ }
+ }
+
+ /// Stop microphone capture and detach player.
+ public func stopMicrophone() {
+ if let engine {
+ engine.inputNode.removeTap(onBus: 0)
+ player.detach(from: engine)
+ engine.stop()
+ }
+ engine = nil
+ audioLevel = 0
+ microphoneState = .stopped
+ }
+
+ deinit {
+ stopMicrophone()
+ }
+}
+#endif
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioModelError.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioModelError.swift
new file mode 100644
index 0000000..b07fb28
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioModelError.swift
@@ -0,0 +1,34 @@
+import Foundation
+
+/// Unified error type for audio model operations.
+public enum AudioModelError: Error, LocalizedError {
+ /// Model failed to load from disk or network.
+ case modelLoadFailed(modelId: String, reason: String, underlying: Error? = nil)
+ /// Weight file could not be read or parsed.
+ case weightLoadingFailed(path: String, underlying: Error? = nil)
+ /// Inference or generation step failed.
+ case inferenceFailed(operation: String, reason: String)
+ /// Model configuration is invalid or incompatible.
+ case invalidConfiguration(model: String, reason: String)
+ /// Voice preset file not found.
+ case voiceNotFound(voice: String, searchPath: String)
+
+ public var errorDescription: String? {
+ switch self {
+ case .modelLoadFailed(let modelId, let reason, let underlying):
+ var msg = "Failed to load model '\(modelId)': \(reason)"
+ if let underlying { msg += " (\(underlying.localizedDescription))" }
+ return msg
+ case .weightLoadingFailed(let path, let underlying):
+ var msg = "Failed to load weights from '\(path)'"
+ if let underlying { msg += ": \(underlying.localizedDescription)" }
+ return msg
+ case .inferenceFailed(let operation, let reason):
+ return "Inference failed during \(operation): \(reason)"
+ case .invalidConfiguration(let model, let reason):
+ return "Invalid configuration for '\(model)': \(reason)"
+ case .voiceNotFound(let voice, let searchPath):
+ return "Voice preset '\(voice)' not found at '\(searchPath)'"
+ }
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioRingBuffer.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioRingBuffer.swift
new file mode 100644
index 0000000..9aad18a
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioRingBuffer.swift
@@ -0,0 +1,75 @@
+import Foundation
+import os
+
+/// Thread-safe ring buffer for passing audio between the audio capture thread and the MLX
+/// inference thread. Writes drop oldest data when full; reads return zeros on underrun.
+///
+/// Uses `os_unfair_lock` for priority inheritance — safe to call `write` from a real-time
+/// Core Audio I/O thread without risking priority inversion.
+public final class AudioRingBuffer: @unchecked Sendable {
+ private var buffer: [Float]
+ private var readPos = 0
+ private var writePos = 0
+ private var count = 0
+ private var _lock = os_unfair_lock()
+ private let capacity: Int
+
+ public init(capacity: Int) {
+ self.capacity = capacity
+ self.buffer = [Float](repeating: 0, count: capacity)
+ }
+
+ /// Called from audio capture thread — non-blocking; drops oldest data if full.
+ public func write(_ samples: [Float]) {
+ os_unfair_lock_lock(&_lock)
+ defer { os_unfair_lock_unlock(&_lock) }
+ for sample in samples {
+ if count == capacity {
+ // Drop oldest sample
+ readPos = (readPos + 1) % capacity
+ count -= 1
+ }
+ buffer[writePos] = sample
+ writePos = (writePos + 1) % capacity
+ count += 1
+ }
+ }
+
+ /// Zero-copy write from a raw pointer — preferred on real-time audio threads
+ /// to avoid heap allocation from `Array(UnsafeBufferPointer(...))`.
+ public func write(from pointer: UnsafePointer, count sampleCount: Int) {
+ os_unfair_lock_lock(&_lock)
+ defer { os_unfair_lock_unlock(&_lock) }
+ for i in 0.. [Float] {
+ os_unfair_lock_lock(&_lock)
+ defer { os_unfair_lock_unlock(&_lock) }
+ var result = [Float](repeating: 0, count: n)
+ let available = min(n, count)
+ for i in 0.. MLComputeUnits {
+ guard let raw = ProcessInfo.processInfo.environment[envKey]?
+ .trimmingCharacters(in: .whitespacesAndNewlines)
+ .lowercased(), !raw.isEmpty
+ else {
+ return fallback
+ }
+ switch raw {
+ case "ane", "cpuandneuralengine", "neuralengine":
+ return .cpuAndNeuralEngine
+ case "gpu", "cpuandgpu":
+ return .cpuAndGPU
+ case "cpu", "cpuonly":
+ return .cpuOnly
+ case "all":
+ return .all
+ default:
+ return fallback
+ }
+ }
+}
+#endif
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/CoreMLLoader.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/CoreMLLoader.swift
new file mode 100644
index 0000000..35e0933
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/CoreMLLoader.swift
@@ -0,0 +1,125 @@
+import CoreML
+import Foundation
+#if canImport(os)
+import os
+#endif
+
+/// CoreML model loader that surfaces Neural Engine fallback.
+///
+/// `MLModel(contentsOf:configuration:)` silently succeeds when
+/// ``MILCompilerForANE`` fails — the model just runs on CPU instead of
+/// ANE. Users only see the performance cliff: RTF jumps from ~0.04 to
+/// ~1.8 on wake-word, ASR slows 5–20×, etc. They have no way to
+/// correlate this with the CoreML runtime's ``E5RT encountered an STL
+/// exception. msg = MILCompilerForANE error`` stderr message.
+///
+/// This helper:
+/// 1. Times the load.
+/// 2. Logs a single structured line per model with name + compute
+/// units + elapsed ms.
+/// 3. When the requested compute units include `.cpuAndNeuralEngine`
+/// (or `.all`) and the load completes faster than a typical ANE
+/// compile, emits a one-time warning pointing users at the
+/// fallback diagnostic.
+///
+/// Usage:
+/// ```swift
+/// let encoder = try CoreMLLoader.load(
+/// url: cacheDir.appendingPathComponent("encoder.mlmodelc"),
+/// computeUnits: .cpuAndNeuralEngine,
+/// name: "parakeet-eou-encoder"
+/// )
+/// ```
+public enum CoreMLLoader {
+
+ /// Seconds under which an ANE-eligible load is considered suspicious
+ /// (likely CPU fallback). Calibrated against observed behaviour:
+ /// - Successful ANE compile: ~200–800 ms on cold cache, ~20–50 ms
+ /// cached.
+ /// - CPU fallback after ANE compile failure: <10 ms regardless of
+ /// cache state.
+ ///
+ /// Picking 15 ms keeps false positives low on warm caches while
+ /// still catching the silent-fallback case on cold systems.
+ private static let aneCompileFloorSeconds: Double = 0.015
+
+ /// Track which model names we've already warned about so we don't
+ /// spam the log. Protected by ``warnedQueue``.
+ private static var warnedNames = Set()
+ private static let warnedQueue = DispatchQueue(
+ label: "com.qwen3speech.coreml-loader.warned"
+ )
+
+ /// Load a compiled CoreML model with instrumentation.
+ public static func load(
+ url: URL,
+ computeUnits: MLComputeUnits,
+ name: String? = nil
+ ) throws -> MLModel {
+ let config = MLModelConfiguration()
+ config.computeUnits = computeUnits
+ return try load(url: url, configuration: config, name: name)
+ }
+
+ /// Load with an explicit ``MLModelConfiguration``.
+ public static func load(
+ url: URL,
+ configuration: MLModelConfiguration,
+ name: String? = nil
+ ) throws -> MLModel {
+ // Honor the SPEECH_COREML_COMPUTE_UNITS override (CI forces cpuOnly to
+ // skip the runner's hanging ANE/GPU first-load compile). No-op on device.
+ configuration.computeUnits = CoreMLComputeUnitsResolver.resolved(
+ default: configuration.computeUnits)
+ let label = name ?? url.deletingPathExtension().lastPathComponent
+ let unitsLabel = describe(units: configuration.computeUnits)
+ let start = Date()
+ let model = try MLModel(contentsOf: url, configuration: configuration)
+ let elapsed = Date().timeIntervalSince(start)
+ let ms = Int((elapsed * 1000).rounded())
+ AudioLog.modelLoading.info("CoreML loaded \(label) in \(ms)ms (units=\(unitsLabel))")
+
+ let aneEligible =
+ configuration.computeUnits == .cpuAndNeuralEngine ||
+ configuration.computeUnits == .all
+ if aneEligible && elapsed < aneCompileFloorSeconds {
+ maybeWarn(
+ name: label,
+ message: """
+ CoreML model '\(label)' loaded in \(ms)ms with compute units \
+ \(unitsLabel). This is faster than a typical Neural Engine \
+ compile (~200–800 ms cold, ~20–50 ms cached). If console logs \
+ show 'MILCompilerForANE error', the model has fallen back to \
+ CPU and inference may be 5–20× slower than expected.
+ """
+ )
+ }
+ return model
+ }
+
+ // MARK: - Private
+
+ private static func maybeWarn(name: String, message: String) {
+ warnedQueue.sync {
+ guard !warnedNames.contains(name) else { return }
+ warnedNames.insert(name)
+ AudioLog.modelLoading.warning("\(message)")
+ }
+ }
+
+ private static func describe(units: MLComputeUnits) -> String {
+ switch units {
+ case .cpuOnly: return "cpuOnly"
+ case .cpuAndGPU: return "cpuAndGPU"
+ case .all: return "all"
+ case .cpuAndNeuralEngine: return "cpuAndNeuralEngine"
+ @unknown default: return "unknown(\(units.rawValue))"
+ }
+ }
+
+ /// Reset the per-process warning set. Exposed for tests so a fresh
+ /// run of the helper can emit a warning again.
+ public static func resetWarningState() {
+ warnedQueue.sync { warnedNames.removeAll() }
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/HuggingFaceDownloader.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/HuggingFaceDownloader.swift
new file mode 100644
index 0000000..cc8d95f
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/HuggingFaceDownloader.swift
@@ -0,0 +1,428 @@
+import Foundation
+import Hub
+import os
+
+/// Download errors
+public enum DownloadError: Error, LocalizedError {
+ case failedToDownload(String)
+ case invalidRemoteFileName(String)
+ /// A download attempt made no progress for `seconds` and was aborted
+ /// so the caller's retry loop can fire instead of hanging.
+ case stalled(modelId: String, seconds: Int)
+
+ public var errorDescription: String? {
+ switch self {
+ case .failedToDownload(let file):
+ return "Failed to download: \(file)"
+ case .invalidRemoteFileName(let file):
+ return "Refusing to write unsafe remote file name: \(file)"
+ case .stalled(let modelId, let seconds):
+ return "Download stalled for \(modelId): no progress in \(seconds)s"
+ }
+ }
+}
+
+/// HuggingFace model downloader — shared between ASR, TTS, VAD, etc.
+///
+/// Uses `HubApi` from the swift-transformers `Hub` module for downloads,
+/// which provides HF token auth and metadata tracking. Files that finished
+/// downloading are skipped on retry (etag/commit-hash check), but a file
+/// interrupted mid-transfer restarts from byte 0 — there is no usable
+/// mid-file resume in the current Hub stack, which is why the stall guard
+/// and retry ladder below favor patience over fast abort.
+public enum HuggingFaceDownloader {
+
+ // MARK: - Cache Directory
+
+ /// Get cache directory for a model.
+ ///
+ /// Returns the old flat cache path if it already contains model files (preserving
+ /// ~10 GB of existing cached models), otherwise returns the new Hub-style path.
+ public static func getCacheDirectory(for modelId: String, basePath: URL? = nil, cacheDirName: String = "qwen3-speech") throws -> URL {
+ let base = basePath ?? resolveBaseCacheDir(cacheDirName: cacheDirName)
+ let fm = FileManager.default
+
+ // Check old (flat) cache path for backward compat:
+ // ~/Library/Caches/qwen3-speech/aufklarer_Qwen3-ASR-0.6B-MLX-4bit/
+ let oldDir = base.appendingPathComponent(sanitizedCacheKey(for: modelId), isDirectory: true)
+ if weightsExist(in: oldDir) {
+ return oldDir
+ }
+
+ // New Hub-style path:
+ // ~/Library/Caches/qwen3-speech/models/aufklarer/Qwen3-ASR-0.6B-MLX-4bit/
+ let hub = HubApi(downloadBase: base)
+ let repo = Hub.Repo(id: modelId)
+ let dir = hub.localRepoLocation(repo)
+ try fm.createDirectory(at: dir, withIntermediateDirectories: true)
+ return dir
+ }
+
+ // MARK: - Weight Existence Check
+
+ /// Extensions recognised as cached model weights: the canonical
+ /// HF `.safetensors` layout plus Apple CoreML bundle directories
+ /// (`.mlmodelc`, `.mlpackage`) shipped by CoreML-only repos.
+ public static let weightFileExtensions: Set = [
+ "safetensors", "mlmodelc", "mlpackage"
+ ]
+
+ /// Returns `true` when `directory` contains at least one entry
+ /// whose extension matches `weightFileExtensions`. Used by
+ /// `downloadWeights` to short-circuit network requests when
+ /// `offlineMode: true` is set on caches that contain only CoreML
+ /// bundles and no `.safetensors` files.
+ public static func weightsExist(in directory: URL) -> Bool {
+ let fm = FileManager.default
+ guard fm.fileExists(atPath: directory.path) else { return false }
+ let contents: [URL]
+ do {
+ contents = try fm.contentsOfDirectory(at: directory, includingPropertiesForKeys: nil)
+ } catch {
+ AudioLog.download.debug("Could not list directory \(directory.path): \(error)")
+ contents = []
+ }
+ return contents.contains { weightFileExtensions.contains($0.pathExtension) }
+ }
+
+ // MARK: - Download
+
+ /// Download model files from HuggingFace using `HubApi.snapshot()`.
+ ///
+ /// Builds glob patterns from the file list:
+ /// - Always includes `config.json`
+ /// - If `additionalFiles` doesn't contain `.safetensors` files, adds `*.safetensors`
+ /// and `model.safetensors.index.json` to discover sharded weights automatically
+ /// - All entries in `additionalFiles` are added as-is (they work as glob patterns)
+ public static func downloadWeights(
+ modelId: String,
+ to directory: URL,
+ additionalFiles: [String] = [],
+ offlineMode: Bool = false,
+ hubEndpoint: String? = nil,
+ retryDelaysSeconds: [Int]? = nil,
+ progressHandler: ((Double) -> Void)? = nil
+ ) async throws {
+ // Skip network requests when weights are already cached
+ if offlineMode && weightsExist(in: directory) {
+ progressHandler?(1.0)
+ return
+ }
+
+ prepareRepoDirectoryForDownload(at: directory)
+
+ var globs: [String] = ["config.json"]
+
+ let hasExplicitWeights = additionalFiles.contains { $0.hasSuffix(".safetensors") }
+ if !hasExplicitWeights {
+ globs.append("*.safetensors")
+ globs.append("model.safetensors.index.json")
+ }
+ for file in additionalFiles where !globs.contains(file) {
+ globs.append(file)
+ }
+
+ // Derive the download base from the directory.
+ // getCacheDirectory returns either:
+ // old: base/cacheKey (flat, already has weights — won't reach here)
+ // new: base/models/org/model (Hub-style)
+ // For Hub API we need `base` as downloadBase.
+ //
+ // Forward `offlineMode` explicitly so HubApi doesn't fall through to
+ // its internal NWPathMonitor auto-detect, which on macOS can briefly
+ // report `.unsatisfied` and then refuse to download (manifesting as
+ // "Offline mode error: No files available locally for this repository"
+ // for a freshly-requested model).
+ let hub = makeHubApi(for: modelId, repoDir: directory, offlineMode: offlineMode, hubEndpoint: hubEndpoint)
+ let repo = Hub.Repo(id: modelId)
+
+ // Retry with capped backoff — HuggingFace can timeout on slow
+ // connections or rate-limit, and flaky networks (hotspots, captive
+ // portals) drop out for minutes at a time. Each attempt is wrapped
+ // in a progress-stall guard so a wedged mid-transfer (which
+ // `hub.snapshot` won't surface on its own) aborts and retries
+ // instead of hanging until the CI job is killed.
+ //
+ // No retries in offline mode: the failure is a deterministic local
+ // cache miss, and 110 s of backoff can't change what's on disk.
+ let delays = offlineMode ? [] : (retryDelaysSeconds ?? downloadRetryDelaysSeconds)
+ let maxAttempts = delays.count + 1
+ var lastError: Error?
+ for attempt in 1...maxAttempts {
+ do {
+ try await withDownloadStallGuard(modelId: modelId) { reportProgress in
+ try await hub.snapshot(from: repo, matching: globs) { progress in
+ reportProgress(progress.fractionCompleted)
+ progressHandler?(progress.fractionCompleted)
+ }
+ }
+ return // Success
+ } catch {
+ lastError = error
+ if isRecoverableHubCacheError(error) {
+ prepareRepoDirectoryForDownload(at: directory, force: true)
+ }
+ if attempt < maxAttempts {
+ try await Task.sleep(for: .seconds(delays[attempt - 1]))
+ }
+ }
+ }
+ throw DownloadError.failedToDownload(
+ "\(modelId) after \(maxAttempts) attempt\(maxAttempts == 1 ? "" : "s") "
+ + "(target: \(directory.path)): "
+ + (lastError?.localizedDescription ?? "unknown"))
+ }
+
+ /// Download an explicit list of files from HuggingFace without adding any
+ /// implicit weight globs. This is useful for overlaying tokenizer or config
+ /// assets from a second repository on top of an existing cache.
+ public static func downloadFiles(
+ modelId: String,
+ to directory: URL,
+ files: [String],
+ offlineMode: Bool = false,
+ hubEndpoint: String? = nil,
+ retryDelaysSeconds: [Int]? = nil,
+ progressHandler: ((Double) -> Void)? = nil
+ ) async throws {
+ if files.isEmpty {
+ progressHandler?(1.0)
+ return
+ }
+
+ prepareRepoDirectoryForDownload(at: directory)
+
+ let hub = makeHubApi(for: modelId, repoDir: directory, offlineMode: offlineMode, hubEndpoint: hubEndpoint)
+ let repo = Hub.Repo(id: modelId)
+
+ let globs = files.map { $0 }
+ // Same retry semantics as downloadWeights, including the offline
+ // no-retry rule — keep the two loops in lockstep.
+ let delays = offlineMode ? [] : (retryDelaysSeconds ?? downloadRetryDelaysSeconds)
+ let maxAttempts = delays.count + 1
+ var lastError: Error?
+ for attempt in 1...maxAttempts {
+ do {
+ try await withDownloadStallGuard(modelId: modelId) { reportProgress in
+ try await hub.snapshot(from: repo, matching: globs) { progress in
+ reportProgress(progress.fractionCompleted)
+ progressHandler?(progress.fractionCompleted)
+ }
+ }
+ return
+ } catch {
+ lastError = error
+ if isRecoverableHubCacheError(error) {
+ prepareRepoDirectoryForDownload(at: directory, force: true)
+ }
+ if attempt < maxAttempts {
+ try await Task.sleep(for: .seconds(delays[attempt - 1]))
+ }
+ }
+ }
+ throw DownloadError.failedToDownload(
+ "\(modelId) after \(maxAttempts) attempt\(maxAttempts == 1 ? "" : "s") "
+ + "(target: \(directory.path)): "
+ + (lastError?.localizedDescription ?? "unknown"))
+ }
+
+ // MARK: - Retry ladder
+
+ /// Delays between download attempts. One more attempt than entries:
+ /// 5 attempts with 5/15/30/60 s pauses (~110 s of backoff on top of the
+ /// per-attempt stall patience). Generous on purpose — abandoned attempts
+ /// restart files from byte 0 with the current Hub stack, so the cheap
+ /// resource here is wall-clock, not bytes. A network that's down for a
+ /// couple of minutes (AP roam, hotspot sleep, captive-portal re-auth)
+ /// should not kill a 2.75 GB first-run download.
+ static let downloadRetryDelaysSeconds = [5, 15, 30, 60]
+
+ /// Total attempts per download (retries + the initial try).
+ static var downloadMaxAttempts: Int { downloadRetryDelaysSeconds.count + 1 }
+
+ // MARK: - Download stall guard
+
+ /// Seconds of zero download progress after which an attempt is
+ /// considered wedged and aborted. `hub.snapshot` reports
+ /// `fractionCompleted` continuously while bytes flow, so a healthy
+ /// (even slow) transfer keeps resetting the clock; only a genuinely
+ /// stalled connection trips this.
+ ///
+ /// The default is tuned for end users, not CI: aborted attempts restart
+ /// each file from byte 0 (the Hub stack's mid-file resume never engages
+ /// on a fresh download), so firing the guard on a connection that would
+ /// have recovered throws away every byte of that attempt. Flaky networks
+ /// — AP roams, captive-portal re-auth, hotspot sleep — routinely stall
+ /// for 1–3 minutes and then recover, hence 300 s. CI pins
+ /// `HF_DOWNLOAD_STALL_TIMEOUT=90` to keep failing fast (app users can't
+ /// set env vars; CI can).
+ static var downloadStallTimeoutSeconds: Int {
+ if let raw = ProcessInfo.processInfo.environment["HF_DOWNLOAD_STALL_TIMEOUT"],
+ let v = Int(raw), v > 0 {
+ return v
+ }
+ return 300
+ }
+
+ /// Thread-safe last-progress timestamp. `hub.snapshot`'s progress
+ /// callback may fire from a background queue, so guard with a lock.
+ private final class ProgressClock: @unchecked Sendable {
+ private let lock = NSLock()
+ private var last = Date()
+ func tick() { lock.lock(); last = Date(); lock.unlock() }
+ func idleSeconds() -> Double {
+ lock.lock(); defer { lock.unlock() }
+ return Date().timeIntervalSince(last)
+ }
+ }
+
+ /// Run a download `operation` that reports fractional progress, and
+ /// abort it if progress stalls for `downloadStallTimeoutSeconds`.
+ /// On stall the in-flight `hub.snapshot` task is cancelled (URLSession
+ /// honors cancellation) and `DownloadError.stalled` is thrown so the
+ /// caller's retry loop fires instead of hanging indefinitely.
+ static func withDownloadStallGuard(
+ modelId: String,
+ stallTimeoutSeconds: Int? = nil,
+ _ operation: @escaping (@escaping @Sendable (Double) -> Void) async throws -> Void
+ ) async throws {
+ let stall = stallTimeoutSeconds ?? downloadStallTimeoutSeconds
+ let clock = ProgressClock()
+
+ try await withThrowingTaskGroup(of: Void.self) { group in
+ group.addTask {
+ try await operation { _ in clock.tick() }
+ }
+ group.addTask {
+ // Poll on a fraction of the window so we detect a stall
+ // within ~stall..stall+pollStep seconds.
+ let pollStep = max(1, stall / 3)
+ while true {
+ try await Task.sleep(for: .seconds(pollStep))
+ if clock.idleSeconds() >= Double(stall) {
+ throw DownloadError.stalled(modelId: modelId, seconds: stall)
+ }
+ }
+ }
+ // Whichever finishes first wins; cancel the other (the poller
+ // on success, or the download on stall).
+ defer { group.cancelAll() }
+ try await group.next()
+ }
+ }
+
+ // MARK: - Security Helpers (kept for backward compat + security tests)
+
+ /// Convert an arbitrary modelId into a single, safe path component for on-disk caching.
+ public static func sanitizedCacheKey(for modelId: String) -> String {
+ let replaced = modelId.replacingOccurrences(of: "/", with: "_")
+
+ let allowed = CharacterSet(charactersIn: "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-")
+ var scalars: [UnicodeScalar] = []
+ scalars.reserveCapacity(replaced.unicodeScalars.count)
+ for s in replaced.unicodeScalars {
+ scalars.append(allowed.contains(s) ? s : "_")
+ }
+
+ var cleaned = String(String.UnicodeScalarView(scalars))
+ cleaned = cleaned.trimmingCharacters(in: CharacterSet(charactersIn: "._"))
+
+ if cleaned.isEmpty || cleaned == "." || cleaned == ".." {
+ cleaned = "model"
+ }
+
+ return cleaned
+ }
+
+ /// Validate that a remote file name is safe.
+ public static func validatedRemoteFileName(_ file: String) throws -> String {
+ let base = URL(fileURLWithPath: file).lastPathComponent
+ guard base == file else {
+ throw DownloadError.invalidRemoteFileName(file)
+ }
+ guard !base.isEmpty, !base.hasPrefix("."), !base.contains("..") else {
+ throw DownloadError.invalidRemoteFileName(file)
+ }
+ guard base.range(of: #"^[A-Za-z0-9._-]+$"#, options: .regularExpression) != nil else {
+ throw DownloadError.invalidRemoteFileName(file)
+ }
+ return base
+ }
+
+ /// Validate that a local path stays within the expected directory.
+ public static func validatedLocalPath(directory: URL, fileName: String) throws -> URL {
+ let local = directory.appendingPathComponent(fileName, isDirectory: false)
+ let dirPath = directory.standardizedFileURL.path
+ let localPath = local.standardizedFileURL.path
+ let prefix = dirPath.hasSuffix("/") ? dirPath : (dirPath + "/")
+ guard localPath.hasPrefix(prefix) else {
+ throw DownloadError.invalidRemoteFileName(fileName)
+ }
+ return local
+ }
+
+ // MARK: - Private Helpers
+
+ /// Remove a repo folder that has Hub metadata but no complete weights.
+ /// Stale partial caches trigger "File metadata must have been retrieved from server".
+ static func prepareRepoDirectoryForDownload(at directory: URL, force: Bool = false) {
+ let fm = FileManager.default
+ guard fm.fileExists(atPath: directory.path) else { return }
+ if !force && weightsExist(in: directory) { return }
+ try? fm.removeItem(at: directory)
+ try? fm.createDirectory(at: directory, withIntermediateDirectories: true)
+ }
+
+ private static func isRecoverableHubCacheError(_ error: Error) -> Bool {
+ let text = (error as? LocalizedError)?.errorDescription
+ ?? error.localizedDescription
+ return text.localizedCaseInsensitiveContains("metadata")
+ || text.localizedCaseInsensitiveContains("offline mode")
+ }
+
+ /// Resolve the base cache directory from env vars or system default.
+ private static func resolveBaseCacheDir(cacheDirName: String) -> URL {
+ let fm = FileManager.default
+ let root: URL
+ if let override = ProcessInfo.processInfo.environment["QWEN3_CACHE_DIR"],
+ !override.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty {
+ root = URL(fileURLWithPath: override, isDirectory: true)
+ } else if let override = ProcessInfo.processInfo.environment["QWEN3_ASR_CACHE_DIR"],
+ !override.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty {
+ // Legacy env var support
+ root = URL(fileURLWithPath: override, isDirectory: true)
+ } else {
+ root = fm.urls(for: .cachesDirectory, in: .userDomainMask).first!
+ }
+ return root.appendingPathComponent(cacheDirName, isDirectory: true)
+ }
+
+ /// Create a `HubApi` whose `downloadBase` is derived from the repo directory that
+ /// `getCacheDirectory` returned (strips the `models//` suffix).
+ ///
+ /// `offlineMode` is forwarded as `useOfflineMode` so callers get the mode
+ /// they asked for instead of relying on `NWPathMonitor` auto-detection,
+ /// which can spuriously report `.unsatisfied` on macOS.
+ private static func makeHubApi(
+ for modelId: String,
+ repoDir: URL,
+ offlineMode: Bool,
+ hubEndpoint: String?
+ ) -> HubApi {
+ // repoDir is base/models/org/model
+ // We need base
+ let repo = Hub.Repo(id: modelId)
+ let suffix = "/\(repo.type.rawValue)/\(repo.id)"
+ let repoDirPath = repoDir.path
+ let downloadBase: URL
+ if repoDirPath.hasSuffix(suffix) {
+ let basePath = String(repoDirPath.dropLast(suffix.count))
+ downloadBase = URL(fileURLWithPath: basePath, isDirectory: true)
+ } else {
+ // Fallback: old-style flat dir — use its parent as downloadBase.
+ // Hub won't match this path, so we derive base from env/defaults.
+ downloadBase = resolveBaseCacheDir(cacheDirName: repoDir.deletingLastPathComponent().lastPathComponent)
+ }
+ return HubApi(downloadBase: downloadBase, endpoint: hubEndpoint, useOfflineMode: offlineMode)
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Logging.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Logging.swift
new file mode 100644
index 0000000..204debb
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Logging.swift
@@ -0,0 +1,13 @@
+import os
+
+/// Centralized loggers for audio model subsystems.
+public enum AudioLog {
+ /// Logger for model weight loading and initialization.
+ public static let modelLoading = Logger(subsystem: "com.qwen3speech", category: "ModelLoading")
+ /// Logger for inference and generation.
+ public static let inference = Logger(subsystem: "com.qwen3speech", category: "Inference")
+ /// Logger for HuggingFace downloads and caching.
+ public static let download = Logger(subsystem: "com.qwen3speech", category: "Download")
+ /// Logger for voice pipeline events.
+ public static let pipeline = Logger(subsystem: "com.qwen3speech", category: "Pipeline")
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelLoader.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelLoader.swift
new file mode 100644
index 0000000..67583dd
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelLoader.swift
@@ -0,0 +1,175 @@
+import Foundation
+import os
+
+/// Loaded model set — holds references to all loaded models.
+public struct ModelSet {
+ public let vad: (any StreamingVADProvider)?
+ public let stt: (any SpeechRecognitionModel)?
+ public let tts: (any SpeechGenerationModel)?
+
+ public init(
+ vad: (any StreamingVADProvider)? = nil,
+ stt: (any SpeechRecognitionModel)? = nil,
+ tts: (any SpeechGenerationModel)? = nil
+ ) {
+ self.vad = vad
+ self.stt = stt
+ self.tts = tts
+ }
+}
+
+/// A model to load, with its factory closure and progress weight.
+public struct ModelSpec: Sendable {
+ let name: String
+ let weight: Double
+ let group: Int // 0 = parallel group 1, 1 = sequential group 2
+ let loader: @Sendable (_ progress: @escaping (Double, String) -> Void) async throws -> any Sendable
+
+ /// VAD model spec.
+ public static func vad(
+ _ factory: @escaping @Sendable (_ progress: @escaping (Double, String) -> Void) async throws -> any StreamingVADProvider
+ ) -> ModelSpec {
+ ModelSpec(name: "VAD", weight: 1, group: 0, loader: { progress in
+ try await factory(progress) as any Sendable
+ })
+ }
+
+ /// Speech-to-text model spec.
+ public static func stt(
+ _ factory: @escaping @Sendable (_ progress: @escaping (Double, String) -> Void) async throws -> any SpeechRecognitionModel
+ ) -> ModelSpec {
+ ModelSpec(name: "ASR", weight: 15, group: 0, loader: { progress in
+ try await factory(progress) as any Sendable
+ })
+ }
+
+ /// Text-to-speech model spec.
+ public static func tts(
+ _ factory: @escaping @Sendable (_ progress: @escaping (Double, String) -> Void) async throws -> any SpeechGenerationModel
+ ) -> ModelSpec {
+ ModelSpec(name: "TTS", weight: 20, group: 1, loader: { progress in
+ try await factory(progress) as any Sendable
+ })
+ }
+}
+
+/// Unified model loading orchestrator with aggregated progress.
+///
+/// Loads multiple speech models with coordinated progress reporting.
+/// Group 0 models (VAD, ASR) load in parallel; Group 1 (TTS) loads after
+/// to reduce peak memory.
+///
+/// ```swift
+/// let models = try await ModelLoader.load([
+/// .vad { p in try await SileroVADModel.fromPretrained(engine: .coreml, progressHandler: p) },
+/// .stt { p in try await ParakeetASRModel.fromPretrained(progressHandler: p) },
+/// .tts { p in try await KokoroTTSModel.fromPretrained(progressHandler: p) },
+/// ], onProgress: { progress, stage in
+/// self.loadProgress = progress
+/// self.loadingStatus = stage
+/// })
+/// // models.vad, models.stt, models.tts are ready
+/// ```
+public enum ModelLoader {
+
+ private static let log = Logger(subsystem: "audio.soniqo", category: "ModelLoader")
+
+ /// Load the requested models with aggregated progress reporting.
+ public static func load(
+ _ specs: [ModelSpec],
+ onProgress: @escaping @Sendable (_ progress: Double, _ stage: String) -> Void = { _, _ in }
+ ) async throws -> ModelSet {
+ let totalWeight = specs.reduce(0.0) { $0 + $1.weight }
+ guard totalWeight > 0 else { return ModelSet() }
+
+ let state = LoadState(totalWeight: totalWeight)
+
+ // Group 0: parallel (VAD + ASR)
+ let group0 = specs.filter { $0.group == 0 }
+ // Group 1: sequential after group 0 (TTS — heavy, reduce peak memory)
+ let group1 = specs.filter { $0.group != 0 }
+
+ var results: [(String, any Sendable)] = []
+
+ // Load group 0 in parallel
+ if !group0.isEmpty {
+ try await withThrowingTaskGroup(of: (String, any Sendable).self) { group in
+ for spec in group0 {
+ group.addTask {
+ let model = try await loadSpec(spec, state: state, onProgress: onProgress)
+ return (spec.name, model)
+ }
+ }
+ for try await result in group {
+ results.append(result)
+ }
+ }
+ }
+
+ // Load group 1 sequentially
+ for spec in group1 {
+ let model = try await loadSpec(spec, state: state, onProgress: onProgress)
+ results.append((spec.name, model))
+ }
+
+ onProgress(1.0, "Ready")
+ log.info("All models loaded")
+
+ // Build ModelSet from results
+ var vad: (any StreamingVADProvider)?
+ var stt: (any SpeechRecognitionModel)?
+ var tts: (any SpeechGenerationModel)?
+
+ for (_, model) in results {
+ if let m = model as? any StreamingVADProvider { vad = m }
+ if let m = model as? any SpeechRecognitionModel { stt = m }
+ if let m = model as? any SpeechGenerationModel { tts = m }
+ }
+
+ return ModelSet(vad: vad, stt: stt, tts: tts)
+ }
+
+ // MARK: - Internal
+
+ private final class LoadState: @unchecked Sendable {
+ let totalWeight: Double
+ private var completed: Double = 0
+ private let lock = NSLock()
+
+ init(totalWeight: Double) { self.totalWeight = totalWeight }
+
+ func addCompleted(_ w: Double) {
+ lock.lock(); completed += w; lock.unlock()
+ }
+
+ var completedFraction: Double {
+ lock.lock(); defer { lock.unlock() }
+ return completed / totalWeight
+ }
+
+ func overallProgress(specWeight: Double, localFraction: Double) -> Double {
+ lock.lock(); defer { lock.unlock() }
+ return (completed + localFraction * specWeight) / totalWeight
+ }
+ }
+
+ private static func loadSpec(
+ _ spec: ModelSpec,
+ state: LoadState,
+ onProgress: @escaping @Sendable (Double, String) -> Void
+ ) async throws -> any Sendable {
+ log.info("Loading \(spec.name)...")
+ onProgress(state.completedFraction, "\(spec.name)...")
+
+ let adapter: @Sendable (Double, String) -> Void = { fraction, status in
+ let overall = state.overallProgress(specWeight: spec.weight, localFraction: fraction)
+ let stage = status.isEmpty ? spec.name : "\(spec.name): \(status)"
+ onProgress(overall, stage)
+ }
+
+ let model = try await spec.loader(adapter)
+ state.addCompleted(spec.weight)
+ log.info("\(spec.name) loaded")
+ return model
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelRegistry.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelRegistry.swift
new file mode 100644
index 0000000..9196d08
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelRegistry.swift
@@ -0,0 +1,9 @@
+import Foundation
+
+/// Remote registry used when fetching on-device model weights.
+public enum ModelRegistry: Sendable, Equatable {
+ /// Official Hugging Face Hub (`swift-transformers` / `HubApi`).
+ case huggingFace(hubEndpoint: String? = nil)
+ /// ModelScope.cn — same `owner/model` ids as Hugging Face for aufklarer MLX repos.
+ case modelScope(baseURL: String = ModelScopeDownloader.defaultBaseURL, revision: String = "master")
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelScopeDownloader.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelScopeDownloader.swift
new file mode 100644
index 0000000..5c4a41b
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelScopeDownloader.swift
@@ -0,0 +1,335 @@
+import Foundation
+
+/// Downloads model files from [ModelScope](https://www.modelscope.cn) using the
+/// public repo API. Uses the same `owner/model` ids as Hugging Face for repos
+/// mirrored on ModelScope (e.g. `aufklarer/Qwen3-ASR-0.6B-MLX-4bit`).
+public enum ModelScopeDownloader {
+
+ public static let defaultBaseURL = "https://modelscope.cn"
+
+ private struct FilesPayload: Decodable {
+ struct Entry: Decodable {
+ let Path: String
+ let Size: Int64?
+ let entryType: String?
+
+ enum CodingKeys: String, CodingKey {
+ case Path
+ case Size
+ case entryType = "Type"
+ }
+ }
+ let Files: [Entry]
+ }
+
+ private struct APIResponse: Decodable {
+ let Data: FilesPayload
+ }
+
+ public struct RemoteFile: Sendable {
+ public let path: String
+ public let size: Int64
+ }
+
+ // MARK: - Public API
+
+ /// Mirror of `HuggingFaceDownloader.downloadWeights` for ModelScope.
+ public static func downloadWeights(
+ modelId: String,
+ to directory: URL,
+ additionalFiles: [String] = [],
+ baseURL: String = defaultBaseURL,
+ revision: String = "master",
+ retryDelaysSeconds: [Int]? = nil,
+ progressHandler: ((Double) -> Void)? = nil
+ ) async throws {
+ HuggingFaceDownloader.prepareRepoDirectoryForDownload(at: directory)
+
+ let listed = try await listAllFiles(modelId: modelId, baseURL: baseURL, revision: revision)
+ var selected = Set(["config.json"])
+ for file in additionalFiles {
+ selected.insert(file)
+ }
+
+ let hasExplicitWeights = additionalFiles.contains { $0.hasSuffix(".safetensors") }
+ if !hasExplicitWeights {
+ for file in listed where file.path.hasSuffix(".safetensors") {
+ selected.insert(file.path)
+ }
+ if listed.contains(where: { $0.path == "model.safetensors.index.json" }) {
+ selected.insert("model.safetensors.index.json")
+ }
+ }
+
+ let files = listed.filter { selected.contains($0.path) }.map(\.path)
+ guard !files.isEmpty else {
+ throw DownloadError.failedToDownload("\(modelId): no matching files on ModelScope")
+ }
+
+ try await downloadFiles(
+ modelId: modelId,
+ to: directory,
+ files: files,
+ fileSizes: Dictionary(uniqueKeysWithValues: listed.map { ($0.path, $0.size) }),
+ baseURL: baseURL,
+ revision: revision,
+ retryDelaysSeconds: retryDelaysSeconds,
+ progressHandler: progressHandler
+ )
+ }
+
+ /// Download an explicit list of repo-relative paths into `directory`.
+ public static func downloadFiles(
+ modelId: String,
+ to directory: URL,
+ files: [String],
+ fileSizes: [String: Int64] = [:],
+ baseURL: String = defaultBaseURL,
+ revision: String = "master",
+ retryDelaysSeconds: [Int]? = nil,
+ progressHandler: ((Double) -> Void)? = nil
+ ) async throws {
+ if files.isEmpty {
+ progressHandler?(1.0)
+ return
+ }
+
+ HuggingFaceDownloader.prepareRepoDirectoryForDownload(at: directory)
+
+ let ordered = files.sorted()
+ var sizes = fileSizes
+ for path in ordered where sizes[path] == nil {
+ sizes[path] = 0
+ }
+
+ // Without byte sizes the old logic fell back to `(index + 1) / count`,
+ // which jumps to 50% as soon as two small JSON files finish. Resolve
+ // sizes from the repo listing whenever any entry is missing.
+ if ordered.contains(where: { (sizes[$0] ?? 0) <= 0 }) {
+ let listed = try await listAllFiles(
+ modelId: modelId,
+ baseURL: baseURL,
+ revision: revision
+ )
+ let listedMap = Dictionary(uniqueKeysWithValues: listed.map { ($0.path, $0.size) })
+ for path in ordered where (sizes[path] ?? 0) <= 0 {
+ if let remote = listedMap[path], remote > 0 {
+ sizes[path] = remote
+ }
+ }
+ }
+
+ let totalBytes = max(ordered.reduce(Int64(0)) { $0 + (sizes[$1] ?? 0) }, 1)
+ var completedBytes: Int64 = 0
+
+ let delays = retryDelaysSeconds ?? HuggingFaceDownloader.downloadRetryDelaysSeconds
+ let maxAttempts = delays.count + 1
+
+ for (index, path) in ordered.enumerated() {
+ let destination = directory.appendingPathComponent(path, isDirectory: false)
+ try FileManager.default.createDirectory(
+ at: destination.deletingLastPathComponent(),
+ withIntermediateDirectories: true
+ )
+
+ var lastError: Error?
+ for attempt in 1...maxAttempts {
+ do {
+ try await HuggingFaceDownloader.withDownloadStallGuard(modelId: modelId) { reportProgress in
+ try await fetchFile(
+ modelId: modelId,
+ filePath: path,
+ to: destination,
+ baseURL: baseURL,
+ revision: revision
+ ) { fileBytes, fileExpectedBytes in
+ reportProgress(1.0)
+ let fileSize = sizes[path] ?? 0
+ let expected = fileSize > 0 ? fileSize : fileExpectedBytes
+ let overall: Double
+ if expected > 0, totalBytes > 1 {
+ overall = Double(completedBytes + min(fileBytes, expected)) / Double(totalBytes)
+ } else {
+ // Last resort when listing omits sizes: spread each
+ // file's slice by bytes received vs Content-Length.
+ let slice = 1.0 / Double(ordered.count)
+ let base = Double(index) * slice
+ let inFile = expected > 0
+ ? min(Double(fileBytes) / Double(expected), 1.0) * slice
+ : slice
+ overall = base + inFile
+ }
+ progressHandler?(min(max(overall, 0), 1))
+ }
+ }
+ lastError = nil
+ break
+ } catch {
+ lastError = error
+ try? FileManager.default.removeItem(at: destination)
+ if attempt < maxAttempts {
+ try await Task.sleep(for: .seconds(delays[attempt - 1]))
+ }
+ }
+ }
+
+ if let lastError {
+ throw DownloadError.failedToDownload(
+ "\(modelId)/\(path) on ModelScope: \(lastError.localizedDescription)"
+ )
+ }
+
+ completedBytes += sizes[path] ?? 0
+ progressHandler?(min(Double(completedBytes) / Double(totalBytes), 1))
+ }
+
+ progressHandler?(1.0)
+ }
+
+ // MARK: - Listing
+
+ /// Recursively lists every file in a ModelScope repo (used for CoreML bundles).
+ public static func listAllFiles(
+ modelId: String,
+ baseURL: String,
+ revision: String
+ ) async throws -> [RemoteFile] {
+ var collected: [RemoteFile] = []
+ try await listFiles(
+ modelId: modelId,
+ root: nil,
+ into: &collected,
+ baseURL: baseURL,
+ revision: revision
+ )
+ return collected
+ }
+
+ private static func listFiles(
+ modelId: String,
+ root: String?,
+ into collected: inout [RemoteFile],
+ baseURL: String,
+ revision: String
+ ) async throws {
+ guard let url = listingURL(modelId: modelId, baseURL: baseURL, revision: revision, root: root) else {
+ throw DownloadError.failedToDownload("Invalid ModelScope listing URL for \(modelId)")
+ }
+
+ let (data, response) = try await URLSession.shared.data(from: url)
+ guard let http = response as? HTTPURLResponse, (200...299).contains(http.statusCode) else {
+ throw DownloadError.failedToDownload("ModelScope listing failed for \(modelId)")
+ }
+
+ let payload = try JSONDecoder().decode(APIResponse.self, from: data)
+ for entry in payload.Data.Files {
+ if isDirectoryEntry(entry) {
+ try await listFiles(
+ modelId: modelId,
+ root: entry.Path,
+ into: &collected,
+ baseURL: baseURL,
+ revision: revision
+ )
+ } else {
+ collected.append(RemoteFile(path: entry.Path, size: entry.Size ?? 0))
+ }
+ }
+ }
+
+ private static func isDirectoryEntry(_ entry: FilesPayload.Entry) -> Bool {
+ if entry.entryType?.lowercased() == "tree" { return true }
+ let size = entry.Size ?? 0
+ return size == 0 && !entry.Path.contains(".")
+ }
+
+ // MARK: - Transfer
+
+ /// Streams a single repo file. `onBytes` receives `(bytesWritten, expectedBytes)`.
+ private static func fetchFile(
+ modelId: String,
+ filePath: String,
+ to destination: URL,
+ baseURL: String,
+ revision: String,
+ onBytes: @escaping (Int64, Int64) -> Void
+ ) async throws {
+ guard let url = fileURL(modelId: modelId, baseURL: baseURL, revision: revision, filePath: filePath) else {
+ throw DownloadError.invalidRemoteFileName(filePath)
+ }
+
+ var request = URLRequest(url: url)
+ request.timeoutInterval = 3600
+
+ let (asyncBytes, response) = try await URLSession.shared.bytes(for: request)
+ guard let http = response as? HTTPURLResponse else {
+ throw DownloadError.failedToDownload(filePath)
+ }
+ guard (200...299).contains(http.statusCode) else {
+ throw DownloadError.failedToDownload("\(filePath) HTTP \(http.statusCode)")
+ }
+
+ let expectedBytes = http.value(forHTTPHeaderField: "Content-Length")
+ .flatMap(Int64.init) ?? 0
+
+ if FileManager.default.fileExists(atPath: destination.path) {
+ try FileManager.default.removeItem(at: destination)
+ }
+ FileManager.default.createFile(atPath: destination.path, contents: nil)
+ let handle = try FileHandle(forWritingTo: destination)
+ defer { try? handle.close() }
+
+ var buffer = Data()
+ buffer.reserveCapacity(1_048_576)
+ var written: Int64 = 0
+
+ for try await byte in asyncBytes {
+ try Task.checkCancellation()
+ buffer.append(byte)
+ if buffer.count >= 1_048_576 {
+ try handle.write(contentsOf: buffer)
+ written += Int64(buffer.count)
+ buffer.removeAll(keepingCapacity: true)
+ onBytes(written, expectedBytes)
+ }
+ }
+ if !buffer.isEmpty {
+ try handle.write(contentsOf: buffer)
+ written += Int64(buffer.count)
+ }
+ onBytes(written, expectedBytes)
+ }
+
+ // MARK: - URLs
+
+ private static func listingURL(
+ modelId: String,
+ baseURL: String,
+ revision: String,
+ root: String?
+ ) -> URL? {
+ var components = URLComponents(string: "\(baseURL)/api/v1/models/\(modelId)/repo/files")
+ var items = [
+ URLQueryItem(name: "Revision", value: revision),
+ ]
+ if let root, !root.isEmpty {
+ items.append(URLQueryItem(name: "Root", value: root))
+ }
+ components?.queryItems = items
+ return components?.url
+ }
+
+ private static func fileURL(
+ modelId: String,
+ baseURL: String,
+ revision: String,
+ filePath: String
+ ) -> URL? {
+ var components = URLComponents(string: "\(baseURL)/api/v1/models/\(modelId)/repo")
+ components?.queryItems = [
+ URLQueryItem(name: "Revision", value: revision),
+ URLQueryItem(name: "FilePath", value: filePath),
+ ]
+ return components?.url
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/PipelineLLM.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/PipelineLLM.swift
new file mode 100644
index 0000000..17b77fc
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/PipelineLLM.swift
@@ -0,0 +1,53 @@
+// MARK: - LLM Protocol
+
+/// Protocol for language model integration with voice pipelines.
+///
+/// Conforming types bridge an LLM (local or remote) to the VoicePipeline's
+/// ASR → LLM → TTS flow. The pipeline calls `chat()` on a background thread
+/// and expects blocking behavior (return when generation is complete).
+public protocol PipelineLLM: AnyObject {
+ /// Generate a response given conversation messages.
+ ///
+ /// Called on the pipeline's worker thread (blocking). Emit tokens via
+ /// `onToken(text, isFinal)` — the pipeline forwards them to TTS.
+ func chat(messages: [(role: MessageRole, content: String)],
+ onToken: @escaping (String, Bool) -> Void)
+
+ /// Cancel in-progress generation. Thread-safe.
+ func cancel()
+}
+
+/// Message roles for LLM conversation.
+public enum MessageRole: Int, Sendable {
+ case system = 0
+ case user = 1
+ case assistant = 2
+ case tool = 3
+}
+
+// MARK: - Tool Calling
+
+/// A tool that can be invoked by the LLM during voice pipeline execution.
+public struct PipelineTool {
+ public let name: String
+ public let description: String
+ public let handler: (String) -> String
+ public let cooldown: Int
+
+ /// - Parameters:
+ /// - name: Tool name (used by LLM to invoke)
+ /// - description: What the tool does (included in LLM system prompt)
+ /// - cooldown: Minimum seconds between invocations (0 = no limit)
+ /// - handler: Synchronous handler `(arguments) -> result`. Called on pipeline worker thread.
+ public init(
+ name: String,
+ description: String,
+ cooldown: Int = 0,
+ handler: @escaping (String) -> String
+ ) {
+ self.name = name
+ self.description = description
+ self.cooldown = cooldown
+ self.handler = handler
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Protocols.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Protocols.swift
new file mode 100644
index 0000000..0762968
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Protocols.swift
@@ -0,0 +1,282 @@
+import Foundation
+
+// MARK: - Model Memory Management
+
+/// Memory statistics for a loaded model.
+public struct ModelMemoryStats: Sendable {
+ /// Estimated weight memory in bytes
+ public let weightMemory: Int
+ /// Current active GPU memory in bytes (MLX only)
+ public let activeMemory: Int
+
+ public init(weightMemory: Int, activeMemory: Int = 0) {
+ self.weightMemory = weightMemory
+ self.activeMemory = activeMemory
+ }
+}
+
+/// A model that supports explicit memory management.
+///
+/// Call `unload()` to release model weights and free GPU memory.
+/// After unloading, the model cannot be used for inference until re-loaded.
+public protocol ModelMemoryManageable: AnyObject {
+ /// Whether the model is currently loaded and ready for inference.
+ var isLoaded: Bool { get }
+
+ /// Release model weights and free GPU memory.
+ ///
+ /// After calling this, `isLoaded` returns false and inference methods will fail.
+ /// To use the model again, create a new instance via `fromPretrained()`.
+ func unload()
+
+ /// Estimated memory footprint of the loaded model weights in bytes.
+ /// Returns 0 if the model is not loaded.
+ var memoryFootprint: Int { get }
+}
+
+// MARK: - Unified Audio Chunk
+
+/// A chunk of audio produced during streaming synthesis or generation.
+public struct AudioChunk: Sendable {
+ /// PCM audio samples (Float32)
+ public let samples: [Float]
+ /// Sample rate in Hz (e.g. 24000)
+ public let sampleRate: Int
+ /// Index of the first frame in this chunk
+ public let frameIndex: Int
+ /// True if this is the last chunk
+ public let isFinal: Bool
+ /// Wall-clock seconds since generation started (nil if not tracked)
+ public let elapsedTime: Double?
+ /// Text tokens generated alongside audio (populated on final chunk if available)
+ public let textTokens: [Int32]
+
+ public init(
+ samples: [Float],
+ sampleRate: Int,
+ frameIndex: Int,
+ isFinal: Bool,
+ elapsedTime: Double? = nil,
+ textTokens: [Int32] = []
+ ) {
+ self.samples = samples
+ self.sampleRate = sampleRate
+ self.frameIndex = frameIndex
+ self.isFinal = isFinal
+ self.elapsedTime = elapsedTime
+ self.textTokens = textTokens
+ }
+}
+
+// MARK: - Aligned Word
+
+/// A word with its aligned start and end timestamps (in seconds).
+public struct AlignedWord: Sendable {
+ public let text: String
+ public let startTime: Float
+ public let endTime: Float
+
+ public init(text: String, startTime: Float, endTime: Float) {
+ self.text = text
+ self.startTime = startTime
+ self.endTime = endTime
+ }
+}
+
+// MARK: - Speech Generation (TTS)
+
+/// A text-to-speech model that generates audio from text.
+public protocol SpeechGenerationModel: AnyObject {
+ /// Output sample rate in Hz
+ var sampleRate: Int { get }
+ /// Synthesize audio from text (returns full waveform)
+ func generate(text: String, language: String?) async throws -> [Float]
+ /// Synthesize audio from text with streaming output.
+ /// Default implementation wraps `generate()` as a single chunk.
+ func generateStream(text: String, language: String?) -> AsyncThrowingStream
+}
+
+extension SpeechGenerationModel {
+ /// Default: wraps `generate()` as a single-chunk stream.
+ public func generateStream(text: String, language: String?) -> AsyncThrowingStream {
+ let rate = sampleRate
+ return AsyncThrowingStream { continuation in
+ Task {
+ do {
+ let samples = try await self.generate(text: text, language: language)
+ continuation.yield(AudioChunk(samples: samples, sampleRate: rate, frameIndex: 0, isFinal: true))
+ continuation.finish()
+ } catch {
+ continuation.finish(throwing: error)
+ }
+ }
+ }
+ }
+}
+
+// MARK: - Speech Recognition (STT)
+
+/// A word with its confidence score.
+public struct WordConfidence: Sendable {
+ public let word: String
+ /// Confidence score (0.0–1.0) derived from mean token log-probability.
+ public let confidence: Float
+
+ public init(word: String, confidence: Float) {
+ self.word = word
+ self.confidence = confidence
+ }
+}
+
+/// Result of speech recognition including detected language.
+public struct TranscriptionResult: Sendable {
+ public let text: String
+ /// Detected language (e.g. "english", "russian"). Nil if model doesn't detect.
+ public let language: String?
+ /// Confidence score (0.0–1.0). Higher = more confident transcription.
+ /// Derived from average token log-probability. 0.0 if model doesn't provide.
+ public let confidence: Float
+ /// Per-word confidence scores. Nil if model doesn't provide.
+ public let words: [WordConfidence]?
+
+ public init(text: String, language: String? = nil, confidence: Float = 0.0, words: [WordConfidence]? = nil) {
+ self.text = text
+ self.language = language
+ self.confidence = confidence
+ self.words = words
+ }
+}
+
+/// A speech-to-text model that transcribes audio.
+public protocol SpeechRecognitionModel: AnyObject {
+ /// Expected input sample rate in Hz
+ var inputSampleRate: Int { get }
+ /// Transcribe audio to text
+ func transcribe(audio: [Float], sampleRate: Int, language: String?) -> String
+ /// Transcribe audio to text with language detection
+ func transcribeWithLanguage(audio: [Float], sampleRate: Int, language: String?) -> TranscriptionResult
+}
+
+/// Default implementation: delegates to transcribe() with no language detection.
+public extension SpeechRecognitionModel {
+ func transcribeWithLanguage(audio: [Float], sampleRate: Int, language: String?) -> TranscriptionResult {
+ TranscriptionResult(text: transcribe(audio: audio, sampleRate: sampleRate, language: language))
+ }
+}
+
+// MARK: - Forced Alignment
+
+/// A model that aligns text to audio at the word level.
+public protocol ForcedAlignmentModel: AnyObject {
+ /// Align text to audio, returning word-level timestamps
+ func align(audio: [Float], text: String, sampleRate: Int, language: String?) -> [AlignedWord]
+}
+
+// MARK: - Speech-to-Speech
+
+/// A speech-to-speech model that generates a spoken response to spoken input.
+public protocol SpeechToSpeechModel: AnyObject {
+ /// Output sample rate in Hz
+ var sampleRate: Int { get }
+ /// Generate response audio from input audio (blocking)
+ func respond(userAudio: [Float]) -> [Float]
+ /// Generate response audio from input audio with streaming output
+ func respondStream(userAudio: [Float]) -> AsyncThrowingStream
+}
+
+// MARK: - Voice Activity Detection
+
+/// A time segment where speech was detected.
+public struct SpeechSegment: Sendable {
+ /// Start time in seconds
+ public let startTime: Float
+ /// End time in seconds
+ public let endTime: Float
+
+ public init(startTime: Float, endTime: Float) {
+ self.startTime = startTime
+ self.endTime = endTime
+ }
+
+ /// Duration in seconds
+ public var duration: Float { endTime - startTime }
+}
+
+/// A model that detects speech activity regions in audio.
+public protocol VoiceActivityDetectionModel: AnyObject {
+ /// Expected input sample rate in Hz
+ var inputSampleRate: Int { get }
+ /// Detect speech segments in audio
+ func detectSpeech(audio: [Float], sampleRate: Int) -> [SpeechSegment]
+}
+
+/// A streaming VAD that processes fixed-size audio chunks and returns speech probability.
+///
+/// Maps directly to speech-core's `sc_vad_vtable_t` for pipeline integration.
+public protocol StreamingVADProvider: AnyObject {
+ /// Expected input sample rate in Hz
+ var inputSampleRate: Int { get }
+ /// Number of samples per chunk
+ var chunkSize: Int { get }
+ /// Process a single audio chunk, returns speech probability in [0, 1]
+ func processChunk(_ samples: [Float]) -> Float
+ /// Reset internal state (LSTM hidden state, context buffer, etc.)
+ func resetState()
+}
+
+// MARK: - Speaker Diarization
+
+/// A speech segment with an assigned speaker identity.
+public struct DiarizedSegment: Sendable {
+ /// Start time in seconds
+ public let startTime: Float
+ /// End time in seconds
+ public let endTime: Float
+ /// Speaker identifier (0-based)
+ public let speakerId: Int
+
+ public init(startTime: Float, endTime: Float, speakerId: Int) {
+ self.startTime = startTime
+ self.endTime = endTime
+ self.speakerId = speakerId
+ }
+
+ /// Duration in seconds
+ public var duration: Float { endTime - startTime }
+}
+
+/// A model that produces speaker embeddings from audio.
+public protocol SpeakerEmbeddingModel: AnyObject {
+ /// Expected input sample rate in Hz
+ var inputSampleRate: Int { get }
+ /// Embedding vector dimension
+ var embeddingDimension: Int { get }
+ /// Extract a speaker embedding from audio
+ func embed(audio: [Float], sampleRate: Int) -> [Float]
+}
+
+// MARK: - Speech Enhancement
+
+/// A model that enhances speech by removing noise.
+public protocol SpeechEnhancementModel: AnyObject {
+ /// Expected input sample rate in Hz
+ var inputSampleRate: Int { get }
+ /// Enhance audio by removing noise
+ func enhance(audio: [Float], sampleRate: Int) throws -> [Float]
+}
+
+/// A model that assigns speaker identities to speech segments.
+public protocol SpeakerDiarizationModel: AnyObject {
+ /// Expected input sample rate in Hz
+ var inputSampleRate: Int { get }
+ /// Diarize audio into speaker-labeled segments
+ func diarize(audio: [Float], sampleRate: Int) -> [DiarizedSegment]
+}
+
+/// A diarization model that also supports extracting a specific speaker's segments
+/// using a reference embedding. Not all engines support this (e.g. Sortformer is
+/// end-to-end and does not produce speaker embeddings).
+public protocol SpeakerExtractionCapable: SpeakerDiarizationModel {
+ /// Extract segments belonging to a target speaker identified by a reference embedding.
+ func extractSpeaker(audio: [Float], sampleRate: Int, targetEmbedding: [Float]) -> [SpeechSegment]
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/SentencePieceModel.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/SentencePieceModel.swift
new file mode 100644
index 0000000..28b9211
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/SentencePieceModel.swift
@@ -0,0 +1,182 @@
+import Foundation
+
+/// Minimal SentencePiece `.model` (`sentencepiece_model.proto`) reader.
+///
+/// Extracts the vocabulary list — `(text, score, type)` for every piece —
+/// without requiring a protobuf runtime dependency. Modules build their own
+/// encode/decode logic on top: this struct only owns the wire-format parse
+/// and the raw piece array.
+///
+/// `sentencepiece_model.proto` excerpt:
+/// ```
+/// message ModelProto {
+/// repeated SentencePiece pieces = 1; // field 1, length-delimited submsg
+/// ...
+/// }
+/// message SentencePiece {
+/// optional string piece = 1; // field 1, length-delimited string
+/// optional float score = 2; // field 2, fixed32 (wire type 5)
+/// optional Type type = 3; // field 3, varint (wire type 0)
+/// }
+/// ```
+public struct SentencePieceModel: Sendable {
+
+ /// Piece type constants from `sentencepiece_model.proto`. Values not in
+ /// this enum are surfaced as `.unknown(rawValue)` so callers can apply
+ /// their own special-token handling.
+ public enum PieceType: Int32, Sendable {
+ case normal = 1
+ case unknown = 2
+ case control = 3
+ case userDefined = 4
+ case unused = 5
+ case byte = 6
+ }
+
+ public struct Piece: Sendable, Equatable {
+ public let text: String
+ public let score: Float
+ public let type: Int32
+
+ public init(text: String, score: Float, type: Int32) {
+ self.text = text
+ self.score = score
+ self.type = type
+ }
+
+ public var pieceType: PieceType? { PieceType(rawValue: type) }
+
+ public var isControlOrUnknown: Bool {
+ type == PieceType.control.rawValue ||
+ type == PieceType.unknown.rawValue ||
+ type == PieceType.unused.rawValue ||
+ type == PieceType.byte.rawValue
+ }
+ }
+
+ public let pieces: [Piece]
+
+ public var count: Int { pieces.count }
+
+ public subscript(_ id: Int) -> Piece? {
+ guard id >= 0, id < pieces.count else { return nil }
+ return pieces[id]
+ }
+
+ public init(contentsOf url: URL) throws {
+ let data = try Data(contentsOf: url)
+ try self.init(data: data)
+ }
+
+ public init(modelPath: String) throws {
+ try self.init(contentsOf: URL(fileURLWithPath: modelPath))
+ }
+
+ public init(data: Data) throws {
+ var parsed: [Piece] = []
+ var offset = 0
+
+ while offset < data.count {
+ let (fieldNumber, wireType, afterTag) = Self.readTag(data: data, offset: offset)
+ offset = afterTag
+
+ // Top-level field 1 = repeated SentencePiece, length-delimited (wire 2)
+ guard fieldNumber == 1, wireType == 2 else {
+ offset = Self.skipField(data: data, offset: offset, wireType: wireType)
+ continue
+ }
+
+ let (length, afterLen) = Self.readVarint(data: data, offset: afterTag)
+ offset = afterLen
+ let end = offset + length
+
+ var piece = ""
+ var score: Float = 0
+ var type: Int32 = PieceType.normal.rawValue
+
+ var sub = offset
+ while sub < end {
+ let (subField, subWire, afterSubTag) = Self.readTag(data: data, offset: sub)
+ sub = afterSubTag
+ switch (subField, subWire) {
+ case (1, 2): // piece string
+ let (strLen, afterStrLen) = Self.readVarint(data: data, offset: sub)
+ sub = afterStrLen
+ if let s = String(data: data[sub..<(sub + strLen)], encoding: .utf8) {
+ piece = s
+ }
+ sub += strLen
+ case (2, 5): // score (fixed32 / wire type 5)
+ score = data[sub..<(sub + 4)].withUnsafeBytes { $0.loadUnaligned(as: Float.self) }
+ sub += 4
+ case (3, 0): // type varint
+ let (typeValue, afterType) = Self.readVarint(data: data, offset: sub)
+ sub = afterType
+ type = Int32(typeValue)
+ default:
+ sub = Self.skipField(data: data, offset: sub, wireType: subWire)
+ }
+ }
+
+ parsed.append(Piece(text: piece, score: score, type: type))
+ offset = end
+ }
+
+ guard !parsed.isEmpty else {
+ throw SentencePieceModelError.emptyModel
+ }
+ self.pieces = parsed
+ }
+
+ // MARK: - Protobuf wire helpers
+
+ private static func readVarint(data: Data, offset: Int) -> (value: Int, newOffset: Int) {
+ var result = 0
+ var shift = 0
+ var off = offset
+ while off < data.count {
+ let byte = Int(data[off])
+ off += 1
+ result |= (byte & 0x7F) << shift
+ if byte & 0x80 == 0 { break }
+ shift += 7
+ }
+ return (result, off)
+ }
+
+ private static func readTag(data: Data, offset: Int) -> (fieldNumber: Int, wireType: Int, newOffset: Int) {
+ let (tag, newOffset) = readVarint(data: data, offset: offset)
+ return (tag >> 3, tag & 0x07, newOffset)
+ }
+
+ private static func skipField(data: Data, offset: Int, wireType: Int) -> Int {
+ switch wireType {
+ case 0:
+ let (_, newOffset) = readVarint(data: data, offset: offset)
+ return newOffset
+ case 1:
+ return offset + 8
+ case 2:
+ let (length, newOffset) = readVarint(data: data, offset: offset)
+ return newOffset + length
+ case 5:
+ return offset + 4
+ default:
+ return data.count
+ }
+ }
+}
+
+public enum SentencePieceModelError: Error, CustomStringConvertible {
+ case emptyModel
+ case invalidFile(URL)
+
+ public var description: String {
+ switch self {
+ case .emptyModel:
+ return "SentencePiece model contained no pieces"
+ case .invalidFile(let url):
+ return "Could not read SentencePiece model at \(url.path)"
+ }
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/StreamingAudioPlayer.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/StreamingAudioPlayer.swift
new file mode 100644
index 0000000..5be907e
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/StreamingAudioPlayer.swift
@@ -0,0 +1,511 @@
+#if canImport(AVFoundation)
+import AVFoundation
+import os
+
+/// Lock-free SPSC ring buffer for audio samples.
+/// Producer (TTS thread) writes, consumer (audio render thread) reads.
+public final class AudioSampleRingBuffer: @unchecked Sendable {
+ private let buffer: UnsafeMutableBufferPointer
+ private let capacity: Int
+ private var writePos: Int = 0 // only written by producer
+ private var readPos: Int = 0 // only written by consumer
+
+ public init(capacity: Int) {
+ self.capacity = capacity
+ let ptr = UnsafeMutablePointer.allocate(capacity: capacity)
+ ptr.initialize(repeating: 0, count: capacity)
+ self.buffer = UnsafeMutableBufferPointer(start: ptr, count: capacity)
+ }
+
+ deinit {
+ buffer.baseAddress?.deinitialize(count: capacity)
+ buffer.baseAddress?.deallocate()
+ }
+
+ /// Number of samples available to read.
+ public var availableToRead: Int {
+ let w = writePos
+ let r = readPos
+ return w >= r ? w - r : capacity - r + w
+ }
+
+ /// Number of free slots for writing.
+ public var availableToWrite: Int {
+ return capacity - availableToRead - 1
+ }
+
+ /// Write samples into the buffer. Returns number actually written.
+ @discardableResult
+ public func write(_ samples: [Float]) -> Int {
+ let count = min(samples.count, availableToWrite)
+ guard count > 0 else { return 0 }
+
+ samples.withUnsafeBufferPointer { src in
+ let w = writePos
+ let firstChunk = min(count, capacity - w)
+ buffer.baseAddress!.advanced(by: w).update(from: src.baseAddress!, count: firstChunk)
+ if firstChunk < count {
+ buffer.baseAddress!.update(from: src.baseAddress!.advanced(by: firstChunk), count: count - firstChunk)
+ }
+ }
+ writePos = (writePos + count) % capacity
+ return count
+ }
+
+ /// Read samples from the buffer into dst. Returns number actually read.
+ @discardableResult
+ public func read(into dst: UnsafeMutablePointer, count: Int) -> Int {
+ let available = min(count, availableToRead)
+ guard available > 0 else { return 0 }
+
+ let r = readPos
+ let firstChunk = min(available, capacity - r)
+ dst.update(from: buffer.baseAddress!.advanced(by: r), count: firstChunk)
+ if firstChunk < available {
+ dst.advanced(by: firstChunk).update(from: buffer.baseAddress!, count: available - firstChunk)
+ }
+ readPos = (readPos + available) % capacity
+ return available
+ }
+
+ /// Reset both pointers (call when not actively reading/writing).
+ public func reset() {
+ readPos = 0
+ writePos = 0
+ }
+}
+
+/// Streams TTS audio via AVAudioEngine using an event-driven render callback.
+///
+/// Architecture:
+/// ```
+/// TTS (producer) → [Ring Buffer] → AVAudioSourceNode render callback (consumer)
+/// pre-fill N sec hardware pulls when it needs data
+/// ```
+///
+/// The render thread calls our callback when it needs audio. We read from the
+/// ring buffer. If the buffer is empty (underflow), we output silence.
+///
+/// `preBufferDuration` controls how much audio must accumulate before playback
+/// starts. This is the latency-quality tradeoff:
+/// - Higher = more resilient to TTS jitter, but more latency
+/// - Lower = less latency, but risk of underflow gaps
+///
+/// Typical values:
+/// - 0s: single-pass TTS (Kokoro) where all audio arrives at once
+/// - 2s: streaming TTS (Qwen3-TTS, RTF ~0.53)
+public final class StreamingAudioPlayer: @unchecked Sendable {
+ private var engine: AVAudioEngine?
+ private var sourceNode: AVAudioSourceNode?
+ private var format: AVAudioFormat?
+ private let lock = NSLock()
+
+ private var ringBuffer: AudioSampleRingBuffer?
+ private var playbackStarted = false
+ private var generationComplete = false
+ private var isFirstChunk = true
+ private var upsampler: AVAudioConverter?
+ private var preBufferSamples: Int = 0
+ public private(set) var totalWritten: Int = 0
+ /// Number of samples written for external diagnostics.
+ public var totalWrittenSamples: Int { totalWritten }
+ private var totalRead: Int = 0
+
+ public private(set) var isPlaying = false
+ private var playbackFinishedFired = false
+
+ /// Pre-buffer duration in seconds. Playback starts after this much audio accumulates.
+ /// Default 1.0s — sufficient for streaming TTS at RTF < 0.6.
+ public var preBufferDuration: Double = 1.0
+
+ /// Callback when all audio has finished playing.
+ public var onPlaybackFinished: (() -> Void)?
+
+ /// Ring buffer capacity in seconds. Default 30s — enough for any TTS response.
+ public var ringBufferDuration: Double = 30
+
+ public init() {}
+
+ // MARK: - Standalone mode
+
+ /// Start playback engine at the given sample rate.
+ public func start(sampleRate: Double = 24000) throws {
+ stop()
+ let eng = AVAudioEngine()
+ guard let fmt = AVAudioFormat(
+ commonFormat: .pcmFormatFloat32,
+ sampleRate: sampleRate,
+ channels: 1,
+ interleaved: false
+ ) else { return }
+
+ setupSourceNode(engine: eng, format: fmt)
+ try eng.start()
+ self.engine = eng
+ self.format = fmt
+ }
+
+ /// Create a standalone engine at the hardware's native sample rate.
+ public func ensureStandaloneEngine() {
+ guard sourceNode == nil else { return }
+ let eng = AVAudioEngine()
+ let mixerFormat = eng.mainMixerNode.outputFormat(forBus: 0)
+ guard let monoFormat = AVAudioFormat(
+ commonFormat: .pcmFormatFloat32,
+ sampleRate: mixerFormat.sampleRate,
+ channels: 1,
+ interleaved: false
+ ) else { return }
+ setupSourceNode(engine: eng, format: monoFormat)
+ do {
+ try eng.start()
+ self.engine = eng
+ self.format = monoFormat
+ } catch {}
+ }
+
+ // MARK: - Attached mode
+
+ /// Attach to an existing AVAudioEngine.
+ public func attach(to engine: AVAudioEngine, format: AVAudioFormat) {
+ setupSourceNode(engine: engine, format: format)
+ self.format = format
+ }
+
+ /// Start the source node (for use when attaching before engine.start()).
+ public func startPlayback() {
+ // Source node is always running once attached — no-op
+ }
+
+ /// Detach from an external engine.
+ public func detach(from engine: AVAudioEngine) {
+ if let node = sourceNode {
+ engine.disconnectNodeOutput(node)
+ engine.detach(node)
+ }
+ sourceNode = nil
+ format = nil
+ upsampler = nil
+ ringBuffer?.reset()
+ }
+
+ // MARK: - Audio Scheduling
+
+ /// Write a chunk of audio samples into the ring buffer.
+ /// If pre-buffer threshold is reached, playback begins automatically.
+ public func scheduleChunk(_ samples: [Float]) {
+ guard !samples.isEmpty else { return }
+
+ var output = samples
+
+ // Drop near-silent warmup chunks at start of generation
+ if isFirstChunk {
+ var sumSq: Float = 0
+ for s in samples { sumSq += s * s }
+ let rms = sqrt(sumSq / Float(samples.count))
+ if rms < 0.005 { return } // Only drop near-silence (codec init noise)
+ isFirstChunk = false
+ // 5ms fade-in to prevent pop
+ if let fmt = format {
+ let fadeFrames = min(samples.count, Int(fmt.sampleRate * 0.005))
+ for i in 0.. 0 {
+ if (ringBuffer?.availableToRead ?? 0) >= preBufferSamples {
+ playbackStarted = true
+ }
+ } else if preBufferSamples == 0 {
+ playbackStarted = true
+ }
+ lock.unlock()
+ }
+
+ /// Write samples with resampling from sourceSampleRate to the player's rate.
+ public func play(samples: [Float], sampleRate: Int) throws {
+ guard let fmt = format else { return }
+ if Double(sampleRate) == fmt.sampleRate {
+ scheduleChunk(samples)
+ } else {
+ guard let srcFmt = AVAudioFormat(commonFormat: .pcmFormatFloat32, sampleRate: Double(sampleRate), channels: 1, interleaved: false) else { return }
+ if upsampler == nil || upsampler?.inputFormat.sampleRate != Double(sampleRate) {
+ upsampler = AVAudioConverter(from: srcFmt, to: fmt)
+ }
+ guard let converter = upsampler else { return }
+ guard let inputBuffer = AVAudioPCMBuffer(pcmFormat: srcFmt, frameCapacity: AVAudioFrameCount(samples.count)) else { return }
+ inputBuffer.frameLength = AVAudioFrameCount(samples.count)
+ samples.withUnsafeBufferPointer { ptr in
+ inputBuffer.floatChannelData![0].update(from: ptr.baseAddress!, count: samples.count)
+ }
+ let outFrameCount = AVAudioFrameCount(Double(samples.count) * fmt.sampleRate / Double(sampleRate))
+ guard let outputBuffer = AVAudioPCMBuffer(pcmFormat: fmt, frameCapacity: outFrameCount) else { return }
+ var consumed = false
+ var error: NSError?
+ converter.convert(to: outputBuffer, error: &error) { _, outStatus in
+ if consumed { outStatus.pointee = .noDataNow; return nil }
+ consumed = true
+ outStatus.pointee = .haveData
+ return inputBuffer
+ }
+ let count = Int(outputBuffer.frameLength)
+ guard count > 0, let data = outputBuffer.floatChannelData else { return }
+ let resampled = Array(UnsafeBufferPointer(start: data[0], count: count))
+ scheduleChunk(resampled)
+ }
+ }
+
+ // MARK: - Completion
+
+ /// Signal that TTS generation is complete — no more chunks will arrive.
+ /// The render callback will drain remaining samples, then fire onPlaybackFinished.
+ public func markGenerationComplete() {
+ lock.lock()
+ generationComplete = true
+ playbackStarted = true
+ let hasEngine = sourceNode != nil
+ let empty = (ringBuffer?.availableToRead ?? 0) == 0
+ let written = totalWritten
+ lock.unlock()
+
+ // No engine or nothing was written — fire immediately
+ if !hasEngine || (empty && written == 0) {
+ guard !playbackFinishedFired else { return }
+ playbackFinishedFired = true
+ isPlaying = false
+ onPlaybackFinished?()
+ return
+ }
+
+ // Start polling: the render callback normally fires onPlaybackFinished
+ // when the buffer drains, but if the render thread isn't running (e.g.
+ // simulator, or audio route change), we poll the buffer to detect
+ // completion reliably. Works on both device and simulator.
+ startCompletionPolling()
+ }
+
+ private var completionPollTimer: DispatchSourceTimer?
+ private var lastPolledRead: Int = 0
+ private var noProgressPolls: Int = 0
+
+ private func startCompletionPolling() {
+ completionPollTimer?.cancel()
+ lastPolledRead = -1
+ noProgressPolls = 0
+ let timer = DispatchSource.makeTimerSource(queue: .main)
+ timer.schedule(deadline: .now() + 0.2, repeating: 0.2)
+ timer.setEventHandler { [weak self] in
+ guard let self else { return }
+ // Already fired by render callback — stop polling
+ guard !self.playbackFinishedFired else {
+ self.completionPollTimer?.cancel()
+ self.completionPollTimer = nil
+ return
+ }
+ self.lock.lock()
+ let complete = self.generationComplete
+ let remaining = self.ringBuffer?.availableToRead ?? 0
+ let read = self.totalRead
+ let written = self.totalWritten
+ self.lock.unlock()
+
+ // All samples consumed (or render thread never started reading)
+ let drained = remaining == 0 && read >= written && written > 0
+ // Render thread never started — audio engine not running
+ let stalled = complete && read == 0 && written > 0
+
+ // Render thread stalled mid-stream (partial read, no progress for
+ // 3 consecutive polls = 600 ms). Seen on virtualized macOS CI runners
+ // and on real iOS when an audio-session interrupt freezes the
+ // render thread between buffers.
+ if complete && read > 0 && read < written {
+ if read == self.lastPolledRead {
+ self.noProgressPolls += 1
+ } else {
+ self.noProgressPolls = 0
+ self.lastPolledRead = read
+ }
+ }
+ let frozen = complete && self.noProgressPolls >= 3 && read > 0 && read < written
+
+ if complete && (drained || stalled || frozen) {
+ self.completionPollTimer?.cancel()
+ self.completionPollTimer = nil
+ guard !self.playbackFinishedFired else { return }
+ self.playbackFinishedFired = true
+ self.isPlaying = false
+ self.onPlaybackFinished?()
+ }
+ }
+ completionPollTimer = timer
+ timer.resume()
+ }
+
+ /// Reset for a new generation cycle.
+ public func resetGeneration() {
+ completionPollTimer?.cancel()
+ completionPollTimer = nil
+ lastPolledRead = -1
+ noProgressPolls = 0
+ lock.lock()
+ generationComplete = false
+ playbackFinishedFired = false
+ playbackStarted = false
+ isFirstChunk = true
+ totalWritten = 0
+ totalRead = 0
+ ringBuffer?.reset()
+ lock.unlock()
+ }
+
+ /// Wait until all audio has finished playing.
+ public func waitForCompletion() async {
+ while isPlaying {
+ try? await Task.sleep(nanoseconds: 50_000_000) // 50ms poll
+ }
+ }
+
+ /// Stop immediately.
+ public func fadeOutAndStop() {
+ lock.lock()
+ generationComplete = false
+ playbackStarted = false
+ isFirstChunk = true
+ totalWritten = 0
+ totalRead = 0
+ ringBuffer?.reset()
+ lock.unlock()
+ isPlaying = false
+ }
+
+ /// Stop and release resources.
+ public func stop() {
+ completionPollTimer?.cancel()
+ completionPollTimer = nil
+ if let eng = engine, let node = sourceNode {
+ eng.disconnectNodeOutput(node)
+ eng.detach(node)
+ }
+ engine?.stop()
+ engine = nil
+ sourceNode = nil
+ format = nil
+ upsampler = nil
+ lock.lock()
+ generationComplete = false
+ playbackStarted = false
+ isFirstChunk = true
+ totalWritten = 0
+ totalRead = 0
+ ringBuffer?.reset()
+ lock.unlock()
+ isPlaying = false
+ }
+
+ // MARK: - Private
+
+ private func setupSourceNode(engine: AVAudioEngine, format: AVAudioFormat) {
+ let bufferCapacity = Int(format.sampleRate * ringBufferDuration)
+ let rb = AudioSampleRingBuffer(capacity: bufferCapacity)
+ self.ringBuffer = rb
+ self.preBufferSamples = Int(format.sampleRate * preBufferDuration)
+
+ let node = AVAudioSourceNode(format: format) { [weak self] _, _, frameCount, bufferList -> OSStatus in
+ guard let self else { return noErr }
+
+ let ablPointer = UnsafeMutableAudioBufferListPointer(bufferList)
+ guard let dst = ablPointer[0].mData?.assumingMemoryBound(to: Float.self) else {
+ return noErr
+ }
+ let frames = Int(frameCount)
+
+ self.lock.lock()
+ let started = self.playbackStarted
+ let complete = self.generationComplete
+ let available = rb.availableToRead
+ self.lock.unlock()
+
+ if !started {
+ // Pre-buffer not full yet — output silence
+ dst.update(repeating: 0, count: frames)
+ return noErr
+ }
+
+ if available > 0 {
+ let read = rb.read(into: dst, count: min(frames, available))
+ // Zero-fill remainder if not enough
+ if read < frames {
+ dst.advanced(by: read).update(repeating: 0, count: frames - read)
+ }
+ self.lock.lock()
+ self.totalRead += read
+ self.lock.unlock()
+ } else if complete && !self.playbackFinishedFired {
+ // Buffer empty + generation done = playback finished (fire once)
+ self.playbackFinishedFired = true
+ dst.update(repeating: 0, count: frames)
+ DispatchQueue.main.async {
+ self.isPlaying = false
+ self.onPlaybackFinished?()
+ }
+ } else {
+ // Underflow — output silence, keep waiting for more data
+ dst.update(repeating: 0, count: frames)
+ }
+
+ return noErr
+ }
+
+ engine.attach(node)
+ engine.connect(node, to: engine.mainMixerNode, format: format)
+ self.sourceNode = node
+ }
+
+ /// Compress long silent gaps to at most `maxSilence` samples.
+ /// TTS models produce long pauses between sentences (500ms+).
+ /// This shortens them while keeping a natural brief pause.
+ static func compressSilence(_ samples: [Float], maxSilence: Int, threshold: Float) -> [Float] {
+ guard samples.count > maxSilence else { return samples }
+
+ var result = [Float]()
+ result.reserveCapacity(samples.count)
+ var silenceRun = 0
+
+ // Process in small frames (240 samples = 10ms at 24kHz)
+ let frameSize = 240
+ var offset = 0
+
+ while offset < samples.count {
+ let end = min(offset + frameSize, samples.count)
+ let frame = samples[offset..text) and basic BPE encoding (text->ids) via merges.txt
+public class Qwen3Tokenizer {
+ private var idToToken: [Int: String] = [:]
+ private var tokenToId: [String: Int] = [:]
+ private var bpeMerges: [(String, String)] = []
+ private var bpeMergeRanks: [String: Int] = [:]
+
+ public var eosTokenId: Int = 151643
+ public var padTokenId: Int = 151643
+ public var bosTokenId: Int = 151644
+
+ public init() {}
+
+ /// Test-only initializer with pre-built token mappings
+ internal init(idToToken: [Int: String]) {
+ self.idToToken = idToToken
+ for (id, token) in idToToken { tokenToId[token] = id }
+ }
+
+ /// Load tokenizer from vocab.json file (direct token->id mapping)
+ public func load(from url: URL) throws {
+ let data = try Data(contentsOf: url)
+
+ // vocab.json is a direct {token: id} mapping
+ guard let vocab = try JSONSerialization.jsonObject(with: data) as? [String: Int] else {
+ throw TokenizerError.invalidFormat("Expected {token: id} dictionary")
+ }
+
+ for (token, id) in vocab {
+ idToToken[id] = token
+ tokenToId[token] = id
+ }
+
+ // Also load added tokens from tokenizer_config.json if it exists
+ let configUrl = url.deletingLastPathComponent().appendingPathComponent("tokenizer_config.json")
+ if FileManager.default.fileExists(atPath: configUrl.path) {
+ try loadAddedTokens(from: configUrl)
+ }
+
+ // Load BPE merges if available
+ let mergesUrl = url.deletingLastPathComponent().appendingPathComponent("merges.txt")
+ if FileManager.default.fileExists(atPath: mergesUrl.path) {
+ try loadMerges(from: mergesUrl)
+ }
+
+ logTokenizer("Loaded tokenizer with \(idToToken.count) tokens, \(bpeMerges.count) merges")
+ }
+
+ /// Load added tokens from tokenizer_config.json
+ private func loadAddedTokens(from url: URL) throws {
+ let data = try Data(contentsOf: url)
+
+ guard let config = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
+ return // Not a valid config, skip
+ }
+
+ // added_tokens_decoder is a dict with string keys (token IDs) and object values with "content" field
+ if let addedTokens = config["added_tokens_decoder"] as? [String: [String: Any]] {
+ var addedCount = 0
+ for (idString, tokenInfo) in addedTokens {
+ guard let id = Int(idString),
+ let content = tokenInfo["content"] as? String else {
+ continue
+ }
+
+ // Add to our mappings (overwrite if exists)
+ idToToken[id] = content
+ tokenToId[content] = id
+ addedCount += 1
+ }
+ logTokenizer("Loaded \(addedCount) added tokens from tokenizer_config.json")
+ }
+ }
+
+ /// Load BPE merge rules from merges.txt
+ private func loadMerges(from url: URL) throws {
+ let content = try String(contentsOf: url, encoding: .utf8)
+ let lines = content.components(separatedBy: .newlines)
+
+ for (index, line) in lines.enumerated() {
+ // Skip header line and empty lines
+ if line.hasPrefix("#") || line.isEmpty { continue }
+
+ let parts = line.components(separatedBy: " ")
+ guard parts.count == 2 else { continue }
+
+ bpeMerges.append((parts[0], parts[1]))
+ bpeMergeRanks["\(parts[0]) \(parts[1])"] = index
+ }
+ }
+
+ /// Decode token IDs to text using a unified byte buffer.
+ /// Collects all bytes before converting to UTF-8, so multi-byte characters
+ /// split across BPE tokens (e.g. CJK) decode correctly.
+ public func decode(tokens: [Int]) -> String {
+ var buffer: [UInt8] = []
+
+ for tokenId in tokens {
+ guard let token = idToToken[tokenId] else { continue }
+
+ // Skip <|...|> special tokens
+ if token.hasPrefix("<|") && token.hasSuffix("|>") {
+ continue
+ }
+
+ // Keep and similar markers — append their UTF-8 bytes
+ if token.hasPrefix("<") && token.hasSuffix(">") && !token.contains("|") {
+ buffer.append(contentsOf: Array(token.utf8))
+ continue
+ }
+
+ // Convert each char via unicodeToByte (Ġ→0x20 space is handled
+ // automatically since unicodeToByte maps Ġ (U+0120) → byte 32)
+ for char in token {
+ if let byte = Self.unicodeToByte[char] {
+ buffer.append(byte)
+ } else {
+ buffer.append(contentsOf: String(char).utf8)
+ }
+ }
+ }
+
+ let text = String(bytes: buffer, encoding: .utf8)
+ ?? String(decoding: buffer, as: UTF8.self)
+ return text.trimmingCharacters(in: .whitespaces)
+ }
+
+ /// Byte-to-unicode mapping table (GPT-2 style)
+ /// Built lazily on first use
+ private static var byteToUnicode: [UInt8: Character] = {
+ var mapping: [UInt8: Character] = [:]
+ var n = 0
+
+ // Printable ASCII and some extended chars map directly
+ let ranges: [(ClosedRange)] = [
+ (UInt8(ascii: "!")...UInt8(ascii: "~")), // 33-126
+ (0xA1...0xAC), // 161-172
+ (0xAE...0xFF), // 174-255
+ ]
+
+ for range in ranges {
+ for b in range {
+ mapping[b] = Character(UnicodeScalar(b))
+ }
+ }
+
+ // Remaining bytes (0-32, 127-160, 173) map to U+0100 onwards
+ for b: UInt8 in 0...255 {
+ if mapping[b] == nil {
+ mapping[b] = Character(UnicodeScalar(0x100 + n)!)
+ n += 1
+ }
+ }
+
+ return mapping
+ }()
+
+ /// Unicode-to-byte reverse mapping
+ private static var unicodeToByte: [Character: UInt8] = {
+ var reverse: [Character: UInt8] = [:]
+ for (byte, char) in byteToUnicode {
+ reverse[char] = byte
+ }
+ return reverse
+ }()
+
+ /// Encode a byte-level BPE token string from raw text bytes
+ private func encodeByteLevelToken(_ text: String) -> String {
+ var result = ""
+ for byte in text.utf8 {
+ if let char = Self.byteToUnicode[byte] {
+ result.append(char)
+ }
+ }
+ return result
+ }
+
+ /// BPE encode text to token IDs
+ public func encode(_ text: String) -> [Int] {
+ guard !bpeMerges.isEmpty else {
+ // Fallback: character-level encoding
+ return characterEncode(text)
+ }
+
+ // Split text into words (whitespace-aware, GPT-2 style pre-tokenization)
+ // Simple approach: split on word boundaries, preserving leading spaces as Ġ
+ let words = preTokenize(text)
+
+ var tokens: [Int] = []
+ for word in words {
+ // Convert word to byte-level BPE representation
+ let bpeTokens = bpe(word)
+ for bpeToken in bpeTokens {
+ if let id = tokenToId[bpeToken] {
+ tokens.append(id)
+ }
+ }
+ }
+
+ return tokens
+ }
+
+ /// Pre-tokenize text into words (GPT-2 style)
+ private func preTokenize(_ text: String) -> [String] {
+ // Split on whitespace boundaries while preserving leading spaces as part of the next word
+ var words: [String] = []
+ var current = ""
+
+ for char in text {
+ if char == " " || char == "\n" || char == "\t" {
+ if !current.isEmpty {
+ words.append(encodeByteLevelToken(current))
+ current = ""
+ }
+ current = String(char)
+ } else {
+ current.append(char)
+ }
+ }
+ if !current.isEmpty {
+ words.append(encodeByteLevelToken(current))
+ }
+
+ return words
+ }
+
+ /// Apply BPE merges to a word
+ private func bpe(_ word: String) -> [String] {
+ var pieces = word.map { String($0) }
+
+ while pieces.count > 1 {
+ // Find the pair with lowest merge rank
+ var bestPair: (String, String)?
+ var bestRank = Int.max
+
+ for i in 0..<(pieces.count - 1) {
+ let pair = "\(pieces[i]) \(pieces[i + 1])"
+ if let rank = bpeMergeRanks[pair], rank < bestRank {
+ bestRank = rank
+ bestPair = (pieces[i], pieces[i + 1])
+ }
+ }
+
+ guard let (first, second) = bestPair else { break }
+
+ // Merge the pair
+ var newPieces: [String] = []
+ var i = 0
+ while i < pieces.count {
+ if i < pieces.count - 1 && pieces[i] == first && pieces[i + 1] == second {
+ newPieces.append(first + second)
+ i += 2
+ } else {
+ newPieces.append(pieces[i])
+ i += 1
+ }
+ }
+ pieces = newPieces
+ }
+
+ return pieces
+ }
+
+ /// Simple character-level encoding fallback
+ private func characterEncode(_ text: String) -> [Int] {
+ var tokens: [Int] = []
+ for char in text {
+ if let id = tokenToId[String(char)] {
+ tokens.append(id)
+ }
+ }
+ return tokens
+ }
+
+ /// Get token ID for a specific token string
+ public func getTokenId(for token: String) -> Int? {
+ return tokenToId[token]
+ }
+
+ /// Get token string for a specific ID
+ public func getToken(for id: Int) -> String? {
+ return idToToken[id]
+ }
+
+ /// Debug: print token mappings for common words
+ public func debugTokenMappings() {
+ let commonTokens = [
+ "<|im_start|>", "<|im_end|>", "<|audio_start|>", "<|audio_end|>",
+ "<|audio_pad|>", "", "<|endoftext|>",
+ "system", "user", "assistant", "language", "English",
+ "Ġsystem", "Ġuser", "Ġassistant", "Ġlanguage", "ĠEnglish",
+ "\n", "Ċ" // newline representations
+ ]
+
+ print("Token ID mappings:")
+ for token in commonTokens {
+ if let id = tokenToId[token] {
+ print(" '\(token)' -> \(id)")
+ } else {
+ print(" '\(token)' -> NOT FOUND")
+ }
+ }
+ }
+}
+
+/// Protocol for tokenizer to allow different implementations
+public protocol TokenizerProtocol {
+ func decode(tokens: [Int]) -> String
+ func encode(_ text: String) -> [Int]
+}
+
+extension Qwen3Tokenizer: TokenizerProtocol {}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/WAVWriter.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/WAVWriter.swift
new file mode 100644
index 0000000..d190329
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/WAVWriter.swift
@@ -0,0 +1,105 @@
+import Foundation
+
+/// Write float audio samples to WAV file
+public enum WAVWriter {
+
+ /// Write mono float samples to a 16-bit PCM WAV file
+ /// - Parameters:
+ /// - samples: Float audio samples in [-1.0, 1.0] range
+ /// - sampleRate: Sample rate in Hz (default 24000)
+ /// - url: Output file URL
+ public static func write(samples: [Float], sampleRate: Int = 24000, to url: URL) throws {
+ let numChannels: UInt16 = 1
+ let bitsPerSample: UInt16 = 16
+ let bytesPerSample = Int(bitsPerSample) / 8
+ let dataSize = samples.count * bytesPerSample
+ let fileSize = 36 + dataSize
+
+ var data = Data(capacity: fileSize + 8)
+
+ // RIFF header
+ data.append(contentsOf: "RIFF".utf8)
+ appendUInt32(&data, UInt32(fileSize))
+ data.append(contentsOf: "WAVE".utf8)
+
+ // fmt chunk
+ data.append(contentsOf: "fmt ".utf8)
+ appendUInt32(&data, 16) // chunk size
+ appendUInt16(&data, 1) // PCM format
+ appendUInt16(&data, numChannels)
+ appendUInt32(&data, UInt32(sampleRate))
+ appendUInt32(&data, UInt32(sampleRate * Int(numChannels) * bytesPerSample)) // byte rate
+ appendUInt16(&data, numChannels * UInt16(bytesPerSample)) // block align
+ appendUInt16(&data, bitsPerSample)
+
+ // data chunk
+ data.append(contentsOf: "data".utf8)
+ appendUInt32(&data, UInt32(dataSize))
+
+ // Convert float samples to 16-bit PCM
+ for sample in samples {
+ let clamped = max(-1.0, min(1.0, sample))
+ let int16Value = Int16(clamped * 32767.0)
+ appendInt16(&data, int16Value)
+ }
+
+ try data.write(to: url)
+ }
+
+ /// Write stereo float samples to a 16-bit PCM WAV file.
+ /// - Parameters:
+ /// - left: Left channel float samples in [-1.0, 1.0]
+ /// - right: Right channel float samples in [-1.0, 1.0]
+ /// - sampleRate: Sample rate in Hz
+ /// - url: Output file URL
+ public static func writeStereo(left: [Float], right: [Float], sampleRate: Int = 44100, to url: URL) throws {
+ let numChannels: UInt16 = 2
+ let bitsPerSample: UInt16 = 16
+ let bytesPerSample = Int(bitsPerSample) / 8
+ let frameCount = min(left.count, right.count)
+ let dataSize = frameCount * Int(numChannels) * bytesPerSample
+ let fileSize = 36 + dataSize
+
+ var data = Data(capacity: fileSize + 8)
+
+ data.append(contentsOf: "RIFF".utf8)
+ appendUInt32(&data, UInt32(fileSize))
+ data.append(contentsOf: "WAVE".utf8)
+
+ data.append(contentsOf: "fmt ".utf8)
+ appendUInt32(&data, 16)
+ appendUInt16(&data, 1) // PCM
+ appendUInt16(&data, numChannels)
+ appendUInt32(&data, UInt32(sampleRate))
+ appendUInt32(&data, UInt32(sampleRate * Int(numChannels) * bytesPerSample))
+ appendUInt16(&data, numChannels * UInt16(bytesPerSample))
+ appendUInt16(&data, bitsPerSample)
+
+ data.append(contentsOf: "data".utf8)
+ appendUInt32(&data, UInt32(dataSize))
+
+ for i in 0.. MLXArray {
+ // Precompute Wx·x + b over all time steps.
+ let xT = matmul(x, Wx.T) + bias // [B, T, 4*H]
+ let B = x.dim(0)
+ let T = x.dim(1)
+ var h = MLXArray.zeros([B, hiddenSize], dtype: x.dtype)
+ var c = MLXArray.zeros([B, hiddenSize], dtype: x.dtype)
+ var outputs: [MLXArray] = []
+ outputs.reserveCapacity(T)
+
+ let H = hiddenSize
+ for t in 0...{Wx,Wh,bias}` per layer.
+public final class EncodecLSTM: Module {
+ @ModuleInfo public var lstm: [EncodecLSTMCell]
+
+ public init(dimension: Int, numLayers: Int) {
+ self._lstm = ModuleInfo(wrappedValue: (0.. MLXArray {
+ var h = x
+ for cell in lstm { h = cell(h) }
+ return h + x
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/MetalBudget.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/MetalBudget.swift
new file mode 100644
index 0000000..4a104e0
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/MetalBudget.swift
@@ -0,0 +1,58 @@
+import Foundation
+import Cmlx
+import MLX
+
+/// Metal GPU memory budget utilities.
+public enum MetalBudget {
+
+ /// Query real Metal headroom: recommended working set minus active allocations.
+ /// Returns nil if Metal device info is unavailable.
+ public static var availableBytes: Int? {
+ let info = GPU.deviceInfo()
+ let maxWorking = Int(info.maxRecommendedWorkingSetSize)
+ guard maxWorking > 0 else { return nil }
+ let active = Memory.activeMemory
+ let overhead = 256 * 1024 * 1024 // 256 MB safety margin
+ return max(0, maxWorking - active - overhead)
+ }
+
+ /// Total device memory in bytes.
+ public static var totalMemory: Int {
+ GPU.deviceInfo().memorySize
+ }
+
+ /// Maximum recommended working set size in bytes.
+ public static var maxRecommendedWorkingSet: Int {
+ Int(GPU.deviceInfo().maxRecommendedWorkingSetSize)
+ }
+
+ /// Currently active (non-cache) MLX memory in bytes.
+ public static var activeMemory: Int {
+ Memory.activeMemory
+ }
+
+ /// Pin GPU memory to prevent paging under pressure.
+ /// Uses 90% of recommended working set by default.
+ /// Only effective on macOS 15+ / iOS 18+.
+ @discardableResult
+ public static func pinMemory(fraction: Double = 0.9) -> Int {
+ let limit = Int(Double(maxRecommendedWorkingSet) * fraction)
+ var previous: size_t = 0
+ mlx_set_wired_limit(&previous, size_t(limit))
+ return Int(previous)
+ }
+
+ /// Unpin GPU memory (set wired limit to 0).
+ @discardableResult
+ public static func unpinMemory() -> Int {
+ var previous: size_t = 0
+ mlx_set_wired_limit(&previous, 0)
+ return Int(previous)
+ }
+
+ /// Check if a model of the given size (bytes) can fit in available GPU memory.
+ public static func canFit(modelBytes: Int) -> Bool {
+ guard let available = availableBytes else { return true }
+ return modelBytes <= available
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/ModuleMemory.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/ModuleMemory.swift
new file mode 100644
index 0000000..92d999c
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/ModuleMemory.swift
@@ -0,0 +1,26 @@
+import MLX
+import MLXNN
+
+extension Module {
+ /// Estimated memory footprint of all parameters in bytes.
+ public func parameterMemoryBytes() -> Int {
+ var total = 0
+ for array in allParameters() {
+ total += array.nbytes
+ }
+ return total
+ }
+
+ /// Replace all parameters with empty arrays to free GPU memory.
+ /// After calling this, the module is unusable for inference.
+ public func clearParameters() {
+ apply(filter: Self.filterAll) { _ in MLXArray() }
+ Memory.clearCache()
+ }
+
+ /// Collect all parameter arrays (flattened).
+ private func allParameters() -> [MLXArray] {
+ filterMap(filter: Self.filterAll, map: Self.mapParameters())
+ .flattenedValues()
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/PreQuantizedEmbedding.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/PreQuantizedEmbedding.swift
new file mode 100644
index 0000000..0c465e6
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/PreQuantizedEmbedding.swift
@@ -0,0 +1,50 @@
+import Foundation
+import MLX
+import MLXNN
+
+/// Pre-quantized embedding that can be loaded directly from safetensors
+public class PreQuantizedEmbedding: Module {
+ public let groupSize: Int
+ public let bits: Int
+ public let embeddingCount: Int
+ public let dimensions: Int
+
+ @ParameterInfo public var weight: MLXArray
+ @ParameterInfo public var scales: MLXArray
+ @ParameterInfo public var biases: MLXArray
+
+ public init(embeddingCount: Int, dimensions: Int, groupSize: Int = 64, bits: Int = 4) {
+ self.embeddingCount = embeddingCount
+ self.dimensions = dimensions
+ self.groupSize = groupSize
+ self.bits = bits
+
+ // Packed dimensions: 8 values per uint32 for 4-bit
+ let packedDim = dimensions / (32 / bits)
+ let numGroups = dimensions / groupSize
+
+ // Initialize with zeros - will be loaded from weights
+ self._weight.wrappedValue = MLXArray.zeros([embeddingCount, packedDim], dtype: .uint32)
+ self._scales.wrappedValue = MLXArray.zeros([embeddingCount, numGroups], dtype: .bfloat16)
+ self._biases.wrappedValue = MLXArray.zeros([embeddingCount, numGroups], dtype: .bfloat16)
+
+ super.init()
+ self.freeze()
+ }
+
+ public func callAsFunction(_ x: MLXArray) -> MLXArray {
+ let s = x.shape
+ let x = x.flattened()
+ let out = dequantized(
+ weight[x], scales: scales[x], biases: biases[x],
+ groupSize: groupSize, bits: bits)
+ return out.reshaped(s + [-1])
+ }
+
+ /// For use as LM head (matmul with transposed weight)
+ public func asLinear(_ x: MLXArray) -> MLXArray {
+ quantizedMM(
+ x, weight, scales: scales, biases: biases, transpose: true,
+ groupSize: groupSize, bits: bits)
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/QuantizedMLP.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/QuantizedMLP.swift
new file mode 100644
index 0000000..59df7a6
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/QuantizedMLP.swift
@@ -0,0 +1,56 @@
+import Foundation
+import MLX
+import MLXNN
+
+/// SwiGLU MLP — shared by ASR text decoder, TTS Talker, and Code Predictor.
+/// Linear layers are quantized when `bits > 0`, plain Linear otherwise (bf16/fp32 path).
+public class QuantizedMLP: Module {
+ @ModuleInfo public var gateProj: Linear
+ @ModuleInfo public var upProj: Linear
+ @ModuleInfo public var downProj: Linear
+
+ public init(hiddenSize: Int, intermediateSize: Int, groupSize: Int = 64, bits: Int = 4) {
+ self._gateProj.wrappedValue = makeMaybeQuantizedLinear(
+ hiddenSize, intermediateSize, bias: false,
+ groupSize: groupSize, bits: bits)
+ self._upProj.wrappedValue = makeMaybeQuantizedLinear(
+ hiddenSize, intermediateSize, bias: false,
+ groupSize: groupSize, bits: bits)
+ self._downProj.wrappedValue = makeMaybeQuantizedLinear(
+ intermediateSize, hiddenSize, bias: false,
+ groupSize: groupSize, bits: bits)
+
+ super.init()
+ }
+
+ public func callAsFunction(_ x: MLXArray) -> MLXArray {
+ // SwiGLU: down(silu(gate(x)) * up(x))
+ let gate = silu(gateProj(x))
+ let up = upProj(x)
+ return downProj(gate * up)
+ }
+}
+
+/// SwiGLU MLP with `Linear` projections — used by modules that need to
+/// support both bf16/fp16 and quantized bundles. The runtime swaps
+/// Linear → QuantizedLinear in place via `quantize(model:filter:)` when
+/// the loaded weights carry `.scales` for these paths.
+public class MLP: Module {
+ @ModuleInfo public var gateProj: Linear
+ @ModuleInfo public var upProj: Linear
+ @ModuleInfo public var downProj: Linear
+
+ public init(hiddenSize: Int, intermediateSize: Int) {
+ self._gateProj.wrappedValue = Linear(hiddenSize, intermediateSize, bias: false)
+ self._upProj.wrappedValue = Linear(hiddenSize, intermediateSize, bias: false)
+ self._downProj.wrappedValue = Linear(intermediateSize, hiddenSize, bias: false)
+
+ super.init()
+ }
+
+ public func callAsFunction(_ x: MLXArray) -> MLXArray {
+ let gate = silu(gateProj(x))
+ let up = upProj(x)
+ return downProj(gate * up)
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/SDPA.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/SDPA.swift
new file mode 100644
index 0000000..ec8e645
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/SDPA.swift
@@ -0,0 +1,102 @@
+import Foundation
+import MLX
+import MLXFast
+
+/// Multi-head scaled dot-product attention helper used by every attention
+/// module in the project. Takes already-projected per-token Q/K/V tensors of
+/// shape `[B, T, numHeads * headDim]`, reshapes to `[B, H, T, headDim]`,
+/// runs the optimised `MLXFast.scaledDotProductAttention` Metal kernel, then
+/// merges the heads back to `[B, T, numHeads * headDim]` ready for the
+/// output projection.
+///
+/// Each module still owns its own input projections (which vary in width,
+/// quantisation, bias, etc.) and its own output projection — this helper
+/// only collapses the boilerplate around the SDPA call itself.
+public enum SDPA {
+
+ /// Standard attention with optional bool/float mask.
+ public static func multiHead(
+ q: MLXArray, k: MLXArray, v: MLXArray,
+ numHeads: Int, headDim: Int, scale: Float,
+ mask: MLXArray? = nil
+ ) -> MLXArray {
+ let qLen = q.dim(1)
+ let kLen = k.dim(1)
+
+ // Q: [B, T_q, H*D] → [B, H, T_q, D]. Use -1 for the batch dim so the
+ // helper composes with compile(shapeless:) graphs that vary batch.
+ let qHeads = q.reshaped(-1, qLen, numHeads, headDim).transposed(0, 2, 1, 3)
+ let kHeads = k.reshaped(-1, kLen, numHeads, headDim).transposed(0, 2, 1, 3)
+ let vHeads = v.reshaped(-1, kLen, numHeads, headDim).transposed(0, 2, 1, 3)
+
+ let attn = MLXFast.scaledDotProductAttention(
+ queries: qHeads, keys: kHeads, values: vHeads,
+ scale: scale, mask: mask)
+
+ return attn.transposed(0, 2, 1, 3).reshaped(-1, qLen, numHeads * headDim)
+ }
+
+ /// GQA / MQA variant: query and key/value heads can have different
+ /// counts. The kv tensors are repeated to match the query head count
+ /// inside `MLXFast.scaledDotProductAttention`.
+ public static func multiHead(
+ q: MLXArray, k: MLXArray, v: MLXArray,
+ numQueryHeads: Int, numKVHeads: Int, headDim: Int, scale: Float,
+ mask: MLXArray? = nil
+ ) -> MLXArray {
+ let qLen = q.dim(1)
+ let kLen = k.dim(1)
+
+ let qHeads = q.reshaped(-1, qLen, numQueryHeads, headDim).transposed(0, 2, 1, 3)
+ let kHeads = k.reshaped(-1, kLen, numKVHeads, headDim).transposed(0, 2, 1, 3)
+ let vHeads = v.reshaped(-1, kLen, numKVHeads, headDim).transposed(0, 2, 1, 3)
+
+ let attn = MLXFast.scaledDotProductAttention(
+ queries: qHeads, keys: kHeads, values: vHeads,
+ scale: scale, mask: mask)
+
+ return attn.transposed(0, 2, 1, 3).reshaped(-1, qLen, numQueryHeads * headDim)
+ }
+
+ /// Run SDPA on tensors that are already shaped `[B, H, T, D]` (e.g. after
+ /// RoPE / KV-cache concatenation in LLM-style attention) and merge the
+ /// heads back to `[B, T, H * D]`. Saves the boilerplate transpose+reshape
+ /// after the SDPA call without dictating where the projections live.
+ public static func attendAndMerge(
+ qHeads: MLXArray, kHeads: MLXArray, vHeads: MLXArray,
+ scale: Float,
+ mask: MLXArray? = nil
+ ) -> MLXArray {
+ let attn = MLXFast.scaledDotProductAttention(
+ queries: qHeads, keys: kHeads, values: vHeads,
+ scale: scale, mask: mask)
+ return mergeHeads(attn)
+ }
+
+ /// `ScaledDotProductAttentionMaskMode` overload — used by modules that
+ /// pass causal / additive masks via the newer mlx-swift API.
+ public static func attendAndMerge(
+ qHeads: MLXArray, kHeads: MLXArray, vHeads: MLXArray,
+ scale: Float,
+ mask: MLXFast.ScaledDotProductAttentionMaskMode
+ ) -> MLXArray {
+ let attn = MLXFast.scaledDotProductAttention(
+ queries: qHeads, keys: kHeads, values: vHeads,
+ scale: scale, mask: mask)
+ return mergeHeads(attn)
+ }
+
+ /// Merge heads back: `[B, H, T, D] → [B, T, H * D]`.
+ ///
+ /// Uses `-1` for the batch dimension so the result composes with
+ /// `MLX.compile(shapeless: true)` graphs that vary the batch at runtime
+ /// (e.g. Qwen3-TTS Talker autoregressive decode with different batch
+ /// sizes per call).
+ @inline(__always)
+ public static func mergeHeads(_ attn: MLXArray) -> MLXArray {
+ let H = attn.dim(1)
+ let T = attn.dim(2)
+ let D = attn.dim(3)
+ return attn.transposed(0, 2, 1, 3).reshaped(-1, T, H * D)
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/WeightLoading.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/WeightLoading.swift
new file mode 100644
index 0000000..ae314aa
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/WeightLoading.swift
@@ -0,0 +1,292 @@
+import Foundation
+import MLX
+import MLXNN
+
+/// Build a Linear layer that is quantized when bits > 0, plain when bits == 0.
+/// QuantizedLinear inherits from Linear, so the return type is always Linear and
+/// caller code can store it in a single `@ModuleInfo var x: Linear` field.
+public func makeMaybeQuantizedLinear(
+ _ inputDimensions: Int,
+ _ outputDimensions: Int,
+ bias: Bool,
+ groupSize: Int,
+ bits: Int
+) -> Linear {
+ if bits > 0 {
+ return QuantizedLinear(inputDimensions, outputDimensions, bias: bias,
+ groupSize: groupSize, bits: bits)
+ } else {
+ return Linear(inputDimensions, outputDimensions, bias: bias)
+ }
+}
+
+/// Generic weight loading utilities shared between ASR and TTS
+public enum CommonWeightLoader {
+
+ /// Load weights from safetensors file
+ public static func loadSafetensors(url: URL) throws -> [String: MLXArray] {
+ try MLX.loadArrays(url: url)
+ }
+
+ /// Load all safetensors from a directory, optionally filtering by prefix
+ public static func loadAllSafetensors(
+ from directory: URL,
+ prefix: String? = nil,
+ stripPrefix: Bool = true
+ ) throws -> [String: MLXArray] {
+ let fileManager = FileManager.default
+ let contents = try fileManager.contentsOfDirectory(at: directory, includingPropertiesForKeys: nil)
+ let safetensorFiles = contents.filter { $0.pathExtension == "safetensors" }
+
+ guard !safetensorFiles.isEmpty else {
+ throw WeightLoadingError.noWeightsFound(directory)
+ }
+
+ var allWeights: [String: MLXArray] = [:]
+ for file in safetensorFiles {
+ let weights = try loadSafetensors(url: file)
+ allWeights.merge(weights) { _, new in new }
+ }
+
+ // Filter and strip prefix if specified
+ guard let prefix = prefix else { return allWeights }
+
+ var filtered: [String: MLXArray] = [:]
+ for (key, value) in allWeights {
+ if key.hasPrefix(prefix) {
+ let strippedKey = stripPrefix ? String(key.dropFirst(prefix.count)) : key
+ filtered[strippedKey] = value
+ }
+ }
+ return filtered
+ }
+
+ // MARK: - Quantized Weight Application Helpers
+
+ public static func applyQuantizedEmbeddingWeights(
+ to embedding: PreQuantizedEmbedding,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ var params: [String: NestedItem] = [:]
+
+ if let weight = weights["\(prefix).weight"] {
+ params["weight"] = .value(weight)
+ }
+ if let scales = weights["\(prefix).scales"] {
+ params["scales"] = .value(scales)
+ }
+ if let biases = weights["\(prefix).biases"] {
+ params["biases"] = .value(biases)
+ }
+
+ if !params.isEmpty {
+ embedding.update(parameters: ModuleParameters(values: params))
+ }
+ }
+
+ /// Apply weights to a Linear (or QuantizedLinear, since the latter inherits from
+ /// the former). When the layer is a QuantizedLinear, `.scales`/`.biases` are wired
+ /// in addition to `.weight`. For plain Linear those keys are absent in the
+ /// safetensors (bf16/fp32 model) and only `.weight` (+ optional `.bias`) apply.
+ public static func applyQuantizedLinearWeights(
+ to linear: Linear,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ var params: [String: NestedItem] = [:]
+
+ if let weight = weights["\(prefix).weight"] {
+ params["weight"] = .value(weight)
+ }
+ if linear is QuantizedLinear {
+ if let scales = weights["\(prefix).scales"] {
+ params["scales"] = .value(scales)
+ }
+ if let biases = weights["\(prefix).biases"] {
+ params["biases"] = .value(biases)
+ }
+ }
+ // Regular linear bias (separate from quantization biases)
+ if let bias = weights["\(prefix).bias"] {
+ params["bias"] = .value(bias)
+ }
+
+ if !params.isEmpty {
+ linear.update(parameters: ModuleParameters(values: params))
+ }
+ }
+
+ public static func applyRMSNormWeights(
+ to norm: RMSNorm,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ var params: [String: NestedItem] = [:]
+
+ if let weight = weights["\(prefix).weight"] {
+ params["weight"] = .value(weight)
+ }
+
+ if !params.isEmpty {
+ norm.update(parameters: ModuleParameters(values: params))
+ }
+ }
+
+ public static func applyLinearWeights(
+ to linear: Linear,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ var params: [String: NestedItem] = [:]
+
+ if let weight = weights["\(prefix).weight"] {
+ params["weight"] = .value(weight)
+ }
+ if let bias = weights["\(prefix).bias"] {
+ params["bias"] = .value(bias)
+ }
+
+ if !params.isEmpty {
+ linear.update(parameters: ModuleParameters(values: params))
+ }
+ }
+
+ public static func applyLayerNormWeights(
+ to layerNorm: LayerNorm,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ var params: [String: NestedItem] = [:]
+
+ if let weight = weights["\(prefix).weight"] {
+ params["weight"] = .value(weight)
+ }
+ if let bias = weights["\(prefix).bias"] {
+ params["bias"] = .value(bias)
+ }
+
+ if !params.isEmpty {
+ layerNorm.update(parameters: ModuleParameters(values: params))
+ }
+ }
+
+ public static func applyEmbeddingWeights(
+ to embedding: Embedding,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ var params: [String: NestedItem] = [:]
+
+ if let weight = weights["\(prefix).weight"] {
+ params["weight"] = .value(weight)
+ }
+
+ if !params.isEmpty {
+ embedding.update(parameters: ModuleParameters(values: params))
+ }
+ }
+
+ public static func applyConv1dWeights(
+ to conv: Conv1d,
+ prefix: String,
+ from weights: [String: MLXArray],
+ transpose: Bool = false
+ ) {
+ var params: [String: NestedItem] = [:]
+
+ if let weight = weights["\(prefix).weight"] {
+ // PyTorch Conv1d: [out, in, kernel] -> MLX Conv1d: [out, kernel, in]
+ let w = transpose ? weight.transposed(0, 2, 1) : weight
+ params["weight"] = .value(w)
+ }
+ if let bias = weights["\(prefix).bias"] {
+ params["bias"] = .value(bias)
+ }
+
+ if !params.isEmpty {
+ conv.update(parameters: ModuleParameters(values: params))
+ }
+ }
+
+ public static func applyConvTransposed1dWeights(
+ to conv: ConvTransposed1d,
+ prefix: String,
+ from weights: [String: MLXArray],
+ transpose: Bool = false
+ ) {
+ var params: [String: NestedItem] = [:]
+
+ if let weight = weights["\(prefix).weight"] {
+ // PyTorch ConvTranspose1d: [in, out, kernel] -> MLX ConvTransposed1d: [out, kernel, in]
+ let w = transpose ? weight.transposed(1, 2, 0) : weight
+ params["weight"] = .value(w)
+ }
+ if let bias = weights["\(prefix).bias"] {
+ params["bias"] = .value(bias)
+ }
+
+ if !params.isEmpty {
+ conv.update(parameters: ModuleParameters(values: params))
+ }
+ }
+
+ /// Apply QuantizedMLP weights (SwiGLU)
+ public static func applyQuantizedMLPWeights(
+ to mlp: QuantizedMLP,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ applyQuantizedLinearWeights(to: mlp.gateProj, prefix: "\(prefix).gate_proj", from: weights)
+ applyQuantizedLinearWeights(to: mlp.upProj, prefix: "\(prefix).up_proj", from: weights)
+ applyQuantizedLinearWeights(to: mlp.downProj, prefix: "\(prefix).down_proj", from: weights)
+ }
+
+ /// Apply MLP weights (SwiGLU) — dispatches to quantized or plain
+ /// projections per-leaf based on `.scales` presence. Use when the
+ /// surrounding module was declared with `Linear` and may have been
+ /// swapped to `QuantizedLinear` by `quantize(model:filter:)`.
+ public static func applyMLPWeights(
+ to mlp: MLP,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ applyMaybeQuantizedLinearWeights(to: mlp.gateProj, prefix: "\(prefix).gate_proj", from: weights)
+ applyMaybeQuantizedLinearWeights(to: mlp.upProj, prefix: "\(prefix).up_proj", from: weights)
+ applyMaybeQuantizedLinearWeights(to: mlp.downProj, prefix: "\(prefix).down_proj", from: weights)
+ }
+
+ /// Apply weights to a `Linear` that may have been swapped to
+ /// `QuantizedLinear`. Picks the right keys (`weight`, optional `bias`
+ /// for plain Linear; plus `scales`, `biases` when quantized) based on
+ /// what is present in the safetensors.
+ public static func applyMaybeQuantizedLinearWeights(
+ to linear: Linear,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ if weights["\(prefix).scales"] != nil, let q = linear as? QuantizedLinear {
+ applyQuantizedLinearWeights(to: q, prefix: prefix, from: weights)
+ } else {
+ applyLinearWeights(to: linear, prefix: prefix, from: weights)
+ }
+ }
+}
+
+/// Weight loading errors
+public enum WeightLoadingError: Error, LocalizedError {
+ case noWeightsFound(URL)
+ case incompatibleWeights(String)
+ case missingRequiredWeight(String)
+
+ public var errorDescription: String? {
+ switch self {
+ case .noWeightsFound(let url):
+ return "No safetensors files found in: \(url.path)"
+ case .incompatibleWeights(let reason):
+ return "Incompatible weights: \(reason)"
+ case .missingRequiredWeight(let key):
+ return "Missing required weight: \(key)"
+ }
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/AudioEncoder.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/AudioEncoder.swift
new file mode 100644
index 0000000..b26fecb
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/AudioEncoder.swift
@@ -0,0 +1,512 @@
+import Foundation
+import MLX
+import MLXNN
+import MLXFast
+import MLXCommon
+import AudioCommon
+
+/// Audio encoder configuration matching Qwen3-ASR HuggingFace model
+public struct Qwen3AudioEncoderConfig: Sendable {
+ public let dModel: Int // 896
+ public let encoderAttentionHeads: Int // 14
+ public let encoderFFNDim: Int // 3584
+ public let encoderLayers: Int // 18
+ public let numMelBins: Int // 128
+ public let maxSourcePositions: Int // 1500
+ public let outputDim: Int // 1024
+ public let downsampleHiddenSize: Int // 480
+ public let convChunksize: Int // 500
+ public let nWindow: Int // 50 (chunk size = n_window * 2 = 100)
+ public let nWindowInfer: Int // 800
+ public let dropout: Float // 0.0
+ public let attentionDropout: Float // 0.0
+ public let activationDropout: Float // 0.0
+ public let layerNormEps: Float // 1e-5
+ public let convOutInputDim: Int // 7680 (480 channels * 16 spatial positions)
+
+ /// Config for Qwen3-ASR-0.6B (default)
+ public static let `default` = Qwen3AudioEncoderConfig(
+ dModel: 896,
+ encoderAttentionHeads: 14,
+ encoderFFNDim: 3584,
+ encoderLayers: 18,
+ numMelBins: 128,
+ maxSourcePositions: 1500,
+ outputDim: 1024,
+ downsampleHiddenSize: 480,
+ convChunksize: 500,
+ nWindow: 50,
+ nWindowInfer: 800,
+ dropout: 0.0,
+ attentionDropout: 0.0,
+ activationDropout: 0.0,
+ layerNormEps: 1e-5,
+ convOutInputDim: 7680
+ )
+
+ /// Alias for 0.6B config
+ public static let small = `default`
+
+ /// Config for Qwen3-ASR-1.7B
+ public static let large = Qwen3AudioEncoderConfig(
+ dModel: 1024,
+ encoderAttentionHeads: 16,
+ encoderFFNDim: 4096,
+ encoderLayers: 24,
+ numMelBins: 128,
+ maxSourcePositions: 1500,
+ outputDim: 2048,
+ downsampleHiddenSize: 480,
+ convChunksize: 500,
+ nWindow: 50,
+ nWindowInfer: 800,
+ dropout: 0.0,
+ attentionDropout: 0.0,
+ activationDropout: 0.0,
+ layerNormEps: 1e-5,
+ convOutInputDim: 7680
+ )
+
+ /// Config for Qwen3-ForcedAligner-0.6B (large encoder projecting to 1024-dim text decoder)
+ public static let forcedAligner = Qwen3AudioEncoderConfig(
+ dModel: 1024,
+ encoderAttentionHeads: 16,
+ encoderFFNDim: 4096,
+ encoderLayers: 24,
+ numMelBins: 128,
+ maxSourcePositions: 1500,
+ outputDim: 1024,
+ downsampleHiddenSize: 480,
+ convChunksize: 500,
+ nWindow: 50,
+ nWindowInfer: 800,
+ dropout: 0.0,
+ attentionDropout: 0.0,
+ activationDropout: 0.0,
+ layerNormEps: 1e-5,
+ convOutInputDim: 7680
+ )
+}
+
+/// Multi-head self-attention for audio encoder layers
+/// Weight names: self_attn.q_proj, k_proj, v_proj, out_proj
+public class AudioSelfAttention: Module {
+ let numHeads: Int
+ let headDim: Int
+ let scale: Float
+
+ @ModuleInfo(key: "q_proj") public var qProj: Linear
+ @ModuleInfo(key: "k_proj") public var kProj: Linear
+ @ModuleInfo(key: "v_proj") public var vProj: Linear
+ @ModuleInfo(key: "out_proj") public var outProj: Linear
+
+ public init(hiddenSize: Int, numHeads: Int) {
+ self.numHeads = numHeads
+ self.headDim = hiddenSize / numHeads
+ self.scale = 1.0 / sqrt(Float(headDim))
+
+ self._qProj.wrappedValue = Linear(hiddenSize, hiddenSize, bias: true)
+ self._kProj.wrappedValue = Linear(hiddenSize, hiddenSize, bias: true)
+ self._vProj.wrappedValue = Linear(hiddenSize, hiddenSize, bias: true)
+ self._outProj.wrappedValue = Linear(hiddenSize, hiddenSize, bias: true)
+
+ super.init()
+ }
+
+ public func callAsFunction(_ x: MLXArray, attentionMask: MLXArray? = nil) -> MLXArray {
+ let q = qProj(x)
+ let k = kProj(x)
+ let v = vProj(x)
+ let out = SDPA.multiHead(
+ q: q, k: k, v: v,
+ numHeads: numHeads, headDim: headDim, scale: scale,
+ mask: attentionMask)
+ return outProj(out)
+ }
+}
+
+/// Audio encoder transformer layer
+/// Weight names: self_attn, self_attn_layer_norm, fc1, fc2, final_layer_norm
+public class AudioEncoderLayer: Module {
+ @ModuleInfo(key: "self_attn") public var selfAttn: AudioSelfAttention
+ @ModuleInfo(key: "self_attn_layer_norm") public var selfAttnLayerNorm: LayerNorm
+ @ModuleInfo public var fc1: Linear
+ @ModuleInfo public var fc2: Linear
+ @ModuleInfo(key: "final_layer_norm") public var finalLayerNorm: LayerNorm
+
+ public init(hiddenSize: Int, numHeads: Int, ffnDim: Int, layerNormEps: Float) {
+ self._selfAttn.wrappedValue = AudioSelfAttention(hiddenSize: hiddenSize, numHeads: numHeads)
+ self._selfAttnLayerNorm.wrappedValue = LayerNorm(dimensions: hiddenSize, eps: layerNormEps)
+ self._fc1.wrappedValue = Linear(hiddenSize, ffnDim, bias: true)
+ self._fc2.wrappedValue = Linear(ffnDim, hiddenSize, bias: true)
+ self._finalLayerNorm.wrappedValue = LayerNorm(dimensions: hiddenSize, eps: layerNormEps)
+
+ super.init()
+ }
+
+ public func callAsFunction(_ x: MLXArray, attentionMask: MLXArray? = nil) -> MLXArray {
+ // Self attention with residual
+ var residual = x
+ var hidden = selfAttnLayerNorm(x)
+ hidden = selfAttn(hidden, attentionMask: attentionMask)
+ hidden = residual + hidden
+
+ // FFN with residual
+ residual = hidden
+ hidden = finalLayerNorm(hidden)
+ hidden = fc1(hidden)
+ hidden = gelu(hidden)
+ hidden = fc2(hidden)
+ hidden = residual + hidden
+
+ return hidden
+ }
+}
+
+/// Create sinusoidal position embeddings matching Python mlx-audio implementation
+/// - Parameters:
+/// - seqLen: Sequence length
+/// - dModel: Model dimension (channels)
+/// - Returns: Position embeddings [1, seqLen, dModel]
+private func createSinusoidalPositionEmbeddings(seqLen: Int, dModel: Int) -> MLXArray {
+ let halfDim = dModel / 2
+ let maxTimescale: Float = 10000.0
+
+ // Python formula: log_timescale_increment = log(max_timescale) / (channels // 2 - 1)
+ // inv_timescales = exp(-log_timescale_increment * arange(channels // 2))
+ let logTimescaleIncrement = log(maxTimescale) / Float(halfDim - 1)
+ let invTimescales = exp(
+ MLXArray(0.. [seqLen, halfDim]
+ let scaledTime = positions.expandedDimensions(axis: 1) * invTimescales.expandedDimensions(axis: 0)
+
+ // Sin and cos embeddings
+ let sinEmbed = sin(scaledTime) // [seqLen, halfDim]
+ let cosEmbed = cos(scaledTime) // [seqLen, halfDim]
+
+ // Concatenate [sin, cos] along axis 1 (NOT interleave!)
+ // Python: concatenate([sin(scaled_time), cos(scaled_time)], axis=1)
+ let posEmbed = concatenated([sinEmbed, cosEmbed], axis: 1) // [seqLen, dModel]
+
+ // Add batch dimension
+ return posEmbed.expandedDimensions(axis: 0) // [1, seqLen, dModel]
+}
+
+/// Full Qwen3-ASR Audio Encoder (audio_tower)
+/// Matches HuggingFace weight structure exactly
+public class Qwen3AudioEncoder: Module {
+ public let config: Qwen3AudioEncoderConfig
+
+ // Cache for sinusoidal position embeddings keyed by sequence length
+ private var cachedPosEmbeddings: [Int: MLXArray] = [:]
+
+ // Conv frontend - using 2D convolutions
+ // Input: [batch, 1, mel=128, time] (single channel mel spectrogram image)
+ // Weight format in safetensors: [out, in, kH, kW] -> transpose to MLX [out, kH, kW, in]
+ @ModuleInfo public var conv2d1: Conv2d // 1 -> 480 channels, 3x3 kernel
+ @ModuleInfo public var conv2d2: Conv2d // 480 -> 480 channels, 3x3, stride 2
+ @ModuleInfo public var conv2d3: Conv2d // 480 -> 480 channels, 3x3, stride 2
+
+ // Output projection: flattened conv features -> d_model
+ // Weight: [896, 7680] -> Linear(7680, 896)
+ @ModuleInfo(key: "conv_out") public var convOut: Linear
+
+ // Post layer norm
+ @ModuleInfo(key: "ln_post") public var lnPost: LayerNorm
+
+ // Projector to text model dimension
+ // proj1: [896, 896] -> Linear(896, 896)
+ // proj2: [1024, 896] -> Linear(896, 1024)
+ @ModuleInfo public var proj1: Linear
+ @ModuleInfo public var proj2: Linear
+
+ // Transformer layers
+ @ModuleInfo public var layers: [AudioEncoderLayer]
+
+ public init(config: Qwen3AudioEncoderConfig = .default) {
+ self.config = config
+
+ // Conv2D layers for mel spectrogram processing
+ // All three convs have stride 2 for 8x downsampling
+ // Input: [batch, mel_bins=128, time, 1] in NHWC
+ self._conv2d1.wrappedValue = Conv2d(
+ inputChannels: 1,
+ outputChannels: config.downsampleHiddenSize, // 480
+ kernelSize: IntOrPair(3),
+ stride: IntOrPair(2), // Changed from 1 to 2
+ padding: IntOrPair(1)
+ )
+ self._conv2d2.wrappedValue = Conv2d(
+ inputChannels: config.downsampleHiddenSize,
+ outputChannels: config.downsampleHiddenSize,
+ kernelSize: IntOrPair(3),
+ stride: IntOrPair(2),
+ padding: IntOrPair(1)
+ )
+ self._conv2d3.wrappedValue = Conv2d(
+ inputChannels: config.downsampleHiddenSize,
+ outputChannels: config.downsampleHiddenSize,
+ kernelSize: IntOrPair(3),
+ stride: IntOrPair(2),
+ padding: IntOrPair(1)
+ )
+
+ // Output conv projection: flattened features (7680) -> d_model (896)
+ self._convOut.wrappedValue = Linear(config.convOutInputDim, config.dModel, bias: false)
+
+ // Post layer norm
+ self._lnPost.wrappedValue = LayerNorm(dimensions: config.dModel, eps: config.layerNormEps)
+
+ // Projector to text model dimension
+ // proj1: 896 -> 896
+ // proj2: 896 -> 1024
+ self._proj1.wrappedValue = Linear(config.dModel, config.dModel, bias: true)
+ self._proj2.wrappedValue = Linear(config.dModel, config.outputDim, bias: true)
+
+ // Transformer layers
+ self._layers.wrappedValue = (0.. Int {
+ let chunkSize = config.nWindow * 2 // 100
+ let remainder = inputLength % chunkSize
+
+ // Process remainder through conv downsampling formula
+ var featLen = (remainder - 1) / 2 + 1 // First stride-2
+ featLen = (featLen - 1) / 2 + 1 // Second stride-2
+ featLen = (featLen - 1) / 2 + 1 // Third stride-2
+
+ // Full chunks each produce 13 tokens
+ let fullChunkTokens = (inputLength / chunkSize) * 13
+
+ // Handle edge case when remainder is 0
+ let remainderTokens = remainder > 0 ? max(featLen, 1) : 0
+
+ return fullChunkTokens + remainderTokens
+ }
+
+ /// Process a single chunk through conv layers
+ /// Input: [batch, mel=128, time, 1] in NHWC format
+ /// Output: [batch, time_tokens, features=7680]
+ private func processConvChunk(_ chunk: MLXArray) -> MLXArray {
+ var x = chunk
+
+ // Apply conv layers: 128 -> 64 -> 32 -> 16 in mel dimension
+ // Time dimension also downsampled 8x
+ x = conv2d1(x)
+ x = gelu(x)
+ x = conv2d2(x)
+ x = gelu(x)
+ x = conv2d3(x)
+ x = gelu(x)
+
+ // Shape after conv: [batch, mel/8=16, time/8, 480]
+ let batch = x.dim(0)
+ let height = x.dim(1) // 16 (mel after 3x stride-2)
+ let width = x.dim(2) // time/8
+ let channels = x.dim(3) // 480
+
+ // Flatten mel*channels: [batch, 16, time_tokens, 480] -> [batch, time_tokens, 16*480]
+ // Transpose to [batch, time, mel, channels] then flatten last two
+ x = x.transposed(0, 2, 1, 3) // [batch, time, mel, channels]
+ x = x.reshaped(batch, width, height * channels) // [batch, time_tokens, 7680]
+
+ return x
+ }
+
+ /// Create block attention mask for preventing cross-chunk attention
+ /// Each block in cu_seqlens can only attend to itself
+ /// Uses MLXArray broadcast comparison instead of scalar O(n^2) loop
+ private func createBlockAttentionMask(seqLen: Int, cuSeqlens: [Int]) -> MLXArray {
+ // Assign a block ID to each position
+ var blockIds = [Int32](repeating: 0, count: seqLen)
+ for i in 0..<(cuSeqlens.count - 1) {
+ let start = cuSeqlens[i]
+ let end = cuSeqlens[i + 1]
+ for pos in start.. [1, 1, seqLen, seqLen]
+ return mask.expandedDimensions(axes: [0, 1])
+ }
+
+ /// Process mel spectrogram with time chunking (matching Python mlx-audio exactly)
+ /// Input: [batch, mel_bins, time]
+ /// Output: [time', output_dim] (no batch dim, matching Python)
+ public func callAsFunction(_ melFeatures: MLXArray) -> MLXArray {
+ let timeFrames = melFeatures.dim(2)
+ let chunkSize = config.nWindow * 2 // 100
+
+ // Calculate number of chunks
+ let numChunks = (timeFrames + chunkSize - 1) / chunkSize // ceil division
+
+ // Compute chunk lengths (all full except possibly last)
+ var chunkLengths = [Int]()
+ for i in 0.. [numChunks, time, channels*freq]
+ x = x.transposed(0, 2, 3, 1) // [numChunks, time, channels, freq]
+ x = x.reshaped(numChunksBatch, timeAfterConv, channels * freq) // [numChunks, time, 7680]
+
+ // Project through conv_out (7680 -> 896)
+ x = convOut(x)
+
+ // Add sinusoidal position embeddings - same for each chunk!
+ // Cache to avoid recomputing for the same sequence length
+ let posEmbed: MLXArray
+ if let cached = cachedPosEmbeddings[timeAfterConv] {
+ posEmbed = cached
+ } else {
+ let computed = createSinusoidalPositionEmbeddings(seqLen: timeAfterConv, dModel: config.dModel)
+ cachedPosEmbeddings[timeAfterConv] = computed
+ posEmbed = computed
+ }
+ x = x + posEmbed // Broadcasting: [numChunks, time, 896] + [1, time, 896]
+
+ // Calculate valid lengths after CNN for each chunk
+ var featureLensAfterCnn = [Int]()
+ for clen in chunkLengths {
+ // Formula from Python: (((clen-1)//2 + 1 - 1)//2 + 1 - 1)//2 + 1
+ var featLen = (clen - 1) / 2 + 1
+ featLen = (featLen - 1) / 2 + 1
+ featLen = (featLen - 1) / 2 + 1
+ featureLensAfterCnn.append(featLen)
+ }
+
+ // Extract valid portions and concatenate
+ var hiddenList: [MLXArray] = []
+ for i in 0..Mel
+ let logstepMelToHz: Float = log(6.4) / 27.0 // For Mel->Hz
+
+ func hzToMel(_ hz: Float) -> Float {
+ if hz < minLogHertz {
+ return 3.0 * hz / 200.0 // Linear region
+ } else {
+ return minLogMel + log(hz / minLogHertz) * logstepHzToMel // Log region
+ }
+ }
+
+ func melToHz(_ mel: Float) -> Float {
+ if mel < minLogMel {
+ return 200.0 * mel / 3.0 // Linear region
+ } else {
+ return minLogHertz * exp((mel - minLogMel) * logstepMelToHz) // Exp region
+ }
+ }
+
+ // Use paddedFFT for bin count since we zero-pad to 512 for FFT
+ let nBins = paddedFFT / 2 + 1 // 257 for paddedFFT=512
+
+ // FFT bin frequencies: k * fs / paddedFFT (not nFFT, since we zero-pad)
+ var fftFreqs = [Float](repeating: 0, count: nBins)
+ for i in 0.. MLXArray {
+ let nBins = paddedFFT / 2 + 1 // 257 bins for 512-point FFT
+ let halfPadded = paddedFFT / 2 // 256
+
+ // Pad audio with reflect padding (like Whisper/librosa)
+ let padLength = nFFT / 2
+ var paddedAudio = [Float](repeating: 0, count: padLength + audio.count + padLength)
+
+ // Reflect pad left side
+ for i in 0.. filterbankT [nBins, nMels]
+ var filterbankT = [Float](repeating: 0, count: nBins * nMels)
+ vDSP_mtrans(filterbank, 1, &filterbankT, 1, vDSP_Length(nBins), vDSP_Length(nMels))
+
+ // C[nFrames, nMels] = A[nFrames, nBins] * B[nBins, nMels]
+ vDSP_mmul(magnitude, 1, filterbankT, 1, &melSpec, 1,
+ vDSP_Length(nFrames), vDSP_Length(nMels), vDSP_Length(nBins))
+
+ // --- Vectorized log10, clamp, normalize ---
+ let count = melSpec.count
+ var countN = Int32(count)
+
+ // Clamp minimum to epsilon before log
+ var epsilon: Float = 1e-10
+ vDSP_vclip(melSpec, 1, &epsilon, [Float.greatestFiniteMagnitude], &melSpec, 1, vDSP_Length(count))
+
+ // log10 using vForce
+ vvlog10f(&melSpec, melSpec, &countN)
+
+ // Find max value for dynamic range compression
+ var maxVal: Float = -Float.infinity
+ vDSP_maxv(melSpec, 1, &maxVal, vDSP_Length(count))
+
+ // Clamp minimum to max - 8.0
+ var minClamp = maxVal - 8.0
+ var maxClamp = Float.greatestFiniteMagnitude
+ vDSP_vclip(melSpec, 1, &minClamp, &maxClamp, &melSpec, 1, vDSP_Length(count))
+
+ // Normalize: (x + 4.0) / 4.0 = x * 0.25 + 1.0
+ var scale: Float = 0.25
+ var offset: Float = 1.0
+ vDSP_vsmsa(melSpec, 1, &scale, &offset, &melSpec, 1, vDSP_Length(count))
+
+ // CRITICAL: HuggingFace WhisperFeatureExtractor removes the last frame: log_spec[:, :-1]
+ let trimmedFrames = nFrames - 1
+ let trimmedMelSpec = Array(melSpec.prefix(trimmedFrames * nMels))
+
+ // Qwen3-ASR encoder handles arbitrary-length audio via windowed attention.
+ // The chunkLength=30 from preprocessor_config.json is inherited from the HuggingFace
+ // WhisperFeatureExtractor class but is not enforced by the official Qwen3-ASR pipeline.
+ // Cap at 1200 seconds (120000 frames) to match the official pipeline's upper bound
+ // and prevent OOM on extremely long inputs.
+ let maxFrames = 1200 * sampleRate / hopLength // 120000 frames at 16kHz/160hop
+ let finalFrames: Int
+ let finalMelSpec: [Float]
+ if trimmedFrames > maxFrames {
+ finalFrames = maxFrames
+ finalMelSpec = Array(trimmedMelSpec.prefix(maxFrames * nMels))
+ } else {
+ finalFrames = trimmedFrames
+ finalMelSpec = trimmedMelSpec
+ }
+
+ let array = MLXArray(finalMelSpec, [finalFrames, nMels])
+ return array.transposed(1, 0) // [mel_bins, time_frames]
+ }
+
+ /// Process audio for Qwen3-ASR model
+ /// - Parameter audio: Raw audio samples (any sample rate)
+ /// - Parameter inputSampleRate: Sample rate of input audio
+ /// - Returns: Preprocessed mel features ready for the model
+ public func process(_ audio: [Float], sampleRate inputSampleRate: Int) -> MLXArray {
+ var processedAudio = audio
+
+ // Resample if needed
+ if inputSampleRate != sampleRate {
+ processedAudio = AudioFileLoader.resample(audio, from: inputSampleRate, to: sampleRate)
+ }
+
+ // NOTE: HuggingFace WhisperFeatureExtractor does NOT normalize audio amplitude
+ // The model expects raw audio values (typically in [-1, 1] range from int16 conversion)
+ // Do NOT divide by max absolute value!
+
+ // Extract features
+ return extractFeatures(processedAudio)
+ }
+
+ // MARK: - MLX-free variants (for CoreML-only path)
+
+ /// Extract mel spectrogram features without MLXArray dependency.
+ /// Produces identical output to `extractFeatures(_:)` but returns a plain `MelFeatures`
+ /// struct with layout `[melBins, timeFrames]` (transposed, same as the MLXArray version).
+ ///
+ /// All mel computation uses the same Accelerate/vDSP pipeline; only the final
+ /// transpose is done with a pure-Swift loop instead of `MLXArray.transposed`.
+ public func extractFeaturesRaw(_ audio: [Float]) -> MelFeatures {
+ let nBins = paddedFFT / 2 + 1 // 257 bins for 512-point FFT
+ let halfPadded = paddedFFT / 2 // 256
+
+ // Pad audio with reflect padding (like Whisper/librosa)
+ let padLength = nFFT / 2
+ var paddedAudio = [Float](repeating: 0, count: padLength + audio.count + padLength)
+
+ // Reflect pad left side
+ for i in 0.. maxFrames {
+ finalFrames = maxFrames
+ finalMelSpec = Array(trimmedMelSpec.prefix(maxFrames * nMels))
+ } else {
+ finalFrames = trimmedFrames
+ finalMelSpec = trimmedMelSpec
+ }
+
+ // Transpose [timeFrames, melBins] -> [melBins, timeFrames] in pure Swift
+ var transposed = [Float](repeating: 0, count: finalFrames * nMels)
+ for t in 0.. MelFeatures {
+ var processedAudio = audio
+
+ // Resample if needed
+ if inputSampleRate != sampleRate {
+ processedAudio = AudioFileLoader.resample(audio, from: inputSampleRate, to: sampleRate)
+ }
+
+ // NOTE: HuggingFace WhisperFeatureExtractor does NOT normalize audio amplitude
+ // The model expects raw audio values (typically in [-1, 1] range from int16 conversion)
+ // Do NOT divide by max absolute value!
+
+ return extractFeaturesRaw(processedAudio)
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Configuration.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Configuration.swift
new file mode 100644
index 0000000..933b823
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Configuration.swift
@@ -0,0 +1,158 @@
+import Foundation
+import AudioCommon
+
+/// Configuration for Qwen3-ASR audio encoder
+public struct AudioEncoderConfig: Codable, Sendable {
+ public var inputDim: Int = 128 // Mel filterbank bins
+ public var hiddenDim: Int = 1024 // Transformer hidden dim
+ public var numLayers: Int = 18 // Transformer layers
+ public var numHeads: Int = 16 // Attention heads
+ public var kernelSize: Int = 3 // Conv kernel size
+ public var headDim: Int = 64 // Head dimension (hiddenDim / numHeads)
+ public var ffnDim: Int = 4096 // FFN intermediate dim
+ public var maxSourcePositions: Int = 1500
+ public var layerNormEps: Float = 1e-5
+ public var attentionDropout: Float = 0.0
+ public var dropoutRate: Float = 0.0
+ public var layerdrop: Float = 0.0
+ public var numMelBins: Int = 128
+ public var projectorHiddenAct: String = "silu"
+
+ public init() {}
+
+ /// Config for Qwen3-ASR-0.6B
+ public static var small: AudioEncoderConfig {
+ var config = AudioEncoderConfig()
+ config.hiddenDim = 768
+ config.numLayers = 12
+ config.numHeads = 12
+ config.headDim = 64
+ config.ffnDim = 3072
+ return config
+ }
+
+ /// Config for Qwen3-ASR-1.7B
+ public static var large: AudioEncoderConfig {
+ var config = AudioEncoderConfig()
+ config.hiddenDim = 1024
+ config.numLayers = 24
+ config.numHeads = 16
+ config.headDim = 64
+ config.ffnDim = 4096
+ return config
+ }
+}
+
+/// Configuration for Qwen3 text decoder
+public struct TextDecoderConfig: Codable, Sendable {
+ public var vocabSize: Int = 151936
+ public var hiddenSize: Int = 1024 // Model dimension
+ public var numLayers: Int = 28 // Transformer layers
+ public var numHeads: Int = 16 // Attention heads
+ public var numKVHeads: Int = 8 // KV heads for GQA
+ public var headDim: Int = 64 // Head dimension (hiddenSize / numHeads for 0.6B)
+ public var intermediateSize: Int = 3072 // FFN intermediate size
+ public var maxPositionEmbeddings: Int = 65536
+ public var rmsNormEps: Float = 1e-6
+ public var ropeTheta: Float = 1000000.0
+ public var ropeScaling: RopeScaling? = nil
+ public var tieWordEmbeddings: Bool = true
+
+ // Quantization config
+ public var groupSize: Int = 64
+ public var bits: Int = 4
+
+ public init() {}
+
+ /// Config for Qwen3-ASR-0.6B decoder, 4-bit (from HuggingFace model config)
+ public static var small: TextDecoderConfig {
+ var config = TextDecoderConfig()
+ config.hiddenSize = 1024
+ config.numLayers = 28
+ config.numHeads = 16
+ config.numKVHeads = 8
+ config.headDim = 128 // From config.json: head_dim = 128
+ config.intermediateSize = 3072
+ config.groupSize = 64
+ config.bits = 4
+ return config
+ }
+
+ /// Config for Qwen3-ASR-0.6B decoder, 8-bit
+ public static var small8bit: TextDecoderConfig {
+ var config = small
+ config.bits = 8
+ return config
+ }
+
+ /// Config for Qwen3-ASR-1.7B decoder, 4-bit
+ public static var large: TextDecoderConfig {
+ var config = TextDecoderConfig()
+ config.hiddenSize = 2048
+ config.numLayers = 28
+ config.numHeads = 16
+ config.numKVHeads = 8
+ config.headDim = 128
+ config.intermediateSize = 6144
+ config.groupSize = 64
+ config.bits = 4
+ return config
+ }
+
+ /// Config for Qwen3-ASR-1.7B decoder, 8-bit
+ public static var large8bit: TextDecoderConfig {
+ var config = large
+ config.bits = 8
+ return config
+ }
+}
+
+/// RoPE scaling configuration
+public struct RopeScaling: Codable, Sendable {
+ public var type: String
+ public var factor: Float?
+ public var originalMaxPositionEmbeddings: Int?
+
+ enum CodingKeys: String, CodingKey {
+ case type
+ case factor
+ case originalMaxPositionEmbeddings = "original_max_position_embeddings"
+ }
+}
+
+/// Combined Qwen3-ASR model configuration
+public struct Qwen3ASRConfig: Codable, Sendable {
+ public var audioEncoder: AudioEncoderConfig
+ public var textDecoder: TextDecoderConfig
+ public var audioTokenIndex: Int = 151646
+ public var eosTokenId: Int = 151645
+ public var padTokenId: Int = 151643
+
+ // ForcedAligner-specific config
+ public var classifyNum: Int = 5000
+ public var timestampSegmentTime: Float = 0.08 // 80ms per timestamp class
+
+ public init(
+ audioEncoder: AudioEncoderConfig = AudioEncoderConfig(),
+ textDecoder: TextDecoderConfig = TextDecoderConfig()
+ ) {
+ self.audioEncoder = audioEncoder
+ self.textDecoder = textDecoder
+ }
+
+ /// Config for Qwen3-ASR-0.6B
+ public static var small: Qwen3ASRConfig {
+ Qwen3ASRConfig(
+ audioEncoder: .small,
+ textDecoder: .small
+ )
+ }
+
+ /// Config for Qwen3-ASR-1.7B
+ public static var large: Qwen3ASRConfig {
+ Qwen3ASRConfig(
+ audioEncoder: .large,
+ textDecoder: .large
+ )
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/CoreMLASRModel.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/CoreMLASRModel.swift
new file mode 100644
index 0000000..dbc3f41
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/CoreMLASRModel.swift
@@ -0,0 +1,389 @@
+#if canImport(CoreML)
+import CoreML
+import Foundation
+import MLX
+import AudioCommon
+
+/// Full CoreML ASR model: CoreML encoder + CoreML text decoder.
+///
+/// Runs the entire Qwen3-ASR pipeline on CoreML (Neural Engine + CPU),
+/// eliminating the MLX GPU dependency. Requires macOS 15+ / iOS 18+
+/// for MLState KV cache support.
+public class CoreMLASRModel {
+ public let encoder: CoreMLASREncoder
+ public let decoder: CoreMLTextDecoder
+ public let featureExtractor: WhisperFeatureExtractor
+ private var tokenizer: Qwen3Tokenizer?
+
+ public init(encoder: CoreMLASREncoder, decoder: CoreMLTextDecoder) {
+ self.encoder = encoder
+ self.decoder = decoder
+ self.featureExtractor = WhisperFeatureExtractor()
+ }
+
+ /// Load full CoreML ASR from HuggingFace.
+ ///
+ /// Downloads encoder and decoder models from `aufklarer/Qwen3-ASR-CoreML`.
+ ///
+ /// Compute units are split per-component because the encoder and decoder
+ /// have different optimal backends: the encoder defaults to `.all` (the
+ /// 30 s fixed-shape graph runs well on GPU on most Macs), while the
+ /// decoder defaults to `.cpuAndNeuralEngine` (the autoregressive
+ /// MLState path is ANE-friendly and ~7× faster there than GPU per the
+ /// rebuilt-encoder PR). A single `computeUnits` parameter that
+ /// propagated into both calls silently overrode the decoder's stated
+ /// `.cpuAndNeuralEngine` default with `.all`, costing real ANE
+ /// throughput on the full-pipeline path.
+ public static func fromPretrained(
+ encoderModelId: String = CoreMLASREncoder.defaultModelId,
+ decoderModelId: String = CoreMLASREncoder.defaultModelId,
+ tokenizerModelId: String = "aufklarer/Qwen3-ASR-0.6B-MLX-4bit",
+ encoderComputeUnits: MLComputeUnits = CoreMLComputeUnitsResolver.resolved(default: .all),
+ decoderComputeUnits: MLComputeUnits = CoreMLComputeUnitsResolver.resolved(default: .cpuAndNeuralEngine),
+ cacheDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> CoreMLASRModel {
+ // Download encoder (0-30%)
+ progressHandler?(0.0, "Loading CoreML encoder...")
+ let enc = try await CoreMLASREncoder.fromPretrained(
+ modelId: encoderModelId,
+ computeUnits: encoderComputeUnits,
+ cacheDir: cacheDir,
+ offlineMode: offlineMode
+ ) { p, msg in
+ progressHandler?(p * 0.3, msg)
+ }
+
+ // Download decoder (30-80%)
+ progressHandler?(0.3, "Loading CoreML decoder...")
+ let dec = try await CoreMLTextDecoder.fromPretrained(
+ modelId: decoderModelId,
+ computeUnits: decoderComputeUnits,
+ cacheDir: cacheDir,
+ offlineMode: offlineMode
+ ) { p, msg in
+ progressHandler?(0.3 + p * 0.5, msg)
+ }
+
+ // Download tokenizer (80-90%)
+ progressHandler?(0.8, "Loading tokenizer...")
+ let tokenizerDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: tokenizerModelId)
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: tokenizerModelId,
+ to: tokenizerDir,
+ additionalFiles: ["vocab.json", "merges.txt", "tokenizer_config.json"],
+ offlineMode: offlineMode
+ )
+
+ let model = CoreMLASRModel(encoder: enc, decoder: dec)
+
+ let vocabPath = tokenizerDir.appendingPathComponent("vocab.json")
+ if FileManager.default.fileExists(atPath: vocabPath.path) {
+ let tokenizer = Qwen3Tokenizer()
+ try tokenizer.load(from: vocabPath)
+ model.tokenizer = tokenizer
+ }
+
+ progressHandler?(1.0, "Ready")
+ return model
+ }
+
+ /// Warm up both encoder and decoder.
+ public func warmUp() throws {
+ try encoder.warmUp()
+ try decoder.warmUp()
+ }
+
+ /// Transcribe audio to text using full CoreML pipeline.
+ ///
+ /// The entire inference runs on CoreML (Neural Engine + CPU) without MLX GPU.
+ public func transcribe(
+ audio: [Float],
+ sampleRate: Int = 16000,
+ language: String? = nil,
+ maxTokens: Int = 448
+ ) throws -> String {
+ let profile = ProcessInfo.processInfo.environment["COREML_ASR_PROFILE"] == "1"
+ let t0 = CFAbsoluteTimeGetCurrent()
+
+ let melFeatures = featureExtractor.process(audio, sampleRate: sampleRate)
+ let t1 = CFAbsoluteTimeGetCurrent()
+
+ // The encoder pads mel to a fixed 30 s shape and reports the real
+ // (un-padded) audio-token count via ``output_length``. We feed only
+ // the first ``numAudioTokens`` of the padded embeddings to the
+ // decoder so trailing zero-derived tokens don't pollute attention.
+ let (audioEmbeds, numAudioTokens) = try encoder.encode(melFeatures)
+ let t2 = CFAbsoluteTimeGetCurrent()
+
+ decoder.resetCache()
+
+ // Build chat template token sequence
+ let imStartId: Int32 = 151644
+ let imEndId: Int32 = 151645
+ let audioStartId: Int32 = 151669
+ let audioEndId: Int32 = 151670
+ let asrTextId: Int32 = 151704
+ let newlineId: Int32 = 198
+ let systemId: Int32 = 8948
+ let userId: Int32 = 872
+ let assistantId: Int32 = 77091
+
+ // <|im_start|>system\n<|im_end|>\n
+ var prefixTokens: [Int32] = [imStartId, systemId, newlineId, imEndId, newlineId]
+ // <|im_start|>user\n<|audio_start|>
+ prefixTokens += [imStartId, userId, newlineId, audioStartId]
+
+ // <|audio_end|><|im_end|>\n<|im_start|>assistant\n
+ var suffixTokens: [Int32] = [audioEndId, imEndId, newlineId, imStartId, assistantId, newlineId]
+
+ // Language hint + <|asr_text|>
+ if let lang = language, let tokenizer = tokenizer {
+ let langPrefix = "language \(lang)"
+ let langTokens = tokenizer.encode(langPrefix)
+ suffixTokens += langTokens.map { Int32($0) }
+ }
+ suffixTokens.append(asrTextId)
+
+ // ── Prefill: process all prefix tokens (one batched call) ──
+ var lastLogits: MLMultiArray?
+ lastLogits = try decoder.decoderPrefillTokens(prefixTokens)
+ let t3 = CFAbsoluteTimeGetCurrent()
+
+ // ── Prefill: process audio embeddings ──
+ // Bulk-extract the MLX audio embeddings once (single Metal sync)
+ // then feed them to the decoder in batched chunks of
+ // ``prefillBatchSize`` tokens. The fixed-T CoreML decoder packs
+ // T tokens per ANE dispatch, so a 20 s / 250-token clip becomes
+ // ~2 calls instead of 250 (each ANE dispatch costs ~30 ms — the
+ // dispatch overhead dominated the per-step cost in profiling).
+ let _ = audioEmbeds.dim(2) // sanity-check hidden dim
+ let audioEmbedsFlat: [Float] = audioEmbeds.asArray(Float.self)
+ let chunk = decoder.prefillBatchSize
+ var consumed = 0
+ while consumed < numAudioTokens {
+ let n = min(chunk, numAudioTokens - consumed)
+ lastLogits = try decoder.decoderPrefill(
+ flatEmbeddings: audioEmbedsFlat,
+ offset: consumed,
+ realCount: n,
+ )
+ consumed += n
+ }
+ let t4 = CFAbsoluteTimeGetCurrent()
+
+ // ── Prefill: process suffix tokens (one batched call) ──
+ lastLogits = try decoder.decoderPrefillTokens(suffixTokens)
+ let t5 = CFAbsoluteTimeGetCurrent()
+
+ // ── Autoregressive generation ──
+ guard var logits = lastLogits else {
+ return "[CoreML decoder: no output]"
+ }
+
+ // Known WER issue with this path (multi-sentence utterances on
+ // LibriSpeech test-clean): the CoreML port emits ``<|im_end|>``
+ // after the first sentence-final period with a wide logit margin
+ // (~6+ nats over the runner-up). The MLX path at the same
+ // effective bit width keeps generating. We tried both a
+ // logit-margin guard and a force-first-EOS-suppression; neither
+ // helps — the model's runner-up at the truncation point is also
+ // wrong (e.g. " The" instead of " on"), so substituting it just
+ // trades deletions for substitutions. The root cause is upstream
+ // — encoder INT8 quantization / mel padding leakage into audio
+ // embeddings, or position drift in chunked prefill. Tracked as
+ // a separate fix that requires model re-export, not a sampler
+ // change. The ``argmax(skipping:)`` / ``logit(_:at:)`` helpers
+ // stay for that future work.
+ var generatedTokens: [Int32] = []
+ var nextToken = decoder.argmax(logits: logits)
+ // First-token EOS would mean the model thinks the audio yielded an
+ // empty transcript — never right on real speech. Cheap to guard.
+ if nextToken == imEndId {
+ nextToken = decoder.argmax(logits: logits, skipping: imEndId)
+ }
+ generatedTokens.append(nextToken)
+
+ for _ in 1..") {
+ return String(rawText[range.upperBound...]).trimmingCharacters(in: .whitespaces)
+ }
+ return rawText
+ } else {
+ return generatedTokens.map { String($0) }.joined(separator: " ")
+ }
+ }
+
+ // MARK: - MLX-Free Transcription
+
+ /// Transcribe audio to text without any MLX/Metal dependency.
+ ///
+ /// Uses `featureExtractor.processRaw()` (CPU via Accelerate) and
+ /// `encoder.encode(melData:melBins:timeFrames:)` (CoreML) to produce
+ /// MLMultiArray embeddings, then decodes using `audioEmbeddingFromMultiArray()`.
+ ///
+ /// This method is safe for iOS background execution where Metal GPU eval
+ /// (triggered by MLXArray operations) would cause a crash.
+ ///
+ /// - Note: Requires `processRaw()` on WhisperFeatureExtractor and
+ /// `encode(melData:melBins:timeFrames:)` on CoreMLASREncoder, both added by T2.
+ public func transcribeWithoutMLX(
+ audio: [Float],
+ sampleRate: Int = 16000,
+ language: String? = nil,
+ maxTokens: Int = 448
+ ) throws -> String {
+ // 1. Extract mel features (pure CPU via Accelerate — no MLXArray)
+ let melFeatures = featureExtractor.processRaw(audio, sampleRate: sampleRate)
+
+ // 2. Encode audio → MLMultiArray embeddings + real (un-padded)
+ // audio-token count from the encoder's ``output_length``.
+ let encoded = try encoder.encode(
+ melData: melFeatures.data,
+ melBins: melFeatures.melBins,
+ timeFrames: melFeatures.timeFrames
+ )
+ let audioEmbeds = encoded.embeddings
+ let numAudioTokens = encoded.outputLength
+
+ // 3. Reset decoder KV cache
+ decoder.resetCache()
+
+ // 4. Build chat template token sequence (identical to transcribe())
+ let imStartId: Int32 = 151644
+ let imEndId: Int32 = 151645
+ let audioStartId: Int32 = 151669
+ let audioEndId: Int32 = 151670
+ let asrTextId: Int32 = 151704
+ let newlineId: Int32 = 198
+ let systemId: Int32 = 8948
+ let userId: Int32 = 872
+ let assistantId: Int32 = 77091
+
+ // <|im_start|>system\n<|im_end|>\n
+ var prefixTokens: [Int32] = [imStartId, systemId, newlineId, imEndId, newlineId]
+ // <|im_start|>user\n<|audio_start|>
+ prefixTokens += [imStartId, userId, newlineId, audioStartId]
+
+ // <|audio_end|><|im_end|>\n<|im_start|>assistant\n
+ var suffixTokens: [Int32] = [audioEndId, imEndId, newlineId, imStartId, assistantId, newlineId]
+
+ // Language hint + <|asr_text|>
+ if let lang = language, let tokenizer = tokenizer {
+ let langPrefix = "language \(lang)"
+ let langTokens = tokenizer.encode(langPrefix)
+ suffixTokens += langTokens.map { Int32($0) }
+ }
+ suffixTokens.append(asrTextId)
+
+ // 5. Prefill: process all prefix tokens
+ var lastLogits: MLMultiArray?
+
+ for token in prefixTokens {
+ let embedding = try decoder.embed(tokenId: token)
+ lastLogits = try decoder.decoderStep(embedding: embedding)
+ }
+
+ // 6. Prefill: process audio embeddings (MLX-free path)
+ for i in 0..") {
+ return String(rawText[range.upperBound...]).trimmingCharacters(in: .whitespaces)
+ }
+ return rawText
+ } else {
+ return generatedTokens.map { String($0) }.joined(separator: " ")
+ }
+ }
+}
+
+// MARK: - SpeechRecognitionModel
+
+extension CoreMLASRModel: SpeechRecognitionModel {
+ public var inputSampleRate: Int { 16000 }
+
+ public func transcribe(audio: [Float], sampleRate: Int, language: String?) -> String {
+ do {
+ return try transcribe(audio: audio, sampleRate: sampleRate, language: language, maxTokens: 448)
+ } catch {
+ return "[CoreML error: \(error.localizedDescription)]"
+ }
+ }
+}
+
+// MARK: - Background-Safe Transcription
+
+extension CoreMLASRModel {
+ /// Background-safe transcription (no MLX/Metal dependency).
+ ///
+ /// Uses `transcribeWithoutMLX()` which avoids all MLXArray operations
+ /// that would trigger Metal GPU eval. Safe to call from iOS background
+ /// audio processing where GPU access is prohibited.
+ public func transcribeBackgroundSafe(audio: [Float], sampleRate: Int, language: String?) -> String {
+ do {
+ return try transcribeWithoutMLX(audio: audio, sampleRate: sampleRate, language: language)
+ } catch {
+ return "[CoreML error: \(error.localizedDescription)]"
+ }
+ }
+}
+#endif
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/CoreMLEncoder.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/CoreMLEncoder.swift
new file mode 100644
index 0000000..01b2e29
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/CoreMLEncoder.swift
@@ -0,0 +1,231 @@
+#if canImport(CoreML)
+import CoreML
+import Foundation
+import MLX
+import AudioCommon
+
+/// CoreML audio encoder for Qwen3-ASR.
+///
+/// Runs the audio encoder on Neural Engine via CoreML instead of GPU via MLX.
+/// Produces audio embeddings that feed into the MLX text decoder. This enables
+/// lower power consumption on macOS and is a step toward full iOS deployment.
+///
+/// The encoder uses a single fixed 30 s mel shape ``[1, 128, 3000]`` and
+/// applies upstream's chunked block-attention (100-frame chunks → 13 tokens
+/// each, 8-chunk attention windows). Mel input is zero-padded to 3000 frames
+/// and the real length is signaled via a separate ``mel_length`` input so
+/// the in-graph block-attention bias can mask out the padded frames; the
+/// model returns the matching real audio-token count via ``output_length``.
+public class CoreMLASREncoder {
+ private let model: MLModel
+ /// Fixed mel length the chunked-attention encoder is exported with.
+ /// 3000 mel frames = 30 s @ 100 Hz hop, matching upstream training.
+ public static let paddedMelLength: Int = 3000
+ /// Max audio tokens out of the padded encoder (3000 mel / 8 conv stride ≈
+ /// 30 chunks × 13 tokens). The model writes the real count to ``output_length``.
+ public static let paddedAudioTokens: Int = 390
+
+ public static let defaultModelId = "aufklarer/Qwen3-ASR-CoreML"
+
+ /// Embeddings + the real, un-padded audio-token count (from the model's
+ /// ``output_length`` output). Callers should iterate only the first
+ /// ``outputLength`` tokens of ``embeddings``.
+ public struct EncodedAudio {
+ public let embeddings: MLMultiArray
+ public let outputLength: Int
+ }
+
+ public init(model: MLModel) {
+ self.model = model
+ }
+
+ /// Load encoder from a directory containing `encoder.mlmodelc`.
+ public static func load(
+ from directory: URL,
+ computeUnits: MLComputeUnits = CoreMLComputeUnitsResolver.resolved(default: .all)
+ ) throws -> CoreMLASREncoder {
+ let modelURL = directory.appendingPathComponent("encoder.mlmodelc", isDirectory: true)
+ guard FileManager.default.fileExists(atPath: modelURL.path) else {
+ throw AudioModelError.modelLoadFailed(
+ modelId: "encoder",
+ reason: "CoreML encoder not found at \(modelURL.path)")
+ }
+
+ let config = MLModelConfiguration()
+ config.computeUnits = computeUnits
+ let model = try MLModel(contentsOf: modelURL, configuration: config)
+ return CoreMLASREncoder(model: model)
+ }
+
+ /// Load encoder from HuggingFace.
+ public static func fromPretrained(
+ modelId: String = defaultModelId,
+ computeUnits: MLComputeUnits = CoreMLComputeUnitsResolver.resolved(default: .all),
+ cacheDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> CoreMLASREncoder {
+ let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
+
+ progressHandler?(0.0, "Downloading CoreML encoder...")
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: modelId,
+ to: cacheDir,
+ additionalFiles: ["encoder.mlmodelc/**", "config.json"],
+ offlineMode: offlineMode
+ ) { fraction in
+ progressHandler?(fraction * 0.8, "Downloading CoreML encoder...")
+ }
+
+ progressHandler?(0.9, "Loading CoreML encoder...")
+ let encoder = try load(from: cacheDir, computeUnits: computeUnits)
+ progressHandler?(1.0, "Ready")
+ return encoder
+ }
+
+ /// Warm up the encoder with a short dummy input to trigger CoreML compilation.
+ public func warmUp() throws {
+ // Use a small fake length so warmup doesn't depend on a real clip.
+ _ = try encodeRaw(melData: [Float](repeating: 0, count: 128 * Self.paddedMelLength),
+ melBins: 128, realFrames: 100)
+ }
+
+ /// Encode a mel spectrogram and return embeddings as an MLXArray.
+ ///
+ /// - Parameter melFeatures: Mel spectrogram as MLXArray `[128, T]`
+ /// - Returns: `(embeddings: [1, paddedAudioTokens, 1024], outputLength)`
+ public func encode(_ melFeatures: MLXArray) throws -> (embeddings: MLXArray, outputLength: Int) {
+ let melBins = melFeatures.dim(0)
+ let melTime = melFeatures.dim(1)
+ let melData: [Float] = melFeatures.asArray(Float.self)
+ let raw = try encodeRaw(melData: melData, melBins: melBins, realFrames: melTime)
+ return (multiArrayToMLXArray(raw.embeddings), raw.outputLength)
+ }
+
+ // MARK: - MLX-free encoding (for iOS background / pure CoreML path)
+
+ /// Encode mel spectrogram to audio embeddings without any MLXArray dependency.
+ ///
+ /// Accepts raw `[Float]` mel data in `[melBins, timeFrames]` layout (the same
+ /// layout produced by `WhisperFeatureExtractor.extractFeaturesRaw`).
+ /// Returns the encoder output as `MLMultiArray` directly, avoiding the
+ /// Metal GPU eval that `MLXArray` would trigger.
+ ///
+ /// - Parameters:
+ /// - melData: Flat float array in row-major `[melBins, timeFrames]` order
+ /// - melBins: Number of mel frequency bins (typically 128)
+ /// - timeFrames: Number of time frames
+ /// - Returns: Audio embeddings as `MLMultiArray` with shape `[1, T/8, 1024]`
+ public func encode(melData: [Float], melBins: Int, timeFrames: Int) throws -> EncodedAudio {
+ return try encodeRaw(melData: melData, melBins: melBins, realFrames: timeFrames)
+ }
+
+ /// Convenience: encode a `MelFeatures` struct directly.
+ public func encode(melFeatures: MelFeatures) throws -> EncodedAudio {
+ return try encodeRaw(melData: melFeatures.data,
+ melBins: melFeatures.melBins,
+ realFrames: melFeatures.timeFrames)
+ }
+
+ /// Shared core: zero-pads ``melData`` to the fixed ``paddedMelLength``,
+ /// runs the two-input/two-output graph, and returns the model's reported
+ /// ``output_length`` alongside the full padded embeddings.
+ private func encodeRaw(
+ melData: [Float], melBins: Int, realFrames: Int
+ ) throws -> EncodedAudio {
+ let padded = Self.paddedMelLength
+ guard realFrames <= padded else {
+ throw AudioModelError.inferenceFailed(
+ operation: "CoreML encoder",
+ reason: "Audio too long: \(realFrames) mel frames exceeds fixed shape \(padded). Segment with SpeechVAD or process in 30s windows.")
+ }
+ // Mel input: [1, melBins, paddedMelLength], zero-padded past realFrames.
+ let melArray = try MLMultiArray(
+ shape: [1, melBins as NSNumber, padded as NSNumber], dataType: .float32)
+ let mptr = melArray.dataPointer.assumingMemoryBound(to: Float.self)
+ for bin in 0.. MLXArray {
+ let shape = array.shape.map { $0.intValue }
+ let count = array.count
+
+ switch array.dataType {
+ case .float16:
+ let src = array.dataPointer.assumingMemoryBound(to: Float16.self)
+ var floats = [Float](repeating: 0, count: count)
+ for i in 0.. String {
+ let melFeatures = featureExtractor.process(audio, sampleRate: sampleRate)
+
+ // CoreML encoder returns the padded `[1, paddedAudioTokens, 1024]`
+ // embeddings + the real audio-token count via ``outputLength``.
+ let (audioEmbeds, audioLength) = try coremlEncoder.encode(melFeatures)
+
+ guard let textDecoder = textDecoder else {
+ return "[CoreML encoded: \(audioEmbeds.shape) length=\(audioLength)] - Text decoder not loaded"
+ }
+
+ // Slice the embeddings to the real (un-padded) length before passing
+ // to the MLX text decoder, otherwise trailing zero-derived tokens
+ // would pollute decoder cross-attention.
+ let realEmbeds = audioEmbeds[0..., 0.. CoreMLTextDecoder {
+ let config = MLModelConfiguration()
+ config.computeUnits = computeUnits
+
+ var maxSeq = 1024
+ var vocabSize = 151936
+ var hiddenSize = 1024
+ var batchSize = 128
+ let configPath = directory.appendingPathComponent("config.json")
+ if let data = try? Data(contentsOf: configPath),
+ let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] {
+ maxSeq = json["max_seq_length"] as? Int ?? 1024
+ vocabSize = json["vocab_size"] as? Int ?? 151936
+ hiddenSize = json["hidden_size"] as? Int ?? 1024
+ if let ts = json["enumerated_t"] as? [Int], let t = ts.first {
+ batchSize = t
+ }
+ }
+
+ let embURL = findModel(named: "embedding", in: directory)
+ let p1URL = findModel(named: "decoder_part1", in: directory)
+ let p2URL = findModel(named: "decoder_part2", in: directory)
+
+ guard let embURL else {
+ throw AudioModelError.modelLoadFailed(
+ modelId: "embedding",
+ reason: "CoreML embedding not found in \(directory.path)")
+ }
+ guard let p1URL else {
+ throw AudioModelError.modelLoadFailed(
+ modelId: "decoder_part1",
+ reason: "CoreML decoder_part1 not found in \(directory.path)")
+ }
+ guard let p2URL else {
+ throw AudioModelError.modelLoadFailed(
+ modelId: "decoder_part2",
+ reason: "CoreML decoder_part2 not found in \(directory.path)")
+ }
+
+ let embModel = try MLModel(contentsOf: embURL, configuration: config)
+ let p1Model = try MLModel(contentsOf: p1URL, configuration: config)
+ let p2Model = try MLModel(contentsOf: p2URL, configuration: config)
+
+ return CoreMLTextDecoder(
+ embeddingModel: embModel,
+ decoderPart1Model: p1Model,
+ decoderPart2Model: p2Model,
+ maxSeqLength: maxSeq,
+ vocabSize: vocabSize,
+ hiddenSize: hiddenSize,
+ batchSize: batchSize
+ )
+ }
+
+ /// Load from HuggingFace.
+ public static func fromPretrained(
+ modelId: String = defaultModelId,
+ computeUnits: MLComputeUnits = CoreMLComputeUnitsResolver.resolved(default: .cpuAndNeuralEngine),
+ cacheDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> CoreMLTextDecoder {
+ let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
+
+ progressHandler?(0.0, "Downloading CoreML decoder...")
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: modelId,
+ to: cacheDir,
+ additionalFiles: [
+ "embedding.mlmodelc/**",
+ "decoder_part1.mlmodelc/**",
+ "decoder_part2.mlmodelc/**",
+ "config.json",
+ ],
+ offlineMode: offlineMode
+ ) { fraction in
+ progressHandler?(fraction * 0.8, "Downloading CoreML decoder...")
+ }
+
+ progressHandler?(0.9, "Loading CoreML decoder...")
+ let decoder = try load(from: cacheDir, computeUnits: computeUnits)
+ progressHandler?(1.0, "Ready")
+ return decoder
+ }
+
+ /// Warm up the three models so the first real call doesn't pay
+ /// the ANE compile / load latency. Uses throwaway MLStates so the
+ /// live KV cache stays untouched.
+ public func warmUp() throws {
+ let dummyToken = try MLMultiArray(shape: [1, 1], dataType: .int32)
+ dummyToken[0] = 0
+ _ = try embeddingModel.prediction(from: MLDictionaryFeatureProvider(dictionary: [
+ "token_id": MLFeatureValue(multiArray: dummyToken),
+ ]))
+
+ let warmP1 = decoderPart1Model.makeState()
+ let warmP2 = decoderPart2Model.makeState()
+ let dummyEmbeds = try MLMultiArray(shape: [1, batchSize as NSNumber, hiddenSize as NSNumber],
+ dataType: .float32)
+ let dummyPositions = try MLMultiArray(shape: [batchSize as NSNumber], dataType: .int32)
+ for i in 0.. MLMultiArray {
+ let tokenArray = try MLMultiArray(shape: [1, 1], dataType: .int32)
+ tokenArray[0] = NSNumber(value: tokenId)
+
+ let input = try MLDictionaryFeatureProvider(dictionary: [
+ "token_id": MLFeatureValue(multiArray: tokenArray),
+ ])
+ let output = try embeddingModel.prediction(from: input)
+
+ guard let embedding = output.featureValue(for: "embedding")?.multiArrayValue else {
+ throw AudioModelError.inferenceFailed(
+ operation: "CoreML embedding", reason: "Missing embedding output")
+ }
+ return embedding
+ }
+
+ /// Run one decoder step on a single embedding.
+ ///
+ /// Packs the one real token into the last input slot and fills the
+ /// remaining ``batchSize - 1`` slots with scratch positions whose
+ /// outputs and cache writes are discarded by masking. Single-token
+ /// decode pays the same ANE dispatch cost as a full chunked prefill.
+ public func decoderStep(embedding: MLMultiArray) throws -> MLMultiArray {
+ let bufs = try writeChunk(realEmbeddingsSource: { (slot, dstPtr) in
+ Self.copyRow(from: embedding, sourceRow: 0, hidden: self.hiddenSize,
+ to: dstPtr, destSlot: slot)
+ }, realCount: 1)
+ return try runParts(embeds: bufs.embeds, positions: bufs.positions, mask: bufs.mask)
+ }
+
+ /// Copy one ``hidden``-length row from an MLMultiArray (any float
+ /// dtype, any strides) into a contiguous Float32 destination row.
+ /// The decoder input is declared Float32, but CoreML model *outputs*
+ /// (embedding lookup, part1 hidden) may come back as Float16 with
+ /// padded strides on ANE; a raw ``assumingMemoryBound(to: Float)``
+ /// copy would then read garbage.
+ private static func copyRow(from src: MLMultiArray, sourceRow: Int, hidden: Int,
+ to dst: UnsafeMutablePointer, destSlot: Int) {
+ let rowStride = src.strides.count >= 2 ? src.strides[src.strides.count - 2].intValue : hidden
+ let lastStride = src.strides.last?.intValue ?? 1
+ let base = sourceRow * rowStride
+ let dstBase = destSlot * hidden
+ switch src.dataType {
+ case .float16:
+ let p = src.dataPointer.assumingMemoryBound(to: Float16.self)
+ for j in 0.. MLMultiArray {
+ precondition(n > 0 && n <= batchSize,
+ "realCount \(n) must be in 1...\(batchSize)")
+ let bufs = try writeChunk(realEmbeddingsSource: { (slot, dstPtr) in
+ let srcPtr = embeddings.dataPointer.assumingMemoryBound(to: Float.self)
+ for t in 0.. MLMultiArray {
+ precondition(!tokenIds.isEmpty, "decoderPrefillTokens requires at least one token")
+ var lastLogits: MLMultiArray!
+ var consumed = 0
+ while consumed < tokenIds.count {
+ let n = min(batchSize, tokenIds.count - consumed)
+ // Pack n token embeddings into a [1, n, hidden] Float32 buffer.
+ let packed = try MLMultiArray(shape: [1, n as NSNumber, hiddenSize as NSNumber],
+ dataType: .float32)
+ let pptr = packed.dataPointer.assumingMemoryBound(to: Float.self)
+ for k in 0.. MLMultiArray {
+ precondition(n > 0 && n <= batchSize,
+ "realCount \(n) must be in 1...\(batchSize)")
+ let bufs = try writeChunk(realEmbeddingsSource: { (slot, dstPtr) in
+ flatEmbeddings.withUnsafeBufferPointer { buf in
+ let src = buf.baseAddress!
+ for t in 0..) -> Void,
+ realCount n: Int
+ ) throws -> (embeds: MLMultiArray, positions: MLMultiArray, mask: MLMultiArray) {
+ precondition(currentPosition + n <= scratchStart,
+ "Cache overflow: would write real position \(currentPosition + n - 1) into scratch range starting at \(scratchStart)")
+
+ let T = batchSize
+ let firstRealSlot = T - n
+
+ let embeds = try MLMultiArray(shape: [1, T as NSNumber, hiddenSize as NSNumber],
+ dataType: .float32)
+ let positions = try MLMultiArray(shape: [T as NSNumber], dataType: .int32)
+ let mask = try MLMultiArray(shape: [1, 1, T as NSNumber, maxSeqLength as NSNumber],
+ dataType: .float32)
+
+ // Positions: scratch slots fill 0..firstRealSlot-1, real fills firstRealSlot..T-1
+ for i in 0.. MLMultiArray {
+ let p1Input = try MLDictionaryFeatureProvider(dictionary: [
+ "input_embeds": MLFeatureValue(multiArray: embeds),
+ "positions": MLFeatureValue(multiArray: positions),
+ "attention_mask": MLFeatureValue(multiArray: mask),
+ ])
+ let p1Out = try decoderPart1Model.prediction(from: p1Input, using: part1State)
+ guard let hidden = p1Out.featureValue(for: "hidden_state")?.multiArrayValue else {
+ throw AudioModelError.inferenceFailed(
+ operation: "CoreML decoder part1",
+ reason: "Missing hidden_state output")
+ }
+
+ let p2Input = try MLDictionaryFeatureProvider(dictionary: [
+ "input_embeds": MLFeatureValue(multiArray: hidden),
+ "positions": MLFeatureValue(multiArray: positions),
+ "attention_mask": MLFeatureValue(multiArray: mask),
+ ])
+ let p2Out = try decoderPart2Model.prediction(from: p2Input, using: part2State)
+ guard let logits = p2Out.featureValue(for: "logits")?.multiArrayValue else {
+ throw AudioModelError.inferenceFailed(
+ operation: "CoreML decoder part2",
+ reason: "Missing logits output")
+ }
+ return logits
+ }
+
+ /// Expose the fixed batch size so callers can chunk audio prefill.
+ public var prefillBatchSize: Int { batchSize }
+
+ /// Get argmax token ID, optionally excluding a single token.
+ ///
+ /// When ``skipToken`` is non-nil, the slot at that index is ignored —
+ /// used by the ASR generation loop to suppress premature ``<|im_end|>``
+ /// (see `CoreMLASRModel.transcribe`). When it's nil this is the plain
+ /// argmax over the full vocab.
+ public func argmax(logits: MLMultiArray, skipping skipToken: Int32? = nil) -> Int32 {
+ return argmaxImpl(logits: logits, skipIdx: skipToken.map { Int($0) })
+ }
+
+ /// Read a single logit by token index. Stride-aware (matches `argmax`).
+ public func logit(_ logits: MLMultiArray, at index: Int32) -> Float {
+ let lastStride = logits.strides.last?.intValue ?? 1
+ let i = Int(index) * lastStride
+ switch logits.dataType {
+ case .float16:
+ let ptr = logits.dataPointer.assumingMemoryBound(to: Float16.self)
+ return Float(ptr[i])
+ case .float32:
+ let ptr = logits.dataPointer.assumingMemoryBound(to: Float.self)
+ return ptr[i]
+ default:
+ let ptr = logits.dataPointer.assumingMemoryBound(to: Float.self)
+ return ptr[i]
+ }
+ }
+
+ /// Get argmax token ID from logits.
+ ///
+ /// Stride-aware: walks ``vocabSize`` (the logical last-dim length)
+ /// using ``strides.last`` as the step, correct for CoreML outputs
+ /// that may be strided (e.g. ANE padding). NaN-safe — NaN values
+ /// are skipped, so one bad logit can't poison the argmax (the
+ /// previous flat ``ptr[i]`` loop with ``maxVal = -Float.infinity``
+ /// would silently keep ``maxIdx = 0`` since IEEE-754 ``NaN > x``
+ /// is always false).
+ private func argmaxImpl(logits: MLMultiArray, skipIdx: Int?) -> Int32 {
+ let vocab = logits.shape.last?.intValue ?? logits.count
+ let lastStride = logits.strides.last?.intValue ?? 1
+ var maxVal: Float = -Float.infinity
+ var maxIdx: Int32 = 0
+ var nanCount: Int = 0
+
+ switch logits.dataType {
+ case .float16:
+ let ptr = logits.dataPointer.assumingMemoryBound(to: Float16.self)
+ for i in 0.. maxVal {
+ maxVal = val
+ maxIdx = Int32(i)
+ }
+ }
+ case .float32:
+ let ptr = logits.dataPointer.assumingMemoryBound(to: Float.self)
+ for i in 0.. maxVal {
+ maxVal = val
+ maxIdx = Int32(i)
+ }
+ }
+ default:
+ let ptr = logits.dataPointer.assumingMemoryBound(to: Float.self)
+ for i in 0.. maxVal {
+ maxVal = val
+ maxIdx = Int32(i)
+ }
+ }
+ }
+
+ return maxIdx
+ }
+
+ // MARK: - Audio Embedding Injection
+
+ /// Convert MLXArray audio embeddings to MLMultiArray for decoder input.
+ public func audioEmbeddingToMultiArray(_ embedding: MLXArray, at index: Int) throws -> MLMultiArray {
+ let hidden = embedding.dim(2)
+ let result = try MLMultiArray(shape: [1, 1, hidden as NSNumber], dataType: .float32)
+ let ptr = result.dataPointer.assumingMemoryBound(to: Float.self)
+ let slice = embedding[0..., index..<(index + 1), 0...]
+ let data: [Float] = slice.asArray(Float.self)
+ for i in 0.. MLMultiArray {
+ let hidden = embeddings.shape[2].intValue
+ let result = try MLMultiArray(shape: [1, 1, hidden as NSNumber], dataType: .float32)
+ let srcPtr = embeddings.dataPointer.assumingMemoryBound(to: Float.self)
+ let dstPtr = result.dataPointer.assumingMemoryBound(to: Float.self)
+ let offset = index * hidden
+ for i in 0.. URL? {
+ let compiled = directory.appendingPathComponent("\(name).mlmodelc", isDirectory: true)
+ if FileManager.default.fileExists(atPath: compiled.path) {
+ return compiled
+ }
+ return nil
+ }
+}
+#endif
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/ExportedImports.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/ExportedImports.swift
new file mode 100644
index 0000000..29500d8
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/ExportedImports.swift
@@ -0,0 +1,3 @@
+// Re-export AudioCommon so host apps linking Qwen3ASR can use
+// `ModelRegistry` and other download helpers without a separate product.
+@_exported import AudioCommon
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/FloatTextDecoder.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/FloatTextDecoder.swift
new file mode 100644
index 0000000..d6bdd04
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/FloatTextDecoder.swift
@@ -0,0 +1,240 @@
+import Foundation
+import MLX
+import MLXNN
+import MLXFast
+import MLXCommon
+import AudioCommon
+
+/// Protocol abstracting the text decoder for the ForcedAligner,
+/// allowing both quantized and float (bf16) implementations.
+public protocol ForcedAlignerTextDecoding: AnyObject {
+ func embeddings(for inputIds: MLXArray) -> MLXArray
+ func decode(
+ inputsEmbeds: MLXArray,
+ attentionMask: MLXArray?,
+ cache: [(MLXArray, MLXArray)]?
+ ) -> (MLXArray, [(MLXArray, MLXArray)])
+}
+
+extension QuantizedTextModel: ForcedAlignerTextDecoding {
+ public func embeddings(for inputIds: MLXArray) -> MLXArray {
+ embedTokens(inputIds)
+ }
+
+ public func decode(
+ inputsEmbeds: MLXArray,
+ attentionMask: MLXArray?,
+ cache: [(MLXArray, MLXArray)]?
+ ) -> (MLXArray, [(MLXArray, MLXArray)]) {
+ self(inputIds: nil, inputsEmbeds: inputsEmbeds, attentionMask: attentionMask, cache: cache)
+ }
+}
+
+// MARK: - Float (non-quantized) text decoder
+
+public class FloatTextAttention: Module {
+ let numHeads: Int
+ let numKVHeads: Int
+ let headDim: Int
+ let scale: Float
+
+ @ModuleInfo var qProj: Linear
+ @ModuleInfo var kProj: Linear
+ @ModuleInfo var vProj: Linear
+ @ModuleInfo var oProj: Linear
+ @ModuleInfo var qNorm: RMSNorm
+ @ModuleInfo var kNorm: RMSNorm
+
+ let rope: MLXNN.RoPE
+
+ public init(config: TextDecoderConfig) {
+ self.numHeads = config.numHeads
+ self.numKVHeads = config.numKVHeads
+ self.headDim = config.headDim
+ self.scale = 1.0 / sqrt(Float(headDim))
+
+ let hiddenSize = config.hiddenSize
+
+ self._qProj.wrappedValue = Linear(hiddenSize, numHeads * headDim, bias: false)
+ self._kProj.wrappedValue = Linear(hiddenSize, numKVHeads * headDim, bias: false)
+ self._vProj.wrappedValue = Linear(hiddenSize, numKVHeads * headDim, bias: false)
+ self._oProj.wrappedValue = Linear(numHeads * headDim, hiddenSize, bias: false)
+
+ self._qNorm.wrappedValue = RMSNorm(dimensions: headDim, eps: config.rmsNormEps)
+ self._kNorm.wrappedValue = RMSNorm(dimensions: headDim, eps: config.rmsNormEps)
+
+ self.rope = MLXNN.RoPE(dimensions: headDim, traditional: false, base: config.ropeTheta)
+ super.init()
+ }
+
+ public func callAsFunction(
+ _ hiddenStates: MLXArray,
+ attentionMask: MLXArray? = nil,
+ cache: (MLXArray, MLXArray)? = nil
+ ) -> (MLXArray, (MLXArray, MLXArray)) {
+ let (batch, seqLen, _) = (hiddenStates.dim(0), hiddenStates.dim(1), hiddenStates.dim(2))
+
+ var queries = qProj(hiddenStates)
+ var keys = kProj(hiddenStates)
+ var values = vProj(hiddenStates)
+
+ queries = queries.reshaped(batch, seqLen, numHeads, headDim)
+ keys = keys.reshaped(batch, seqLen, numKVHeads, headDim)
+ values = values.reshaped(batch, seqLen, numKVHeads, headDim)
+
+ queries = qNorm(queries)
+ keys = kNorm(keys)
+
+ queries = queries.transposed(0, 2, 1, 3)
+ keys = keys.transposed(0, 2, 1, 3)
+ values = values.transposed(0, 2, 1, 3)
+
+ let offset = cache?.0.dim(2) ?? 0
+ queries = rope(queries, offset: offset)
+ keys = rope(keys, offset: offset)
+
+ var cachedKeys = keys
+ var cachedValues = values
+
+ if let (prevKeys, prevValues) = cache {
+ cachedKeys = concatenated([prevKeys, keys], axis: 2)
+ cachedValues = concatenated([prevValues, values], axis: 2)
+ }
+
+ let merged = SDPA.attendAndMerge(
+ qHeads: queries, kHeads: cachedKeys, vHeads: cachedValues,
+ scale: scale, mask: attentionMask)
+ let output = oProj(merged)
+ return (output, (cachedKeys, cachedValues))
+ }
+}
+
+public class FloatTextMLP: Module {
+ @ModuleInfo var gateProj: Linear
+ @ModuleInfo var upProj: Linear
+ @ModuleInfo var downProj: Linear
+
+ public init(config: TextDecoderConfig) {
+ let hiddenSize = config.hiddenSize
+ let intermediateSize = config.intermediateSize
+
+ self._gateProj.wrappedValue = Linear(hiddenSize, intermediateSize, bias: false)
+ self._upProj.wrappedValue = Linear(hiddenSize, intermediateSize, bias: false)
+ self._downProj.wrappedValue = Linear(intermediateSize, hiddenSize, bias: false)
+
+ super.init()
+ }
+
+ public func callAsFunction(_ x: MLXArray) -> MLXArray {
+ let gate = silu(gateProj(x))
+ let up = upProj(x)
+ return downProj(gate * up)
+ }
+}
+
+public class FloatTextDecoderLayer: Module {
+ @ModuleInfo var selfAttn: FloatTextAttention
+ @ModuleInfo var mlp: FloatTextMLP
+ @ModuleInfo var inputLayerNorm: RMSNorm
+ @ModuleInfo var postAttentionLayerNorm: RMSNorm
+
+ public init(config: TextDecoderConfig) {
+ self._selfAttn.wrappedValue = FloatTextAttention(config: config)
+ self._mlp.wrappedValue = FloatTextMLP(config: config)
+ self._inputLayerNorm.wrappedValue = RMSNorm(dimensions: config.hiddenSize, eps: config.rmsNormEps)
+ self._postAttentionLayerNorm.wrappedValue = RMSNorm(dimensions: config.hiddenSize, eps: config.rmsNormEps)
+ super.init()
+ }
+
+ public func callAsFunction(
+ _ hiddenStates: MLXArray,
+ attentionMask: MLXArray? = nil,
+ cache: (MLXArray, MLXArray)? = nil
+ ) -> (MLXArray, (MLXArray, MLXArray)) {
+ let residual = hiddenStates
+ var hidden = inputLayerNorm(hiddenStates)
+ let (attnOutput, newCache) = selfAttn(hidden, attentionMask: attentionMask, cache: cache)
+ hidden = residual + attnOutput
+
+ let residual2 = hidden
+ hidden = postAttentionLayerNorm(hidden)
+ hidden = mlp(hidden)
+ hidden = residual2 + hidden
+
+ return (hidden, newCache)
+ }
+}
+
+public class FloatTextModel: Module {
+ public let config: TextDecoderConfig
+
+ @ModuleInfo public var embedTokens: Embedding
+ @ModuleInfo var layers: [FloatTextDecoderLayer]
+ @ModuleInfo var norm: RMSNorm
+
+ public init(config: TextDecoderConfig) {
+ self.config = config
+ self._embedTokens.wrappedValue = Embedding(embeddingCount: config.vocabSize, dimensions: config.hiddenSize)
+ self._layers.wrappedValue = (0.. (MLXArray, [(MLXArray, MLXArray)]) {
+ var hiddenStates: MLXArray
+ if let embeds = inputsEmbeds {
+ hiddenStates = embeds
+ } else if let ids = inputIds {
+ hiddenStates = embedTokens(ids)
+ } else {
+ fatalError("Either inputIds or inputsEmbeds must be provided")
+ }
+
+ let seqLen = hiddenStates.dim(1)
+
+ let mask: MLXArray?
+ if let providedMask = attentionMask {
+ mask = providedMask
+ } else if seqLen == 1 {
+ mask = nil
+ } else {
+ let cacheLen = cache?.first?.0.dim(2) ?? 0
+ let totalLen = seqLen + cacheLen
+ let rows = (MLXArray(0.. rows, MLXArray(Float(-1e9)), MLXArray(Float(0)))
+ .expandedDimensions(axes: [0, 1])
+ .asType(hiddenStates.dtype)
+ }
+
+ var newCache: [(MLXArray, MLXArray)] = []
+ for (i, layer) in layers.enumerated() {
+ let layerCache = cache?[i]
+ let (output, updatedCache) = layer(hiddenStates, attentionMask: mask, cache: layerCache)
+ hiddenStates = output
+ newCache.append(updatedCache)
+ }
+
+ hiddenStates = norm(hiddenStates)
+ return (hiddenStates, newCache)
+ }
+}
+
+extension FloatTextModel: ForcedAlignerTextDecoding {
+ public func embeddings(for inputIds: MLXArray) -> MLXArray {
+ embedTokens(inputIds)
+ }
+
+ public func decode(
+ inputsEmbeds: MLXArray,
+ attentionMask: MLXArray?,
+ cache: [(MLXArray, MLXArray)]?
+ ) -> (MLXArray, [(MLXArray, MLXArray)]) {
+ self(inputIds: nil, inputsEmbeds: inputsEmbeds, attentionMask: attentionMask, cache: cache)
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/ForcedAligner+Protocols.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/ForcedAligner+Protocols.swift
new file mode 100644
index 0000000..9d3cb09
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/ForcedAligner+Protocols.swift
@@ -0,0 +1,9 @@
+import AudioCommon
+
+// MARK: - ForcedAlignmentModel
+
+extension Qwen3ForcedAligner: ForcedAlignmentModel {
+ public func align(audio: [Float], text: String, sampleRate: Int, language: String?) -> [AlignedWord] {
+ align(audio: audio, text: text, sampleRate: sampleRate, language: language ?? "English")
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/ForcedAligner.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/ForcedAligner.swift
new file mode 100644
index 0000000..3a3a836
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/ForcedAligner.swift
@@ -0,0 +1,482 @@
+import Foundation
+import MLXCommon
+import MLX
+import MLXNN
+import MLXFast
+import AudioCommon
+
+// AlignedWord is defined in AudioCommon/Protocols.swift and re-exported via AudioCommon import above.
+
+/// Forced aligner model variant
+public enum ForcedAlignerVariant: String, CaseIterable, Sendable {
+ case mlx4bit = "aufklarer/Qwen3-ForcedAligner-0.6B-4bit"
+ case mlx8bit = "aufklarer/Qwen3-ForcedAligner-0.6B-8bit"
+ case bf16 = "aufklarer/Qwen3-ForcedAligner-0.6B-bf16"
+
+ /// Detect variant from model ID string
+ public static func detect(from modelId: String) -> ForcedAlignerVariant? {
+ if let exact = Self.allCases.first(where: { $0.rawValue == modelId }) {
+ return exact
+ }
+ if modelId.contains("bf16") || modelId.contains("float") { return .bf16 }
+ if modelId.contains("8bit") { return .mlx8bit }
+ if modelId.contains("4bit") { return .mlx4bit }
+ return nil
+ }
+
+ public var textConfig: TextDecoderConfig {
+ switch self {
+ case .mlx4bit:
+ var cfg = TextDecoderConfig.small
+ cfg.bits = 4
+ cfg.groupSize = 64
+ return cfg
+ case .mlx8bit:
+ var cfg = TextDecoderConfig.small
+ cfg.bits = 8
+ cfg.groupSize = 64
+ return cfg
+ case .bf16:
+ return .small
+ }
+ }
+
+ public var usesFloatTextDecoder: Bool {
+ self == .bf16
+ }
+}
+
+/// Qwen3 Forced Aligner — predicts word-level timestamps for audio+text pairs.
+///
+/// Uses the same encoder-decoder architecture as Qwen3-ASR but replaces the
+/// vocab lm_head with a 5000-class timestamp classification head.
+/// Inference is non-autoregressive (single forward pass).
+public class Qwen3ForcedAligner {
+ public let audioEncoder: Qwen3AudioEncoder
+ public let textDecoder: any ForcedAlignerTextDecoding
+ public let classifyHead: Linear
+ public let featureExtractor: WhisperFeatureExtractor
+ public var tokenizer: Qwen3Tokenizer?
+
+ private let config: Qwen3ASRConfig
+
+ public init(
+ audioConfig: Qwen3AudioEncoderConfig = .forcedAligner,
+ textConfig: TextDecoderConfig = .small,
+ classifyNum: Int = 5000,
+ useFloatTextDecoder: Bool = false
+ ) {
+ self.audioEncoder = Qwen3AudioEncoder(config: audioConfig)
+ if useFloatTextDecoder {
+ self.textDecoder = FloatTextModel(config: textConfig)
+ } else {
+ self.textDecoder = QuantizedTextModel(config: textConfig)
+ }
+ self.classifyHead = Linear(textConfig.hiddenSize, classifyNum)
+ self.featureExtractor = WhisperFeatureExtractor()
+
+ var cfg = Qwen3ASRConfig()
+ cfg.classifyNum = classifyNum
+ self.config = cfg
+ }
+
+ /// Align text to audio with automatic chunking for long inputs.
+ ///
+ /// The underlying classifier head emits a fixed-resolution timestamp
+ /// index (default `classifyNum=5000` × `0.08s` per slot = 400s
+ /// addressable range) but in practice the model's reliable range is
+ /// shorter (~270s on Qwen3-ForcedAligner-0.6B-4bit observed on TED-Ed
+ /// material). Past that, it produces low/non-monotonic indices that
+ /// LIS correction collapses into a flat plateau — every trailing word
+ /// shares the same timestamp.
+ ///
+ /// `alignLong` runs `align` on the full audio, detects the trailing
+ /// plateau, keeps the reliable prefix, then re-aligns the remaining
+ /// audio + remaining words and offsets timestamps. Iterates until no
+ /// plateau remains or the remaining work is below a minimum chunk size.
+ ///
+ /// For audio shorter than the threshold this is a one-pass call into
+ /// `align`. For longer audio it pays one extra align pass per chunk.
+ public func alignLong(
+ audio: [Float],
+ text: String,
+ sampleRate: Int = 16000,
+ language: String = "English",
+ progressHandler: ((String) -> Void)? = nil
+ ) -> [AlignedWord] {
+ // The model is reliable up to ~270s on the bundles we ship; we
+ // don't try to be too aggressive with the threshold so the
+ // single-pass case stays the common path. The plateau detector
+ // does the actual work — this is just a fast bypass when there's
+ // no risk of saturation.
+ let bypassThresholdSeconds: Float = 240
+ let minChunkSeconds: Float = 5
+ let plateauTolerance: Float = 0.1 // seconds; "same start time" if diff < this
+ let plateauMinWords = 5 // need ≥ N stuck words to call it a plateau
+
+ var allAligned: [AlignedWord] = []
+ var remainingAudio = audio
+ var remainingText = text
+ var offsetSec: Float = 0
+ var pass = 1
+
+ while !remainingAudio.isEmpty && !remainingText.isEmpty {
+ let durationSec = Float(remainingAudio.count) / Float(sampleRate)
+ let aligned = align(
+ audio: remainingAudio,
+ text: remainingText,
+ sampleRate: sampleRate,
+ language: language
+ )
+ if aligned.isEmpty { break }
+
+ // Skip plateau detection on small chunks — the model is reliable
+ // there, and detecting plateau on tiny outputs creates spurious
+ // splits.
+ if durationSec <= bypassThresholdSeconds || aligned.count < plateauMinWords * 2 {
+ allAligned.append(contentsOf: Self.offsetWords(aligned, by: offsetSec))
+ break
+ }
+
+ let plateauStart = Self.findTrailingPlateauStart(
+ aligned, tolerance: plateauTolerance, minSize: plateauMinWords
+ )
+ if plateauStart == aligned.count {
+ // No plateau — alignment looks healthy.
+ allAligned.append(contentsOf: Self.offsetWords(aligned, by: offsetSec))
+ break
+ }
+
+ // Take the reliable prefix; recurse on the remainder.
+ let reliable = aligned.prefix(plateauStart)
+ let splitTime = reliable.last!.endTime
+ allAligned.append(contentsOf: Self.offsetWords(Array(reliable), by: offsetSec))
+
+ let splitSample = Int(splitTime * Float(sampleRate))
+ guard splitSample < remainingAudio.count else { break }
+ let nextAudio = Array(remainingAudio[splitSample...])
+ let remainingDuration = Float(nextAudio.count) / Float(sampleRate)
+ if remainingDuration < minChunkSeconds { break }
+
+ // Pull the words from the remainder by name. We split on the
+ // same boundary `align` used (whitespace) so words line up.
+ let wordsAll = remainingText.split(separator: " ", omittingEmptySubsequences: true)
+ guard plateauStart < wordsAll.count else { break }
+ let nextText = wordsAll[plateauStart...].joined(separator: " ")
+
+ progressHandler?(
+ "Audio \(String(format: "%.1f", durationSec))s saturated after word \(plateauStart) "
+ + "(\(String(format: "%.1f", splitTime))s); chunking remaining \(String(format: "%.1f", remainingDuration))s "
+ + "(pass \(pass + 1))"
+ )
+
+ remainingAudio = nextAudio
+ remainingText = nextText
+ offsetSec += splitTime
+ pass += 1
+ if pass > 10 { break } // belt-and-braces against pathological loops
+ }
+
+ return allAligned
+ }
+
+ static func offsetWords(_ words: [AlignedWord], by seconds: Float) -> [AlignedWord] {
+ guard seconds != 0 else { return words }
+ return words.map {
+ AlignedWord(text: $0.text, startTime: $0.startTime + seconds, endTime: $0.endTime + seconds)
+ }
+ }
+
+ /// Index of the first word in the trailing "stuck" plateau, or
+ /// `aligned.count` if no plateau is detected.
+ ///
+ /// A plateau is `≥ minSize` consecutive trailing words whose start
+ /// times differ by less than `tolerance`. This is the LIS-clamp
+ /// signature: the model produced low/garbage indices for those
+ /// positions and the monotonicity pass collapsed them onto the last
+ /// reliable anchor.
+ static func findTrailingPlateauStart(
+ _ aligned: [AlignedWord], tolerance: Float, minSize: Int
+ ) -> Int {
+ let n = aligned.count
+ guard n > minSize else { return n }
+ // Walk backward: when `aligned[i].startTime ≈ aligned[i-1].startTime`,
+ // both are in the plateau, so the plateau extends *to* index `i-1`.
+ // Stop at the first big jump.
+ var plateauStart = n
+ for i in (1..= minSize ? plateauStart : n
+ }
+
+ /// Align text to audio, producing word-level timestamps.
+ ///
+ /// - Parameters:
+ /// - audio: Raw audio samples (mono)
+ /// - text: Text to align against the audio
+ /// - sampleRate: Sample rate of the audio (default 16000)
+ /// - language: Language hint for word splitting (default "English")
+ /// - Returns: Array of words with start/end timestamps in seconds
+ public func align(
+ audio: [Float],
+ text: String,
+ sampleRate: Int = 16000,
+ language: String = "English"
+ ) -> [AlignedWord] {
+ guard let tokenizer = tokenizer else {
+ print("Error: tokenizer not loaded")
+ return []
+ }
+
+ // 1. Extract mel features → audio encoder → audio embeddings
+ let melFeatures = featureExtractor.process(audio, sampleRate: sampleRate)
+ let batchedFeatures = melFeatures.expandedDimensions(axis: 0)
+ var audioEmbeds = audioEncoder(batchedFeatures)
+ audioEmbeds = audioEmbeds.expandedDimensions(axis: 0) // [1, T_audio, hiddenSize]
+
+ let numAudioTokens = audioEmbeds.dim(1)
+
+ // 2. Prepare text with timestamp slots
+ let slotted = TextPreprocessor.prepareForAlignment(
+ text: text,
+ tokenizer: tokenizer,
+ language: language
+ )
+
+ guard !slotted.words.isEmpty else {
+ print("Warning: no words found in text")
+ return []
+ }
+
+ // 3. Build input_ids with chat template
+ let inputIds = buildInputIds(
+ slottedTokenIds: slotted.tokenIds,
+ numAudioTokens: numAudioTokens,
+ tokenizer: tokenizer,
+ language: language
+ )
+
+ // Track where the slotted text starts in the full sequence
+ let slottedTextStart = inputIds.count - slotted.tokenIds.count
+
+ // 4. Embed all tokens and replace audio_pad with audio embeddings
+ let inputIdsTensor = MLXArray(inputIds.map { Int32($0) }).expandedDimensions(axis: 0)
+ var inputEmbeds = textDecoder.embeddings(for: inputIdsTensor)
+
+ // Find audio_pad range and replace with audio embeddings
+ let audioStartIndex = findAudioPadStart(inputIds)
+ let audioEndIndex = audioStartIndex + numAudioTokens
+
+ let audioEmbedsTyped = audioEmbeds.asType(inputEmbeds.dtype)
+ let beforeAudio = inputEmbeds[0..., 0..= start
+ ))
+ }
+
+ return alignedWords
+ }
+
+ // MARK: - Private Helpers
+
+ /// Build full input_ids sequence with chat template
+ private func buildInputIds(
+ slottedTokenIds: [Int],
+ numAudioTokens: Int,
+ tokenizer: Qwen3Tokenizer,
+ language: String
+ ) -> [Int] {
+ let imStartId = Qwen3ASRTokens.imStartTokenId
+ let imEndId = Qwen3ASRTokens.imEndTokenId
+ let audioStartId = Qwen3ASRTokens.audioStartTokenId
+ let audioEndId = Qwen3ASRTokens.audioEndTokenId
+ let audioPadId = Qwen3ASRTokens.audioTokenId
+ let newlineId = 198
+
+ // Token IDs for role names
+ let systemId = 8948
+ let userId = 872
+ let assistantId = 77091
+
+ var ids: [Int] = []
+
+ // <|im_start|>system\n<|im_end|>\n
+ ids.append(contentsOf: [imStartId, systemId, newlineId, imEndId, newlineId])
+
+ // <|im_start|>user\n<|audio_start|>
+ ids.append(contentsOf: [imStartId, userId, newlineId, audioStartId])
+
+ // <|audio_pad|> * numAudioTokens
+ for _ in 0..<|im_end|>\n
+ ids.append(contentsOf: [audioEndId, imEndId, newlineId])
+
+ // <|im_start|>assistant\n
+ ids.append(contentsOf: [imStartId, assistantId, newlineId])
+
+ // Slotted text with timestamp tokens
+ ids.append(contentsOf: slottedTokenIds)
+
+ return ids
+ }
+
+ /// Find the start index of audio_pad tokens in input_ids
+ private func findAudioPadStart(_ inputIds: [Int]) -> Int {
+ let audioPadId = Qwen3ASRTokens.audioTokenId
+ for (i, id) in inputIds.enumerated() {
+ if id == audioPadId { return i }
+ }
+ return 0
+ }
+}
+
+// MARK: - Model Loading
+
+public extension Qwen3ForcedAligner {
+
+ /// Load forced aligner model from HuggingFace hub
+ static func fromPretrained(
+ modelId: String = "aufklarer/Qwen3-ForcedAligner-0.6B-4bit",
+ cacheDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> Qwen3ForcedAligner {
+ progressHandler?(0.0, "Downloading model...")
+
+ let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
+
+ // Download weights and tokenizer files
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: modelId,
+ to: cacheDir,
+ additionalFiles: ["vocab.json", "merges.txt", "tokenizer_config.json",
+ "quantize_config.json"],
+ offlineMode: offlineMode,
+ progressHandler: { progress in
+ progressHandler?(progress * 0.8, "Downloading weights...")
+ }
+ )
+
+ progressHandler?(0.80, "Loading tokenizer...")
+
+ // Detect variant: try quantize_config.json first, fall back to model ID
+ let variant: ForcedAlignerVariant
+ let packaging = detectPackaging(in: cacheDir)
+ if let detected = packaging {
+ variant = detected
+ } else if let detected = ForcedAlignerVariant.detect(from: modelId) {
+ variant = detected
+ } else {
+ variant = .mlx4bit
+ }
+
+ let model = Qwen3ForcedAligner(
+ audioConfig: .forcedAligner,
+ textConfig: variant.textConfig,
+ useFloatTextDecoder: variant.usesFloatTextDecoder
+ )
+
+ // Load tokenizer
+ let vocabPath = cacheDir.appendingPathComponent("vocab.json")
+ if FileManager.default.fileExists(atPath: vocabPath.path) {
+ let tokenizer = Qwen3Tokenizer()
+ try tokenizer.load(from: vocabPath)
+ model.tokenizer = tokenizer
+ }
+
+ progressHandler?(0.85, "Loading audio encoder weights...")
+
+ // Load weights
+ try WeightLoader.loadForcedAlignerWeights(into: model, from: cacheDir)
+
+ progressHandler?(1.0, "Ready")
+
+ return model
+ }
+
+ /// Detect model variant from quantize_config.json
+ private static func detectPackaging(in cacheDir: URL) -> ForcedAlignerVariant? {
+ struct QuantizationFile: Decodable {
+ struct Quantization: Decodable {
+ let bits: Int?
+ let groupSize: Int?
+ enum CodingKeys: String, CodingKey {
+ case bits
+ case groupSize = "group_size"
+ }
+ }
+ let quantization: Quantization?
+ }
+
+ let configPath = cacheDir.appendingPathComponent("quantize_config.json")
+ guard let data = try? Data(contentsOf: configPath),
+ let file = try? JSONDecoder().decode(QuantizationFile.self, from: data),
+ let quant = file.quantization,
+ let bits = quant.bits else {
+ return nil
+ }
+
+ switch bits {
+ case 0: return .bf16
+ case 4: return .mlx4bit
+ case 8: return .mlx8bit
+ default: return nil
+ }
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/QuantizedTextDecoder.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/QuantizedTextDecoder.swift
new file mode 100644
index 0000000..bd0afa1
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/QuantizedTextDecoder.swift
@@ -0,0 +1,252 @@
+import Foundation
+import MLXCommon
+import MLX
+import MLXNN
+import MLXFast
+import AudioCommon
+
+/// Multi-head attention for Qwen3 text decoder with GQA and RoPE (quantized version)
+public class QuantizedTextAttention: Module {
+ let numHeads: Int
+ let numKVHeads: Int
+ let headDim: Int
+ let scale: Float
+
+ @ModuleInfo var qProj: QuantizedLinear
+ @ModuleInfo var kProj: QuantizedLinear
+ @ModuleInfo var vProj: QuantizedLinear
+ @ModuleInfo var oProj: QuantizedLinear
+ @ModuleInfo var qNorm: RMSNorm
+ @ModuleInfo var kNorm: RMSNorm
+
+ let rope: MLXNN.RoPE
+
+ public init(config: TextDecoderConfig) {
+ self.numHeads = config.numHeads
+ self.numKVHeads = config.numKVHeads
+ self.headDim = config.headDim
+ self.scale = 1.0 / sqrt(Float(headDim))
+
+ let hiddenSize = config.hiddenSize
+
+ // Create quantized linear layers
+ self._qProj.wrappedValue = QuantizedLinear(
+ hiddenSize, numHeads * headDim, bias: false,
+ groupSize: config.groupSize, bits: config.bits)
+ self._kProj.wrappedValue = QuantizedLinear(
+ hiddenSize, numKVHeads * headDim, bias: false,
+ groupSize: config.groupSize, bits: config.bits)
+ self._vProj.wrappedValue = QuantizedLinear(
+ hiddenSize, numKVHeads * headDim, bias: false,
+ groupSize: config.groupSize, bits: config.bits)
+ self._oProj.wrappedValue = QuantizedLinear(
+ numHeads * headDim, hiddenSize, bias: false,
+ groupSize: config.groupSize, bits: config.bits)
+
+ // Q/K normalization (Qwen3 specific)
+ self._qNorm.wrappedValue = RMSNorm(dimensions: headDim, eps: config.rmsNormEps)
+ self._kNorm.wrappedValue = RMSNorm(dimensions: headDim, eps: config.rmsNormEps)
+
+ // MLXFast RoPE: split-half rotation (traditional=false), base from config
+ self.rope = MLXNN.RoPE(dimensions: headDim, traditional: false, base: config.ropeTheta)
+
+ super.init()
+ }
+
+ public func callAsFunction(
+ _ hiddenStates: MLXArray,
+ attentionMask: MLXArray? = nil,
+ cache: (MLXArray, MLXArray)? = nil
+ ) -> (MLXArray, (MLXArray, MLXArray)) {
+ let (batch, seqLen, _) = (hiddenStates.dim(0), hiddenStates.dim(1), hiddenStates.dim(2))
+
+ // Project Q, K, V
+ var queries = qProj(hiddenStates)
+ var keys = kProj(hiddenStates)
+ var values = vProj(hiddenStates)
+
+ // Reshape for multi-head attention
+ queries = queries.reshaped(batch, seqLen, numHeads, headDim)
+ keys = keys.reshaped(batch, seqLen, numKVHeads, headDim)
+ values = values.reshaped(batch, seqLen, numKVHeads, headDim)
+
+ // Apply Q/K normalization
+ queries = qNorm(queries)
+ keys = kNorm(keys)
+
+ // Transpose to [batch, heads, seq, head_dim]
+ queries = queries.transposed(0, 2, 1, 3)
+ keys = keys.transposed(0, 2, 1, 3)
+ values = values.transposed(0, 2, 1, 3)
+
+ // Calculate offset for RoPE based on cache
+ let offset = cache?.0.dim(2) ?? 0
+
+ // Apply MLXFast RoPE (handles split-half rotation via optimized Metal kernel)
+ queries = rope(queries, offset: offset)
+ keys = rope(keys, offset: offset)
+
+ // Update cache
+ var cachedKeys = keys
+ var cachedValues = values
+
+ if let (prevKeys, prevValues) = cache {
+ cachedKeys = concatenated([prevKeys, keys], axis: 2)
+ cachedValues = concatenated([prevValues, values], axis: 2)
+ }
+
+ // SDPA handles GQA natively (N_q != N_kv), no need to tile KV heads
+ let merged = SDPA.attendAndMerge(
+ qHeads: queries, kHeads: cachedKeys, vHeads: cachedValues,
+ scale: scale, mask: attentionMask)
+ let output = oProj(merged)
+
+ return (output, (cachedKeys, cachedValues))
+ }
+}
+
+/// MLP for Qwen3 text decoder (SwiGLU activation, quantized)
+/// Wraps the shared QuantizedMLP for backward compatibility
+public class QuantizedTextMLP: Module {
+ @ModuleInfo var gateProj: QuantizedLinear
+ @ModuleInfo var upProj: QuantizedLinear
+ @ModuleInfo var downProj: QuantizedLinear
+
+ public init(config: TextDecoderConfig) {
+ let hiddenSize = config.hiddenSize
+ let intermediateSize = config.intermediateSize
+
+ self._gateProj.wrappedValue = QuantizedLinear(
+ hiddenSize, intermediateSize, bias: false,
+ groupSize: config.groupSize, bits: config.bits)
+ self._upProj.wrappedValue = QuantizedLinear(
+ hiddenSize, intermediateSize, bias: false,
+ groupSize: config.groupSize, bits: config.bits)
+ self._downProj.wrappedValue = QuantizedLinear(
+ intermediateSize, hiddenSize, bias: false,
+ groupSize: config.groupSize, bits: config.bits)
+
+ super.init()
+ }
+
+ public func callAsFunction(_ x: MLXArray) -> MLXArray {
+ // SwiGLU: down(silu(gate(x)) * up(x))
+ let gate = silu(gateProj(x))
+ let up = upProj(x)
+ return downProj(gate * up)
+ }
+}
+
+/// Decoder layer for Qwen3 text model (quantized)
+public class QuantizedTextDecoderLayer: Module {
+ @ModuleInfo var selfAttn: QuantizedTextAttention
+ @ModuleInfo var mlp: QuantizedTextMLP
+ @ModuleInfo var inputLayerNorm: RMSNorm
+ @ModuleInfo var postAttentionLayerNorm: RMSNorm
+
+ public init(config: TextDecoderConfig) {
+ self._selfAttn.wrappedValue = QuantizedTextAttention(config: config)
+ self._mlp.wrappedValue = QuantizedTextMLP(config: config)
+ self._inputLayerNorm.wrappedValue = RMSNorm(dimensions: config.hiddenSize, eps: config.rmsNormEps)
+ self._postAttentionLayerNorm.wrappedValue = RMSNorm(dimensions: config.hiddenSize, eps: config.rmsNormEps)
+
+ super.init()
+ }
+
+ public func callAsFunction(
+ _ hiddenStates: MLXArray,
+ attentionMask: MLXArray? = nil,
+ cache: (MLXArray, MLXArray)? = nil
+ ) -> (MLXArray, (MLXArray, MLXArray)) {
+ // Self attention with pre-norm
+ let residual = hiddenStates
+ var hidden = inputLayerNorm(hiddenStates)
+ let (attnOutput, newCache) = selfAttn(hidden, attentionMask: attentionMask, cache: cache)
+ hidden = residual + attnOutput
+
+ // MLP with pre-norm
+ let residual2 = hidden
+ hidden = postAttentionLayerNorm(hidden)
+ hidden = mlp(hidden)
+ hidden = residual2 + hidden
+
+ return (hidden, newCache)
+ }
+}
+
+/// Full Qwen3 text decoder model (quantized)
+public class QuantizedTextModel: Module {
+ public let config: TextDecoderConfig
+
+ @ModuleInfo public var embedTokens: PreQuantizedEmbedding
+ @ModuleInfo var layers: [QuantizedTextDecoderLayer]
+ @ModuleInfo var norm: RMSNorm
+
+ public init(config: TextDecoderConfig) {
+ self.config = config
+
+ self._embedTokens.wrappedValue = PreQuantizedEmbedding(
+ embeddingCount: config.vocabSize,
+ dimensions: config.hiddenSize,
+ groupSize: config.groupSize,
+ bits: config.bits)
+ self._layers.wrappedValue = (0.. (MLXArray, [(MLXArray, MLXArray)]) {
+ // Get embeddings
+ var hiddenStates: MLXArray
+ if let embeds = inputsEmbeds {
+ hiddenStates = embeds
+ } else if let ids = inputIds {
+ hiddenStates = embedTokens(ids)
+ } else {
+ fatalError("Either inputIds or inputsEmbeds must be provided")
+ }
+
+ let seqLen = hiddenStates.dim(1)
+
+ // Determine attention mask
+ let mask: MLXArray?
+ if let providedMask = attentionMask {
+ mask = providedMask
+ } else if seqLen == 1 {
+ // Autoregressive: single query can attend to all cached positions, no mask needed
+ mask = nil
+ } else {
+ // Prefill: create causal mask using MLX broadcast operations
+ let cacheLen = cache?.first?.0.dim(2) ?? 0
+ let totalLen = seqLen + cacheLen
+ let rows = (MLXArray(0.. rows, MLXArray(Float(-1e9)), MLXArray(Float(0)))
+ .expandedDimensions(axes: [0, 1])
+ .asType(hiddenStates.dtype)
+ }
+
+ // Apply decoder layers
+ var newCache: [(MLXArray, MLXArray)] = []
+ for (i, layer) in layers.enumerated() {
+ let layerCache = cache?[i]
+ let (output, updatedCache) = layer(hiddenStates, attentionMask: mask, cache: layerCache)
+ hiddenStates = output
+ newCache.append(updatedCache)
+ }
+
+ // Final norm
+ hiddenStates = norm(hiddenStates)
+
+ return (hiddenStates, newCache)
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Qwen3ASR+Memory.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Qwen3ASR+Memory.swift
new file mode 100644
index 0000000..b37cea2
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Qwen3ASR+Memory.swift
@@ -0,0 +1,27 @@
+import AudioCommon
+import MLX
+
+extension Qwen3ASRModel: ModelMemoryManageable {
+ public var isLoaded: Bool { _isLoaded }
+
+ public func unload() {
+ guard _isLoaded else { return }
+ audioEncoder.clearParameters()
+ textDecoder?.clearParameters()
+ // Restore the MLX cache limit if we lowered it at load time. This
+ // un-leaks the cap from PersonaPlex / multi-model processes that
+ // co-load this ASR with a Mimi codec, LLM, or TTS that wants the
+ // full default cache budget.
+ if let prior = savedMLXCacheLimit {
+ MLX.Memory.cacheLimit = prior
+ savedMLXCacheLimit = nil
+ }
+ _isLoaded = false
+ }
+
+ public var memoryFootprint: Int {
+ guard _isLoaded else { return 0 }
+ return audioEncoder.parameterMemoryBytes()
+ + (textDecoder?.parameterMemoryBytes() ?? 0)
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Qwen3ASR+Protocols.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Qwen3ASR+Protocols.swift
new file mode 100644
index 0000000..3a7d545
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Qwen3ASR+Protocols.swift
@@ -0,0 +1,11 @@
+import AudioCommon
+
+// MARK: - SpeechRecognitionModel
+
+extension Qwen3ASRModel: SpeechRecognitionModel {
+ public var inputSampleRate: Int { 16000 }
+
+ public func transcribe(audio: [Float], sampleRate: Int, language: String?) -> String {
+ transcribe(audio: audio, sampleRate: sampleRate, language: language, maxTokens: 448)
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Qwen3ASR.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Qwen3ASR.swift
new file mode 100644
index 0000000..5e00bb9
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Qwen3ASR.swift
@@ -0,0 +1,930 @@
+import Foundation
+import MLXCommon
+import MLX
+import MLXNN
+import MLXFast
+import AudioCommon
+
+/// Optional decoder tunables for `Qwen3ASRModel.transcribe(audio:options:)`.
+///
+/// Defaults match the historical greedy behaviour of `transcribe(audio:)`
+/// so existing callers see zero change. Tune these when greedy decoding
+/// collapses onto a single token (typical on silence or ambiguous phonemes).
+///
+/// The struct also carries the "long-input auto-escalation" knobs used by
+/// the public `transcribe(...)` entry points to bound greedy degeneration
+/// on >15 s audio without affecting short-clip behaviour. See
+/// `adaptedFor(audioDurationSeconds:)`.
+public struct Qwen3DecodingOptions: Sendable {
+ /// Cap on decoder output per chunk.
+ public var maxTokens: Int = 448
+
+ /// Optional language hint ("en", "zh", …). `nil` = auto-detect.
+ public var language: String?
+
+ /// Context hint prepended to the decoder prompt.
+ public var context: String?
+
+ /// HuggingFace-style repetition penalty. Divides the logits of tokens
+ /// already generated this chunk by this factor before `argMax`.
+ /// `1.0` disables; `1.1`–`1.3` is the common tuning range.
+ public var repetitionPenalty: Float = 1.0
+
+ /// If > 0, masks any next-token whose emission would form a repeated
+ /// n-gram of this size. `0` disables.
+ public var noRepeatNgramSize: Int = 0
+
+ /// `0` = greedy (argmax). `> 0` = sample with this temperature via
+ /// Gumbel-max. Higher = more random.
+ public var temperature: Float = 0.0
+
+ /// Adaptive decoding threshold. When the input audio is longer than this
+ /// many seconds AND the caller has left `noRepeatNgramSize` at 0 (the
+ /// default greedy path), the public `transcribe(...)` entry points
+ /// auto-escalate `noRepeatNgramSize` to `longInputNoRepeatNgramSize`
+ /// before forwarding into `generateText`. This bounds the 0.6B
+ /// greedy-decode degeneration observed on long-form audio without
+ /// affecting short clips. Set to `.infinity` to disable entirely.
+ public var longInputThresholdSeconds: Double = 15.0
+
+ /// n-gram size applied by the long-input auto-escalation path. Only
+ /// used when the threshold above triggers AND the caller hasn't
+ /// already set a custom `noRepeatNgramSize`. 3 mirrors the slow-path
+ /// default in `E2EQwen3DecodingOptionsTests`.
+ public var longInputNoRepeatNgramSize: Int = 3
+
+ public init(
+ maxTokens: Int = 448,
+ language: String? = nil,
+ context: String? = nil,
+ repetitionPenalty: Float = 1.0,
+ noRepeatNgramSize: Int = 0,
+ temperature: Float = 0.0,
+ longInputThresholdSeconds: Double = 15.0,
+ longInputNoRepeatNgramSize: Int = 3
+ ) {
+ self.maxTokens = maxTokens
+ self.language = language
+ self.context = context
+ self.repetitionPenalty = repetitionPenalty
+ self.noRepeatNgramSize = noRepeatNgramSize
+ self.temperature = temperature
+ self.longInputThresholdSeconds = longInputThresholdSeconds
+ self.longInputNoRepeatNgramSize = longInputNoRepeatNgramSize
+ }
+
+ /// Length-gated auto-escalation. Returns a copy with
+ /// `noRepeatNgramSize` bumped to `longInputNoRepeatNgramSize` IFF
+ /// 1. `audioDurationSeconds > longInputThresholdSeconds`, AND
+ /// 2. the caller left `noRepeatNgramSize` at 0 (default greedy), AND
+ /// 3. `longInputNoRepeatNgramSize > 0` (escalation not disabled).
+ /// Otherwise returns `self` unchanged.
+ ///
+ /// Any caller that has explicitly tuned `noRepeatNgramSize` — including
+ /// setting it to a non-3 value — is honoured. This preserves the
+ /// fast-path / slow-path routing semantics in `isGreedyFastPath`.
+ func adaptedFor(audioDurationSeconds: Double) -> Qwen3DecodingOptions {
+ guard audioDurationSeconds > longInputThresholdSeconds,
+ noRepeatNgramSize == 0,
+ longInputNoRepeatNgramSize > 0 else {
+ return self
+ }
+ var copy = self
+ copy.noRepeatNgramSize = longInputNoRepeatNgramSize
+ return copy
+ }
+}
+
+/// Special token IDs for Qwen3-ASR
+public struct Qwen3ASRTokens: Sendable {
+ public static let audioTokenId = 151676 // <|audio_pad|>
+ public static let audioStartTokenId = 151669 // <|audio_start|>
+ public static let audioEndTokenId = 151670 // <|audio_end|>
+ public static let eosTokenId = 151645 // <|im_end|>
+ public static let padTokenId = 151643 // <|endoftext|>
+ public static let imStartTokenId = 151644 // <|im_start|>
+ public static let imEndTokenId = 151645 // <|im_end|>
+ public static let timestampTokenId = 151705 // <|timestamp|>
+}
+
+/// Main Qwen3-ASR model for speech recognition.
+///
+/// - Warning: This class is not thread-safe. Create separate instances for concurrent use.
+public class Qwen3ASRModel {
+ /// Default HuggingFace model identifier — the 0.6B 4-bit MLX bundle.
+ /// Mirrors `ASRModelSize.small.modelId`; kept as a top-level constant so
+ /// the AudioServer registry and other call sites have a single SSOT.
+ public static let defaultModelId = "aufklarer/Qwen3-ASR-0.6B-MLX-4bit"
+
+ /// 1.7B 8-bit MLX bundle — higher-capacity sibling of the default.
+ public static let largeModelId = "aufklarer/Qwen3-ASR-1.7B-MLX-8bit"
+
+ /// CoreML-packaged variant for Neural Engine deployment.
+ public static let coreMLModelId = "aufklarer/Qwen3-ASR-CoreML"
+
+ public let audioEncoder: Qwen3AudioEncoder
+ public let featureExtractor: WhisperFeatureExtractor
+ public var textDecoder: QuantizedTextModel?
+
+ /// Tokenizer for decoding output tokens
+ private var tokenizer: Qwen3Tokenizer?
+
+ /// Text decoder config
+ public let textConfig: TextDecoderConfig
+
+ /// Whether the model weights are loaded and ready for inference.
+ var _isLoaded = true
+
+ /// MLX cache limit captured at load time for the .large variant. Stored
+ /// per-instance so `unload()` can restore it — preventing the 4 GB cap
+ /// from leaking into co-loaded models (PersonaPlex loads ASR + LM + TTS
+ /// in the same process). `nil` when no cap was applied (small variant
+ /// or already-capped global state).
+ var savedMLXCacheLimit: Int?
+
+ init(
+ audioConfig: Qwen3AudioEncoderConfig = .default,
+ textConfig: TextDecoderConfig = .small
+ ) {
+ self.audioEncoder = Qwen3AudioEncoder(config: audioConfig)
+ self.featureExtractor = WhisperFeatureExtractor()
+ self.textConfig = textConfig
+ // Text decoder will be initialized when loading weights
+ self.textDecoder = nil
+ }
+
+ /// Set tokenizer for text decoding
+ func setTokenizer(_ tokenizer: Qwen3Tokenizer) {
+ self.tokenizer = tokenizer
+ }
+
+ /// Initialize text decoder (called after loading)
+ func initializeTextDecoder() {
+ self.textDecoder = QuantizedTextModel(config: textConfig)
+ }
+
+ /// Transcribe audio to text with explicit decoder options.
+ ///
+ /// The legacy `transcribe(audio:sampleRate:language:maxTokens:context:)`
+ /// overload below forwards into this path with default (greedy) options.
+ ///
+ /// Long-input adaptive decoding: before forwarding into `generateText`,
+ /// `options.adaptedFor(audioDurationSeconds:)` is applied. On audio
+ /// longer than `options.longInputThresholdSeconds` (default 15 s) AND
+ /// when the caller hasn't customized `noRepeatNgramSize`, the options
+ /// are escalated to engage the no-repeat-n-gram slow path. Short clips
+ /// are unaffected; explicit caller settings are honoured.
+ public func transcribe(
+ audio: [Float],
+ sampleRate: Int = 16000,
+ options: Qwen3DecodingOptions
+ ) -> String {
+ let durationSeconds = sampleRate > 0
+ ? Double(audio.count) / Double(sampleRate)
+ : 0.0
+ let effective = options.adaptedFor(audioDurationSeconds: durationSeconds)
+
+ let melFeatures = featureExtractor.process(audio, sampleRate: sampleRate)
+ let batchedFeatures = melFeatures.expandedDimensions(axis: 0)
+ var audioEmbeds = audioEncoder(batchedFeatures)
+ audioEmbeds = audioEmbeds.expandedDimensions(axis: 0)
+ guard let textDecoder = textDecoder else {
+ let shape = audioEmbeds.shape
+ return "[Audio encoded: \(shape)] - Text decoder not loaded"
+ }
+ return generateText(
+ audioEmbeds: audioEmbeds,
+ textDecoder: textDecoder,
+ language: effective.language,
+ maxTokens: effective.maxTokens,
+ context: effective.context,
+ decodingOptions: effective
+ )
+ }
+
+ /// Transcribe audio to text
+ ///
+ /// Long-input adaptive decoding: the legacy overload constructs a
+ /// default `Qwen3DecodingOptions` and routes through the same
+ /// length-gated escalation as the options-based path. Callers who pin
+ /// `noRepeatNgramSize` via `Qwen3DecodingOptions` directly are out of
+ /// scope here (they reach the other overload).
+ public func transcribe(
+ audio: [Float],
+ sampleRate: Int = 16000,
+ language: String? = nil,
+ maxTokens: Int = 448,
+ context: String? = nil
+ ) -> String {
+ let durationSeconds = sampleRate > 0
+ ? Double(audio.count) / Double(sampleRate)
+ : 0.0
+ let baseOptions = Qwen3DecodingOptions(
+ maxTokens: maxTokens, language: language, context: context)
+ let effective = baseOptions.adaptedFor(audioDurationSeconds: durationSeconds)
+
+ // Extract mel features
+ let melFeatures = featureExtractor.process(audio, sampleRate: sampleRate)
+
+ // Add batch dimension: [mel, time] -> [1, mel, time]
+ let batchedFeatures = melFeatures.expandedDimensions(axis: 0)
+
+ // Encode audio - returns [time, features] without batch dim (matching Python)
+ var audioEmbeds = audioEncoder(batchedFeatures)
+
+ // Add batch dimension for consistency: [time, features] -> [1, time, features]
+ audioEmbeds = audioEmbeds.expandedDimensions(axis: 0)
+
+ // Check if text decoder is loaded
+ guard let textDecoder = textDecoder else {
+ let shape = audioEmbeds.shape
+ return "[Audio encoded: \(shape)] - Text decoder not loaded"
+ }
+
+ // Long-form audio that triggered escalation routes through the
+ // options-aware codepath (which calls `isGreedyFastPath` and falls
+ // out to `generateSlow`); short clips with default greedy take the
+ // legacy fast-path call shape, bit-identical to today.
+ // Mirror `isGreedyFastPath` exactly so the two routes stay in sync
+ // even if a fourth decoder knob is added later.
+ if !Self.isGreedyFastPath(effective) {
+ return generateText(
+ audioEmbeds: audioEmbeds,
+ textDecoder: textDecoder,
+ language: effective.language,
+ maxTokens: effective.maxTokens,
+ context: effective.context,
+ decodingOptions: effective
+ )
+ }
+ return generateText(
+ audioEmbeds: audioEmbeds,
+ textDecoder: textDecoder,
+ language: effective.language,
+ maxTokens: effective.maxTokens,
+ context: effective.context
+ )
+ }
+
+ /// Generate text from audio embeddings.
+ ///
+ /// When `decodingOptions` is supplied, the decoder loop applies an
+ /// HF-style repetition penalty, an optional no-repeat n-gram mask, and
+ /// optional temperature sampling before each token selection. With the
+ /// default `Qwen3DecodingOptions()` (repetition=1.0, no-repeat=0,
+ /// temperature=0) behaviour is bit-identical to plain greedy.
+ func generateText(
+ audioEmbeds: MLXArray,
+ textDecoder: QuantizedTextModel,
+ language: String?,
+ maxTokens: Int,
+ context: String? = nil,
+ decodingOptions: Qwen3DecodingOptions = Qwen3DecodingOptions()
+ ) -> String {
+ // Special token IDs
+ let imStartId = 151644
+ let imEndId = 151645
+ let audioStartId = 151669
+ let audioEndId = 151670
+ let audioPadId = 151676
+ let asrTextId = 151704
+ let newlineId = 198
+
+ // Token IDs for "system", "user", "assistant"
+ let systemId = 8948
+ let userId = 872
+ let assistantId = 77091
+
+ // Number of audio tokens (from audio encoder output)
+ let numAudioTokens = audioEmbeds.dim(1)
+
+ // Build input_ids array with audio_pad placeholder tokens
+ var inputIds: [Int32] = []
+
+ // <|im_start|>system\n{context}<|im_end|>\n
+ inputIds.append(contentsOf: [imStartId, systemId, newlineId].map { Int32($0) })
+ if let context = context, !context.isEmpty, let tokenizer = tokenizer {
+ let contextTokens = tokenizer.encode(context)
+ inputIds.append(contentsOf: contextTokens.map { Int32($0) })
+ }
+ inputIds.append(contentsOf: [imEndId, newlineId].map { Int32($0) })
+
+ // <|im_start|>user\n<|audio_start|>
+ inputIds.append(contentsOf: [imStartId, userId, newlineId, audioStartId].map { Int32($0) })
+
+ // <|audio_pad|> * numAudioTokens (placeholder tokens that will be replaced)
+ let audioStartIndex = inputIds.count
+ for _ in 0..<|im_end|>\n
+ inputIds.append(contentsOf: [audioEndId, imEndId, newlineId].map { Int32($0) })
+
+ // <|im_start|>assistant\n
+ inputIds.append(contentsOf: [imStartId, assistantId, newlineId].map { Int32($0) })
+
+ // Add language hint if specified, then always add <|asr_text|> marker.
+ // Without <|asr_text|>, the model doesn't know it should transcribe.
+ // Without language hint, the model auto-detects and prepends "language XX" to output.
+ if let lang = language, let tokenizer = tokenizer {
+ let langPrefix = "language \(lang)"
+ let langTokens = tokenizer.encode(langPrefix)
+ inputIds.append(contentsOf: langTokens.map { Int32($0) })
+ }
+ inputIds.append(Int32(asrTextId))
+
+ // Get text embeddings for all tokens
+ let inputIdsTensor = MLXArray(inputIds).expandedDimensions(axis: 0)
+ var inputEmbeds = textDecoder.embedTokens(inputIdsTensor)
+
+ // Replace audio_pad token positions with actual audio embeddings
+ let audioEmbedsTyped = audioEmbeds.asType(inputEmbeds.dtype)
+ let beforeAudio = inputEmbeds[0..., 0.." prefix if present (auto-detection output)
+ if let range = rawText.range(of: "") {
+ return String(rawText[range.upperBound...]).trimmingCharacters(in: .whitespaces)
+ }
+ return rawText
+ } else {
+ // Fallback: return token IDs
+ return generatedTokens.map { String($0) }.joined(separator: " ")
+ }
+ }
+
+ /// Greedy with default options: temperature 0, no repetition penalty,
+ /// no n-gram blocking. The double-buffered asyncEval loop only kicks
+ /// in for this configuration so we can guarantee bit-identical token
+ /// sequences vs. the legacy `argMax(...).item()` decoder.
+ static func isGreedyFastPath(_ options: Qwen3DecodingOptions) -> Bool {
+ return options.repetitionPenalty == 1.0
+ && options.noRepeatNgramSize == 0
+ && options.temperature == 0.0
+ }
+
+ /// Double-buffered greedy decode loop. The key trick is to keep the
+ /// "next token" as a lazy 0-D `MLXArray` (the result of `argMax`),
+ /// build the *next* step's forward pass on top of it (still lazy),
+ /// then call `MLX.asyncEval` so the GPU starts computing step N+1
+ /// before we sync step N's int32 to CPU. The host-side EOS check
+ /// and `generatedTokens.append` then overlap with the in-flight GPU
+ /// work for step N+1 instead of stalling between every token.
+ ///
+ /// Greedy correctness invariant: argMax is deterministic, so this
+ /// produces the exact same token sequence as the legacy loop on
+ /// matching inputs.
+ static func generateGreedyAsyncEval(
+ textDecoder: QuantizedTextModel,
+ initialLogits: MLXArray,
+ cache initialCache: [(MLXArray, MLXArray)],
+ maxTokens: Int
+ ) -> [Int32] {
+ var generatedTokens: [Int32] = []
+ guard maxTokens > 0 else { return generatedTokens }
+
+ // Stage 0: argmax of the prefill's last logits. Stays lazy until
+ // the first `.item()` below.
+ //
+ // Cast to int32 explicitly: MLX's `argmax` returns uint32, but the
+ // legacy loop fed `embedTokens` an int32 tensor (built from a Swift
+ // `Int32`). Quantized embedding lookup observably dispatches
+ // differently on uint32 vs. int32, producing tokens that diverge
+ // from the legacy path on a small fraction of inputs. Casting here
+ // restores exact dtype parity, so greedy stays token-for-token
+ // identical to the pre-optimisation decoder.
+ var nextTokenArr = argMax(initialLogits, axis: -1).squeezed().asType(.int32)
+ var cache = initialCache
+ // Kick off the GPU on the first token (and the prefill cache that
+ // step 1's graph will read from).
+ asyncEval(nextTokenArr, cache)
+
+ let eosToken = Int32(Qwen3ASRTokens.eosTokenId)
+
+ for step in 0.. [Int32] {
+ var generatedTokens: [Int32] = []
+ guard maxTokens > 0 else { return generatedTokens }
+ var cache: [(MLXArray, MLXArray)]? = initialCache
+
+ var nextToken = Self.pickNextToken(
+ logits: initialLogits,
+ generatedSoFar: generatedTokens,
+ options: options
+ )
+ generatedTokens.append(nextToken)
+
+ for _ in 1.. Int32 {
+ // Fast path — pure greedy, no modifications.
+ if options.repetitionPenalty == 1.0,
+ options.noRepeatNgramSize == 0,
+ options.temperature == 0 {
+ return argMax(logits, axis: -1).squeezed().item(Int32.self)
+ }
+
+ // Pull logits to CPU. `logits` is [1, 1, vocabSize]; after squeeze
+ // and conversion we have a plain `[Float]` of length vocabSize.
+ let flat = logits.squeezed().asType(.float32)
+ let vocabSize = flat.size
+ var scores: [Float] = flat.asArray(Float.self)
+ precondition(scores.count == vocabSize, "pickNextToken: vocab size mismatch")
+
+ // Repetition penalty: divide logits for already-generated tokens.
+ if options.repetitionPenalty > 1.0 && !generatedSoFar.isEmpty {
+ let penalty = options.repetitionPenalty
+ for token in Set(generatedSoFar) {
+ let idx = Int(token)
+ guard idx >= 0, idx < vocabSize else { continue }
+ let v = scores[idx]
+ // Positive logits divide; negative logits multiply — matches
+ // HuggingFace's implementation so the penalty always reduces
+ // the probability of the repeated token.
+ scores[idx] = v > 0 ? v / penalty : v * penalty
+ }
+ }
+
+ // No-repeat-ngram: any next token whose emission would form a
+ // repeated n-gram of size N gets pushed to -infinity.
+ let n = options.noRepeatNgramSize
+ if n > 0 && generatedSoFar.count >= n - 1 {
+ let lastPrefix = Array(generatedSoFar.suffix(n - 1))
+ // Walk every position where `lastPrefix` already appeared —
+ // the token that followed it becomes forbidden as the NEXT
+ // token now.
+ if generatedSoFar.count >= n {
+ for i in 0...(generatedSoFar.count - n) {
+ let window = Array(generatedSoFar[i..<(i + n - 1)])
+ guard window == lastPrefix else { continue }
+ let forbidden = Int(generatedSoFar[i + n - 1])
+ if forbidden >= 0 && forbidden < vocabSize {
+ scores[forbidden] = -.infinity
+ }
+ }
+ }
+ }
+
+ // Temperature sampling via Gumbel-max trick:
+ // argmax(logits/T + Gumbel(0,1)) ~ categorical(softmax(logits/T)).
+ if options.temperature > 0 {
+ let t = options.temperature
+ for i in 0.. bestScore {
+ bestScore = scores[i]
+ bestIdx = i
+ }
+ return Int32(bestIdx)
+ }
+}
+
+// MARK: - Backward Compatibility (delegates to HuggingFaceDownloader)
+
+public extension Qwen3ASRModel {
+ static func sanitizedCacheKey(for modelId: String) -> String {
+ HuggingFaceDownloader.sanitizedCacheKey(for: modelId)
+ }
+
+ static func validatedRemoteFileName(_ file: String) throws -> String {
+ try HuggingFaceDownloader.validatedRemoteFileName(file)
+ }
+
+ static func validatedLocalPath(directory: URL, fileName: String) throws -> URL {
+ try HuggingFaceDownloader.validatedLocalPath(directory: directory, fileName: fileName)
+ }
+}
+
+// MARK: - Model Size Detection
+
+/// Supported ASR model sizes
+public enum ASRModelSize {
+ case small // 0.6B
+ case large // 1.7B
+
+ /// Default model IDs on HuggingFace
+ public var defaultModelId: String {
+ switch self {
+ case .small: return "aufklarer/Qwen3-ASR-0.6B-MLX-4bit"
+ case .large: return "aufklarer/Qwen3-ASR-1.7B-MLX-8bit"
+ }
+ }
+
+ /// Audio encoder config for this model size
+ public var audioConfig: Qwen3AudioEncoderConfig {
+ switch self {
+ case .small: return .small
+ case .large: return .large
+ }
+ }
+
+ /// Text decoder config for this model size and quantization bits
+ public func textConfig(bits: Int) -> TextDecoderConfig {
+ switch (self, bits) {
+ case (.small, 8): return .small8bit
+ case (.small, _): return .small
+ case (.large, 8): return .large8bit
+ case (.large, _): return .large
+ }
+ }
+
+ /// Text decoder config for this model size (default bits)
+ public var textConfig: TextDecoderConfig {
+ switch self {
+ case .small: return .small
+ case .large: return .large
+ }
+ }
+
+ /// Detect model size from a HuggingFace model ID
+ public static func detect(from modelId: String) -> ASRModelSize {
+ if modelId.contains("1.7B") || modelId.contains("1.7b") {
+ return .large
+ }
+ return .small
+ }
+
+ /// Detect quantization bits from a HuggingFace model ID.
+ /// Returns 4 by default for 0.6B, 8 for 1.7B if not specified.
+ public static func detectBits(from modelId: String) -> Int {
+ let lower = modelId.lowercased()
+ if lower.contains("8bit") || lower.contains("8-bit") {
+ return 8
+ }
+ if lower.contains("4bit") || lower.contains("4-bit") {
+ return 4
+ }
+ // Default: 4 for small, 8 for large (backwards-compatible)
+ let size = detect(from: modelId)
+ return size == .large ? 8 : 4
+ }
+}
+
+// MARK: - Memory guards (Bug 4b/4f support)
+
+internal enum Qwen3ASRMemory {
+ /// Threshold below which the 1.7B variant triggers a load-time RAM
+ /// warning. Total physical memory is the pragmatic signal — observed
+ /// hangs cluster on 8/16 GB Macs with other apps open; 24 GB+ has
+ /// consistently completed inference in our benchmarks. Exposed
+ /// `internal` so unit tests can pin the threshold.
+ static let largeModelRAMWarningThresholdGB: Double = 24.0
+
+ /// MLX cache ceiling applied when loading the 1.7B variant. mlx-swift's
+ /// default tracks `recommendedMaxWorkingSetSize` which on a 16 GB Mac
+ /// can grow to several GB under sustained decoding — pushing residency
+ /// past unified-memory headroom and triggering swap. We bound the
+ /// scratch pool to `min(4 GB, 25% of physical RAM)`: well above the
+ /// per-token decoder working set, but small enough to leave room for
+ /// the OS and other apps.
+ static func cacheLimitForLarge(physicalMemoryBytes: Int) -> Int {
+ let fourGB = 4 * 1024 * 1024 * 1024
+ let quarterRAM = physicalMemoryBytes / 4
+ return max(0, min(fourGB, quarterRAM))
+ }
+
+ /// True when the 1.7B variant should print the soft RAM warning. Total
+ /// (not available) RAM is the pragmatic signal — see threshold doc.
+ static func shouldWarnForLarge(physicalMemoryBytes: UInt64) -> Bool {
+ let physicalGB = Double(physicalMemoryBytes) / 1_073_741_824.0
+ return physicalGB < largeModelRAMWarningThresholdGB
+ }
+
+ /// Emit a human-readable RAM-pressure warning to stderr (NDJSON-IPC safe).
+ /// Naming the alternative model IDs so the user can copy-paste.
+ static func emitLargeRAMWarning(physicalMemoryBytes: UInt64) {
+ let physicalGB = Double(physicalMemoryBytes) / 1_073_741_824.0
+ let msg = """
+ [Qwen3ASR] Warning: loading 1.7B variant on \(String(format: "%.0f", physicalGB)) GB Mac.
+ [Qwen3ASR] The 1.7B model has been observed to swap and stall on <\(Int(largeModelRAMWarningThresholdGB)) GB systems
+ [Qwen3ASR] when other apps are running. If you see a hang, consider:
+ [Qwen3ASR] aufklarer/Qwen3-ASR-0.6B-MLX-8bit (recommended for 8-16 GB)
+ [Qwen3ASR] aufklarer/Qwen3-ASR-1.7B-MLX-4bit (smaller, similar quality)
+ """
+ FileHandle.standardError.write(Data((msg + "\n").utf8))
+ }
+
+ /// Format memory readings (active / cache / peak in bytes) for
+ /// human-readable logging. Centralized so the formatting is consistent
+ /// across load-time + transcribe-time telemetry. Sizes are reported in
+ /// MB. This overload takes `Int` directly so unit tests don't depend on
+ /// `MLX.Memory.Snapshot`'s sealed initializer.
+ static func formatSnapshot(active: Int, cache: Int, peak: Int, label: String) -> String {
+ let mb: (Int) -> String = { String(format: "%.0f MB", Double($0) / 1_048_576.0) }
+ return "[Qwen3ASR][mem] \(label): "
+ + "active=\(mb(active)) "
+ + "cache=\(mb(cache)) "
+ + "peak=\(mb(peak))"
+ }
+
+ /// Production-callsite overload that adapts a live `MLX.Memory.Snapshot`.
+ static func formatSnapshot(_ s: MLX.Memory.Snapshot, label: String) -> String {
+ formatSnapshot(
+ active: s.activeMemory, cache: s.cacheMemory, peak: s.peakMemory,
+ label: label)
+ }
+}
+
+// MARK: - Model Loading
+
+public extension Qwen3ASRModel {
+ /// Load model from HuggingFace hub with automatic weight downloading
+ static func fromPretrained(
+ modelId: String = "aufklarer/Qwen3-ASR-0.6B-MLX-4bit",
+ cacheDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> Qwen3ASRModel {
+ progressHandler?(0.0, "Downloading model...")
+
+ // Auto-detect model size and quantization bits from model ID
+ let modelSize = ASRModelSize.detect(from: modelId)
+ let detectedBits = ASRModelSize.detectBits(from: modelId)
+
+ // Bug 4b: soft RAM warning for the 1.7B variant. Emit BEFORE the
+ // download so users see it on the first byte, not after a 1.7 GB
+ // transfer. Routed to stderr to keep stdout clean for NDJSON-IPC
+ // consumers (speech-studio sidecar).
+ if modelSize == .large {
+ let physical = ProcessInfo.processInfo.physicalMemory
+ if Qwen3ASRMemory.shouldWarnForLarge(physicalMemoryBytes: physical) {
+ Qwen3ASRMemory.emitLargeRAMWarning(physicalMemoryBytes: physical)
+ }
+ }
+
+ // Bug 4f: pre-load memory snapshot for telemetry. Cheap to call;
+ // gives us a baseline to compare against the post-load snapshot.
+ let memBeforeLoad = MLX.Memory.snapshot()
+
+ // Get cache directory
+ let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
+
+ // Download weights and tokenizer files (skips files that already exist on disk)
+ // Download is the slowest part — give it 0-80% of progress
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: modelId,
+ to: cacheDir,
+ additionalFiles: ["vocab.json", "merges.txt", "tokenizer_config.json"],
+ offlineMode: offlineMode,
+ progressHandler: { progress in
+ progressHandler?(progress * 0.8, "Downloading weights...")
+ }
+ )
+
+ progressHandler?(0.80, "Loading tokenizer...")
+
+ // Create model with appropriate config for detected size and bits
+ let model = Qwen3ASRModel(
+ audioConfig: modelSize.audioConfig,
+ textConfig: modelSize.textConfig(bits: detectedBits)
+ )
+
+ // Load tokenizer from vocab.json
+ let vocabPath = cacheDir.appendingPathComponent("vocab.json")
+ if FileManager.default.fileExists(atPath: vocabPath.path) {
+ let tokenizer = Qwen3Tokenizer()
+ try tokenizer.load(from: vocabPath)
+ model.setTokenizer(tokenizer)
+ }
+
+ progressHandler?(0.85, "Loading audio encoder weights...")
+
+ // Load audio encoder weights
+ try WeightLoader.loadWeights(into: model.audioEncoder, from: cacheDir)
+
+ progressHandler?(0.92, "Loading text decoder weights...")
+
+ // Initialize and load text decoder
+ model.initializeTextDecoder()
+ if let textDecoder = model.textDecoder {
+ try WeightLoader.loadTextDecoderWeights(into: textDecoder, from: cacheDir)
+ }
+
+ MetalBudget.pinMemory()
+
+ // Bug 4b: cap MLX scratch pool for the 1.7B variant. Default cache
+ // limit tracks `recommendedMaxWorkingSetSize` which on a 16 GB Mac
+ // can grow to several GB during sustained decoding and trigger
+ // swap. Bounding to `min(4 GB, 25% of physical RAM)` leaves enough
+ // headroom for per-token decoder working set while keeping the
+ // total residency under the OS jetsam threshold. 0.6B path is
+ // unchanged.
+ //
+ // Process-global cap leak fix (adversarial review): we save the
+ // prior limit on the model instance and restore it in `unload()`,
+ // so co-loaded models in the same process (e.g. PersonaPlex
+ // loading ASR + LM + TTS) inherit our cap only for the lifetime
+ // of the loaded ASR. Stacks correctly across multiple ASR
+ // instances: each save captures whatever was active when it
+ // loaded, and each unload pops its own saved value.
+ if modelSize == .large {
+ let physical = Int(ProcessInfo.processInfo.physicalMemory)
+ let newCap = Qwen3ASRMemory.cacheLimitForLarge(physicalMemoryBytes: physical)
+ // Only apply the cap if it would lower the current limit —
+ // never raise a limit a caller has already chosen for itself.
+ let currentLimit = MLX.Memory.cacheLimit
+ if newCap > 0 && newCap < currentLimit {
+ model.savedMLXCacheLimit = currentLimit
+ MLX.Memory.cacheLimit = newCap
+ }
+ }
+
+ // Bug 4f: post-load memory snapshot. Difference vs `memBeforeLoad`
+ // is the model's load-time footprint (weights + activations +
+ // metallib JIT). Useful for tuning the cache cap and for spotting
+ // load-time regressions in PRs.
+ let memAfterLoad = MLX.Memory.snapshot()
+ AudioLog.modelLoading.info("\(Qwen3ASRMemory.formatSnapshot(memBeforeLoad, label: "pre-load"))")
+ AudioLog.modelLoading.info("\(Qwen3ASRMemory.formatSnapshot(memAfterLoad, label: "post-load"))")
+ // Display max(0, delta): MLX can free cached weights between
+ // snapshots, which makes "active" go down; clamp at 0 so the
+ // load-delta label stays meaningful in logs.
+ let loadActiveDelta = max(0, memAfterLoad.activeMemory - memBeforeLoad.activeMemory)
+ AudioLog.modelLoading.info(
+ "[Qwen3ASR][mem] load delta (active): \(String(format: "%.0f MB", Double(loadActiveDelta) / 1_048_576.0))")
+
+ progressHandler?(1.0, "Ready")
+
+ return model
+ }
+
+ /// Download tokenizer + weight files only — does not load MLX/Metal.
+ /// Use from Settings manual download; inference still calls `fromPretrained()`.
+ static func downloadWeightsOnly(
+ modelId: String = defaultModelId,
+ cacheDir: URL? = nil,
+ registry: ModelRegistry = .huggingFace(),
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws {
+ progressHandler?(0.0, "Downloading model...")
+ let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
+ switch registry {
+ case .huggingFace(let hubEndpoint):
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: modelId,
+ to: cacheDir,
+ additionalFiles: ["vocab.json", "merges.txt", "tokenizer_config.json"],
+ hubEndpoint: hubEndpoint,
+ progressHandler: { progress in
+ progressHandler?(progress, "Downloading weights...")
+ }
+ )
+ case .modelScope(let baseURL, let revision):
+ try await ModelScopeDownloader.downloadWeights(
+ modelId: modelId,
+ to: cacheDir,
+ additionalFiles: ["vocab.json", "merges.txt", "tokenizer_config.json"],
+ baseURL: baseURL,
+ revision: revision,
+ progressHandler: { progress in
+ progressHandler?(progress, "Downloading weights...")
+ }
+ )
+ }
+ progressHandler?(1.0, "Downloaded")
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/StreamingASR.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/StreamingASR.swift
new file mode 100644
index 0000000..e87ade6
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/StreamingASR.swift
@@ -0,0 +1,277 @@
+import Foundation
+import AudioCommon
+import SpeechVAD
+
+// MARK: - TranscriptionSegment
+
+public struct TranscriptionSegment: Sendable {
+ public let text: String
+ public let startTime: Float
+ public let endTime: Float
+ public let isFinal: Bool
+ public let segmentIndex: Int
+
+ public init(text: String, startTime: Float, endTime: Float, isFinal: Bool, segmentIndex: Int) {
+ self.text = text
+ self.startTime = startTime
+ self.endTime = endTime
+ self.isFinal = isFinal
+ self.segmentIndex = segmentIndex
+ }
+}
+
+// MARK: - StreamingASRConfig
+
+public struct StreamingASRConfig: Sendable {
+ public var maxSegmentDuration: Float
+ public var vadConfig: VADConfig
+ public var language: String?
+ public var maxTokens: Int
+ public var emitPartialResults: Bool
+ public var partialResultInterval: Float
+ public var context: String?
+
+ public init(
+ maxSegmentDuration: Float = 10.0,
+ vadConfig: VADConfig = .sileroDefault,
+ language: String? = nil,
+ maxTokens: Int = 448,
+ emitPartialResults: Bool = false,
+ partialResultInterval: Float = 1.0,
+ context: String? = nil
+ ) {
+ self.maxSegmentDuration = maxSegmentDuration
+ self.vadConfig = vadConfig
+ self.language = language
+ self.maxTokens = maxTokens
+ self.emitPartialResults = emitPartialResults
+ self.partialResultInterval = partialResultInterval
+ self.context = context
+ }
+
+ public static let `default` = StreamingASRConfig()
+}
+
+// MARK: - StreamingASR
+
+/// Streaming ASR with VAD-guided segmentation.
+///
+/// - Warning: This class is not thread-safe. Create separate instances for concurrent use.
+public class StreamingASR {
+ private let asrModel: Qwen3ASRModel
+ private let vadModel: SileroVADModel
+
+ public init(asrModel: Qwen3ASRModel, vadModel: SileroVADModel) {
+ self.asrModel = asrModel
+ self.vadModel = vadModel
+ }
+
+ public static func fromPretrained(
+ asrModelId: String = "aufklarer/Qwen3-ASR-0.6B-MLX-4bit",
+ vadModelId: String = SileroVADModel.defaultModelId,
+ cacheDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> StreamingASR {
+ let asr = try await Qwen3ASRModel.fromPretrained(
+ modelId: asrModelId, cacheDir: cacheDir, offlineMode: offlineMode, progressHandler: progressHandler)
+ let vad = try await SileroVADModel.fromPretrained(
+ modelId: vadModelId, cacheDir: cacheDir, offlineMode: offlineMode, progressHandler: progressHandler)
+ return StreamingASR(asrModel: asr, vadModel: vad)
+ }
+
+ /// Streaming transcription — emits TranscriptionSegments as speech is detected.
+ public func transcribeStream(
+ audio: [Float],
+ sampleRate: Int = 16000,
+ config: StreamingASRConfig = .default
+ ) -> AsyncThrowingStream {
+ AsyncThrowingStream { continuation in
+ let samples: [Float]
+ if sampleRate != 16000 {
+ samples = AudioFileLoader.resample(audio, from: sampleRate, to: 16000)
+ } else {
+ samples = audio
+ }
+
+ let processor = StreamingVADProcessor(model: vadModel, config: config.vadConfig)
+ let chunkSize = SileroVADModel.chunkSize
+ var segmentIndex = 0
+ var speechStartSample: Int?
+
+ // Phase 2 state
+ var lastPartialTime: Float = 0
+
+ var offset = 0
+ while offset < samples.count {
+ let end = min(offset + chunkSize, samples.count)
+ let chunk = Array(samples[offset..= config.partialResultInterval {
+ let endSample = min(Int(currentTime * 16000), samples.count)
+ guard startSample < endSample else {
+ lastPartialTime = currentTime
+ continue
+ }
+ let segmentAudio = Array(samples[startSample..= config.maxSegmentDuration {
+ let endSample = min(Int(currentTime * 16000), samples.count)
+ guard startSample < endSample else {
+ speechStartSample = Int(currentTime * 16000)
+ lastPartialTime = currentTime
+ continue
+ }
+ let segmentAudio = Array(samples[startSample..= config.maxSegmentDuration {
+ let endSample = min(Int(currentTime * 16000), samples.count)
+ guard startSample < endSample else {
+ speechStartSample = Int(currentTime * 16000)
+ continue
+ }
+ let segmentAudio = Array(samples[startSample.. [String] {
+ var result: [String] = []
+ for i in 0..` pairs so the model
+ /// can predict start/end timestamps at those positions.
+ public static func prepareForAlignment(
+ text: String,
+ tokenizer: Qwen3Tokenizer,
+ language: String = "English"
+ ) -> SlottedText {
+ let pairs = splitIntoWordPairs(text, language: language)
+ let tsId = Qwen3ASRTokens.timestampTokenId
+
+ var tokenIds: [Int] = []
+ var timestampPositions: [Int] = []
+ var validWords: [String] = []
+
+ for pair in pairs {
+ let wordTokens = tokenizer.encode(pair.cleaned)
+ guard !wordTokens.isEmpty else {
+ // Cleaned form unencodable: attach surface to previous word
+ // so we don't drop punctuation that anchored to it.
+ if !validWords.isEmpty {
+ validWords[validWords.count - 1] += pair.surface
+ }
+ continue
+ }
+
+ timestampPositions.append(tokenIds.count)
+ tokenIds.append(tsId)
+
+ tokenIds.append(contentsOf: wordTokens)
+
+ timestampPositions.append(tokenIds.count)
+ tokenIds.append(tsId)
+
+ validWords.append(pair.surface)
+ }
+
+ return SlottedText(
+ tokenIds: tokenIds,
+ timestampPositions: timestampPositions,
+ words: validWords
+ )
+ }
+
+ /// Split text into words using the language-appropriate strategy.
+ /// Returned strings are the cleaned (punctuation-stripped) forms,
+ /// intended for callers that don't care about surface preservation.
+ static func splitIntoWords(_ text: String, language: String) -> [String] {
+ return splitIntoWordPairs(text, language: language).map { $0.cleaned }
+ }
+
+ /// Split text into (surface, cleaned) pairs. Surface keeps adjacent
+ /// punctuation; cleaned is what the model tokenizer sees.
+ static func splitIntoWordPairs(_ text: String, language: String) -> [WordPair] {
+ let lang = language.lowercased()
+
+ if lang.contains("japanese") || lang == "ja" {
+ return nlTokenizePairs(text, language: .japanese)
+ }
+ if lang.contains("korean") || lang == "ko" {
+ return nlTokenizePairs(text, language: .korean)
+ }
+ // Scripts without word-level whitespace where Apple's NLTokenizer
+ // provides native segmentation. Without these dispatches the
+ // default whitespace path collapses each sentence to one token.
+ if let nlLang = nlLanguageForUnspaced(lang) {
+ return nlTokenizePairs(text, language: nlLang)
+ }
+ return tokenizeSpaceLangPairs(text)
+ }
+
+ private static func nlLanguageForUnspaced(_ lang: String) -> NLLanguage? {
+ if lang.contains("thai") || lang == "th" { return .thai }
+ if lang.contains("lao") || lang == "lo" { return .lao }
+ if lang.contains("khmer") || lang == "km" { return .khmer }
+ if lang.contains("burmese") || lang.contains("myanmar") || lang == "my" { return .burmese }
+ if lang.contains("tibetan") || lang == "bo" { return .tibetan }
+ return nil
+ }
+
+ // MARK: - Japanese / Korean / unspaced scripts
+
+ /// Apple's `NLTokenizer` reports word ranges (without surrounding
+ /// punctuation). We attach trailing non-letter, non-whitespace
+ /// characters between consecutive word ranges to the preceding word's
+ /// surface so commas, full-width periods, etc. ride along.
+ private static func nlTokenizePairs(_ text: String, language: NLLanguage) -> [WordPair] {
+ let tokenizer = NLTokenizer(unit: .word)
+ tokenizer.setLanguage(language)
+ tokenizer.string = text
+
+ var ranges: [Range] = []
+ tokenizer.enumerateTokens(in: text.startIndex.. [WordPair] {
+ var pairs: [WordPair] = []
+ for raw in text.split(whereSeparator: \.isWhitespace) {
+ let segment = String(raw)
+ let segPairs = pairsForSegment(segment)
+ // Reattach a leading whitespace separator to the first new pair
+ // so the surface reads naturally when concatenated. We don't
+ // actually reinsert spaces — callers can join with spaces — but
+ // we do rejoin punctuation that ended up segment-leading with
+ // no anchor word (rare, e.g. a stray "—") to the previous word.
+ if segPairs.isEmpty {
+ if !pairs.isEmpty {
+ pairs[pairs.count - 1].surface += segment
+ }
+ continue
+ }
+ pairs.append(contentsOf: segPairs)
+ }
+ return pairs
+ }
+
+ /// Convert one whitespace-bounded segment to (surface, cleaned) pairs.
+ /// For non-Han segments: a single pair with surface = segment, cleaned
+ /// = letters/numbers only. For segments containing Han ideographs:
+ /// each Han is its own pair; consecutive non-Han runs become their own
+ /// pair if they contain any letter/number, or attach to the
+ /// neighbouring pair's surface if they are pure punctuation/symbols.
+ private static func pairsForSegment(_ seg: String) -> [WordPair] {
+ let hasHan = seg.unicodeScalars.contains(where: isHanIdeograph)
+ if !hasHan {
+ let cleaned = cleanToken(seg)
+ if cleaned.isEmpty { return [] }
+ return [WordPair(surface: seg, cleaned: cleaned)]
+ }
+ var pairs: [WordPair] = []
+ var nonHanBuf = ""
+
+ func flushNonHan(beforeHan: Bool) {
+ guard !nonHanBuf.isEmpty else { return }
+ let cleaned = cleanToken(nonHanBuf)
+ if cleaned.isEmpty {
+ // Pure punctuation: attach to the previous pair's trailing
+ // surface. If we're at the start with no previous pair,
+ // leave the buffer for the upcoming Han to absorb.
+ if !pairs.isEmpty {
+ pairs[pairs.count - 1].surface += nonHanBuf
+ nonHanBuf = ""
+ } else if !beforeHan {
+ // Trailing pure-punct with no anchor at all — drop.
+ nonHanBuf = ""
+ }
+ return
+ }
+ pairs.append(WordPair(surface: nonHanBuf, cleaned: cleaned))
+ nonHanBuf = ""
+ }
+
+ for scalar in seg.unicodeScalars {
+ if isHanIdeograph(scalar) {
+ flushNonHan(beforeHan: true)
+ let han = String(scalar)
+ if !nonHanBuf.isEmpty {
+ // Leading pure-punct waiting for a Han anchor.
+ pairs.append(WordPair(surface: nonHanBuf + han, cleaned: han))
+ nonHanBuf = ""
+ } else {
+ pairs.append(WordPair(surface: han, cleaned: han))
+ }
+ } else {
+ nonHanBuf.append(Character(scalar))
+ }
+ }
+ flushNonHan(beforeHan: false)
+ return pairs
+ }
+
+ // MARK: - Legacy entry points (kept for tests / external callers)
+
+ static func tokenizeJapanese(_ text: String) -> [String] {
+ return nlTokenizePairs(text, language: .japanese).map { $0.cleaned }
+ }
+
+ static func tokenizeKorean(_ text: String) -> [String] {
+ return nlTokenizePairs(text, language: .korean).map { $0.cleaned }
+ }
+
+ static func tokenizeSpaceLang(_ text: String) -> [String] {
+ return tokenizeSpaceLangPairs(text).map { $0.cleaned }
+ }
+
+ // MARK: - Cleaning + classification
+
+ /// Keep only Unicode Letters (`L*`), Numbers (`N*`), and ASCII apostrophe.
+ /// Strips punctuation (e.g. the full-width period `。`), symbols,
+ /// separators, and marks.
+ static func cleanToken(_ token: String) -> String {
+ var out = ""
+ for scalar in token.unicodeScalars {
+ if isKeptScalar(scalar) {
+ out.unicodeScalars.append(scalar)
+ }
+ }
+ return out
+ }
+
+ private static func isKeptScalar(_ scalar: Unicode.Scalar) -> Bool {
+ if scalar == "'" { return true }
+ let cat = scalar.properties.generalCategory
+ switch cat {
+ case .uppercaseLetter, .lowercaseLetter, .titlecaseLetter,
+ .modifierLetter, .otherLetter,
+ .decimalNumber, .letterNumber, .otherNumber,
+ // Combining marks are essential for scripts like Thai, Lao,
+ // Khmer, Burmese, Tibetan, Devanagari, Bengali, Arabic harakat
+ // — stripping them mangles the word (e.g. "สวัสดี" → "สวสด").
+ .nonspacingMark, .spacingMark, .enclosingMark:
+ return true
+ default:
+ return false
+ }
+ }
+
+ /// Han ideograph ranges only. Notably **excludes** hiragana
+ /// (0x3040–0x309F), katakana (0x30A0–0x30FF), and Hangul syllables
+ /// (0xAC00–0xD7AF), which are handled by the language-specific
+ /// tokenizers (Japanese / Korean) instead.
+ static func isHanIdeograph(_ scalar: Unicode.Scalar) -> Bool {
+ let v = scalar.value
+ if v >= 0x4E00 && v <= 0x9FFF { return true } // CJK Unified
+ if v >= 0x3400 && v <= 0x4DBF { return true } // Extension A
+ if v >= 0x20000 && v <= 0x2A6DF { return true } // Extension B
+ if v >= 0x2A700 && v <= 0x2B73F { return true } // Extension C
+ if v >= 0x2B740 && v <= 0x2B81F { return true } // Extension D
+ if v >= 0x2B820 && v <= 0x2CEAF { return true } // Extension E
+ if v >= 0xF900 && v <= 0xFAFF { return true } // Compatibility
+ return false
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/TimestampCorrection.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/TimestampCorrection.swift
new file mode 100644
index 0000000..1148f4b
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/TimestampCorrection.swift
@@ -0,0 +1,145 @@
+import Foundation
+
+/// Monotonicity correction for forced alignment timestamps using LIS
+public enum TimestampCorrection {
+
+ /// Enforce monotonically increasing timestamps via LIS + interpolation.
+ ///
+ /// 1. Find Longest Increasing Subsequence of raw timestamp indices (O(n log n))
+ /// 2. For positions not in LIS:
+ /// - Small gaps (<=2): nearest-neighbor correction
+ /// - Larger gaps: linear interpolation between LIS anchors
+ ///
+ /// - Parameter rawIndices: Raw timestamp class indices from argmax
+ /// - Returns: Corrected monotonically increasing indices
+ public static func enforceMonotonicity(_ rawIndices: [Int]) -> [Int] {
+ guard rawIndices.count > 1 else { return rawIndices }
+
+ // Find LIS positions
+ let lisPositions = longestIncreasingSubsequencePositions(rawIndices)
+ let lisSet = Set(lisPositions)
+
+ // Build anchor points: (position_in_array, value)
+ var anchors: [(pos: Int, val: Int)] = []
+ for pos in lisPositions {
+ anchors.append((pos, rawIndices[pos]))
+ }
+
+ // If LIS covers everything, already monotonic
+ if anchors.count == rawIndices.count {
+ return rawIndices
+ }
+
+ var corrected = rawIndices
+
+ // Fill gaps between anchors
+ var anchorIdx = 0
+ var i = 0
+ while i < corrected.count {
+ if lisSet.contains(i) {
+ // This position is an anchor, keep it
+ anchorIdx = anchors.firstIndex(where: { $0.pos == i }) ?? anchorIdx
+ i += 1
+ continue
+ }
+
+ // Find surrounding anchors
+ let prevAnchor: (pos: Int, val: Int)?
+ let nextAnchor: (pos: Int, val: Int)?
+
+ if anchorIdx < anchors.count && anchors[anchorIdx].pos < i {
+ prevAnchor = anchors[anchorIdx]
+ } else if anchorIdx > 0 {
+ prevAnchor = anchors[anchorIdx - 1]
+ } else {
+ prevAnchor = nil
+ }
+
+ // Find next anchor after position i
+ var nextIdx = anchorIdx
+ while nextIdx < anchors.count && anchors[nextIdx].pos <= i {
+ nextIdx += 1
+ }
+ nextAnchor = nextIdx < anchors.count ? anchors[nextIdx] : nil
+
+ // Interpolate
+ if let prev = prevAnchor, let next = nextAnchor {
+ let gapSize = next.pos - prev.pos
+ if gapSize <= 3 {
+ // Small gap: nearest neighbor
+ let distToPrev = i - prev.pos
+ let distToNext = next.pos - i
+ corrected[i] = distToPrev <= distToNext ? prev.val : next.val
+ } else {
+ // Linear interpolation
+ let t = Float(i - prev.pos) / Float(next.pos - prev.pos)
+ corrected[i] = prev.val + Int(t * Float(next.val - prev.val))
+ }
+ } else if let prev = prevAnchor {
+ // After last anchor: clamp to last anchor value
+ corrected[i] = prev.val
+ } else if let next = nextAnchor {
+ // Before first anchor: clamp to first anchor value
+ corrected[i] = next.val
+ }
+
+ i += 1
+ }
+
+ // Final pass: ensure strict monotonicity
+ for i in 1.. [Int] {
+ guard !arr.isEmpty else { return [] }
+
+ let n = arr.count
+ // tails[i] = smallest tail element for increasing subsequence of length i+1
+ var tails: [Int] = []
+ // tailIndices[i] = index in arr where tails[i] comes from
+ var tailIndices: [Int] = []
+ // parent[i] = index of previous element in LIS ending at arr[i]
+ var parent = [Int](repeating: -1, count: n)
+
+ for i in 0.. 0 ? tailIndices[lo - 1] : -1
+ }
+
+ // Reconstruct LIS positions
+ var positions: [Int] = []
+ var idx = tailIndices[tails.count - 1]
+ while idx != -1 {
+ positions.append(idx)
+ idx = parent[idx]
+ }
+
+ positions.reverse()
+ return positions
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/WeightLoading.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/WeightLoading.swift
new file mode 100644
index 0000000..a6732c9
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/WeightLoading.swift
@@ -0,0 +1,321 @@
+import Foundation
+import MLXCommon
+import MLX
+import MLXNN
+import AudioCommon
+
+/// Weight loading utilities for Qwen3-ASR
+/// Uses direct HuggingFace key paths — model structure must match exactly.
+///
+/// All loaders stream weights per safetensors shard rather than accumulating
+/// every tensor from every file into a single `[String: MLXArray]` dict before
+/// applying. This keeps transient load-time peak memory at roughly
+/// `model_size + one_shard` instead of `model_size + checkpoint_size` (the
+/// shards plus the assembled dict were duplicating the entire model in RAM
+/// during load — observed ~1.5–2.0 GB peak on 1.7B). `Module.update(parameters:)`
+/// is documented as partial-safe (mlx-swift `Module.swift:423`: "any omitted
+/// values will be unchanged"), so applying every component against every shard
+/// is correct even when a single layer's tensors are split across files.
+public enum WeightLoader {
+
+ /// Load weights from safetensors file
+ public static func loadSafetensors(url: URL) throws -> [String: MLXArray] {
+ try CommonWeightLoader.loadSafetensors(url: url)
+ }
+
+ /// Load and apply weights to model using HuggingFace key paths directly.
+ /// Streams per shard — see file docstring.
+ public static func loadWeights(
+ into audioEncoder: Qwen3AudioEncoder,
+ from directory: URL
+ ) throws {
+ let files = try safetensorFiles(in: directory)
+ print("Found \(files.count) safetensor files")
+
+ var appliedTotal = 0
+ for file in files {
+ print("Loading: \(file.lastPathComponent)")
+ let raw = try loadSafetensors(url: file)
+ let audioTowerWeights = stripPrefix(raw, prefix: "audio_tower.")
+ if audioTowerWeights.isEmpty { continue }
+ applyAudioEncoderComponents(
+ to: audioEncoder, weights: audioTowerWeights,
+ transposeConv2dPyTorch: false)
+ appliedTotal += audioTowerWeights.count
+ // `raw` and `audioTowerWeights` go out of scope here; their
+ // MLXArray references release once each call to
+ // `update(parameters:)` above has adopted the tensors the
+ // model actually needed.
+ }
+ print("Applied weights to audio encoder (\(audioEncoder.layers.count) layers, \(appliedTotal) tensors)")
+ }
+
+ /// Load and apply weights to quantized text decoder. Per-shard streaming.
+ public static func loadTextDecoderWeights(
+ into textModel: QuantizedTextModel,
+ from directory: URL
+ ) throws {
+ let files = try safetensorFiles(in: directory)
+ var appliedTotal = 0
+ for file in files {
+ let raw = try loadSafetensors(url: file)
+ let textWeights = stripPrefix(raw, prefix: "model.")
+ if textWeights.isEmpty { continue }
+ applyQuantizedTextDecoderComponents(to: textModel, weights: textWeights)
+ appliedTotal += textWeights.count
+ }
+ print("Applied weights to text decoder (\(textModel.layers.count) layers, \(appliedTotal) tensors)")
+ }
+
+ // MARK: - Forced Aligner Weight Loading
+
+ /// Load weights for the forced aligner model. Per-shard streaming.
+ ///
+ /// Weight key structure (under optional `thinker.` prefix):
+ /// - `audio_tower.*` → audio encoder
+ /// - `model.*` → text decoder (quantized or float)
+ /// - `lm_head.*` → classify head (Linear, NOT quantized)
+ public static func loadForcedAlignerWeights(
+ into model: Qwen3ForcedAligner,
+ from directory: URL
+ ) throws {
+ let files = try safetensorFiles(in: directory)
+
+ var audioApplied = 0
+ var textApplied = 0
+ var headApplied = 0
+
+ for file in files {
+ print("Loading: \(file.lastPathComponent)")
+ let raw = try loadSafetensors(url: file)
+
+ // Strip `thinker.` if present so downstream prefix-strip logic
+ // sees a uniform key space.
+ let normalized = stripPrefix(raw, prefix: "thinker.", keepUnprefixed: true)
+
+ let audioTowerWeights = stripPrefix(normalized, prefix: "audio_tower.")
+ if !audioTowerWeights.isEmpty {
+ applyAudioEncoderComponents(
+ to: model.audioEncoder, weights: audioTowerWeights,
+ transposeConv2dPyTorch: true)
+ audioApplied += audioTowerWeights.count
+ }
+
+ let textWeights = stripPrefix(normalized, prefix: "model.")
+ if !textWeights.isEmpty {
+ if let quantized = model.textDecoder as? QuantizedTextModel {
+ applyQuantizedTextDecoderComponents(to: quantized, weights: textWeights)
+ } else if let floatModel = model.textDecoder as? FloatTextModel {
+ applyFloatTextDecoderComponents(to: floatModel, weights: textWeights)
+ }
+ textApplied += textWeights.count
+ }
+
+ let classifyWeights = filterPrefix(normalized, keepingPrefix: "lm_head.")
+ if !classifyWeights.isEmpty {
+ CommonWeightLoader.applyLinearWeights(
+ to: model.classifyHead, prefix: "lm_head", from: classifyWeights)
+ headApplied += classifyWeights.count
+ }
+ }
+
+ print("Audio tower: \(audioApplied), Text decoder: \(textApplied), Classify head: \(headApplied)")
+ print("Applied audio encoder weights (\(model.audioEncoder.layers.count) layers)")
+ if let quantized = model.textDecoder as? QuantizedTextModel {
+ print("Applied quantized text decoder weights (\(quantized.layers.count) layers)")
+ } else if let floatModel = model.textDecoder as? FloatTextModel {
+ print("Applied float text decoder weights (\(floatModel.layers.count) layers)")
+ }
+ print("Applied classify head weights")
+ }
+
+ // MARK: - Shard discovery + prefix filter helpers
+
+ private static func safetensorFiles(in directory: URL) throws -> [URL] {
+ let fileManager = FileManager.default
+ let contents = try fileManager.contentsOfDirectory(at: directory, includingPropertiesForKeys: nil)
+ let files = contents.filter { $0.pathExtension == "safetensors" }
+ guard !files.isEmpty else {
+ throw WeightLoadingError.noWeightsFound(directory)
+ }
+ // Sort lexicographically so the load order is deterministic across
+ // runs — purely cosmetic for the logs, but it makes regression
+ // diffs against captured stderr stable.
+ return files.sorted { $0.lastPathComponent < $1.lastPathComponent }
+ }
+
+ /// Return a new dict containing only keys with `prefix`, with the prefix
+ /// stripped off. `keepUnprefixed: true` retains keys that DON'T have the
+ /// prefix (used by the forced aligner where `thinker.` is optional).
+ private static func stripPrefix(
+ _ weights: [String: MLXArray],
+ prefix: String,
+ keepUnprefixed: Bool = false
+ ) -> [String: MLXArray] {
+ var out: [String: MLXArray] = [:]
+ out.reserveCapacity(weights.count)
+ for (key, value) in weights {
+ if key.hasPrefix(prefix) {
+ out[String(key.dropFirst(prefix.count))] = value
+ } else if keepUnprefixed {
+ out[key] = value
+ }
+ }
+ return out
+ }
+
+ /// Return a new dict containing only keys with `keepingPrefix`,
+ /// preserving the prefix on the kept keys.
+ private static func filterPrefix(
+ _ weights: [String: MLXArray],
+ keepingPrefix: String
+ ) -> [String: MLXArray] {
+ var out: [String: MLXArray] = [:]
+ for (key, value) in weights where key.hasPrefix(keepingPrefix) {
+ out[key] = value
+ }
+ return out
+ }
+
+ // MARK: - Per-shard component application
+
+ /// Apply every audio-encoder component slot against `weights`. Components
+ /// whose tensors aren't present in this shard are no-ops (each
+ /// `apply…Weights` helper guards its `weights[key]` lookups).
+ private static func applyAudioEncoderComponents(
+ to audioEncoder: Qwen3AudioEncoder,
+ weights: [String: MLXArray],
+ transposeConv2dPyTorch: Bool
+ ) {
+ applyConv2dWeights(to: audioEncoder.conv2d1, prefix: "conv2d1", from: weights, transposePyTorch: transposeConv2dPyTorch)
+ applyConv2dWeights(to: audioEncoder.conv2d2, prefix: "conv2d2", from: weights, transposePyTorch: transposeConv2dPyTorch)
+ applyConv2dWeights(to: audioEncoder.conv2d3, prefix: "conv2d3", from: weights, transposePyTorch: transposeConv2dPyTorch)
+ CommonWeightLoader.applyLinearWeights(to: audioEncoder.convOut, prefix: "conv_out", from: weights)
+ CommonWeightLoader.applyLayerNormWeights(to: audioEncoder.lnPost, prefix: "ln_post", from: weights)
+ CommonWeightLoader.applyLinearWeights(to: audioEncoder.proj1, prefix: "proj1", from: weights)
+ CommonWeightLoader.applyLinearWeights(to: audioEncoder.proj2, prefix: "proj2", from: weights)
+ for (index, layer) in audioEncoder.layers.enumerated() {
+ applyEncoderLayerWeights(to: layer, prefix: "layers.\(index)", from: weights)
+ }
+ }
+
+ private static func applyQuantizedTextDecoderComponents(
+ to textModel: QuantizedTextModel,
+ weights: [String: MLXArray]
+ ) {
+ CommonWeightLoader.applyQuantizedEmbeddingWeights(
+ to: textModel.embedTokens, prefix: "embed_tokens", from: weights)
+ CommonWeightLoader.applyRMSNormWeights(
+ to: textModel.norm, prefix: "norm", from: weights)
+ for (index, layer) in textModel.layers.enumerated() {
+ applyQuantizedDecoderLayerWeights(
+ to: layer, prefix: "layers.\(index)", from: weights)
+ }
+ }
+
+ private static func applyFloatTextDecoderComponents(
+ to textModel: FloatTextModel,
+ weights: [String: MLXArray]
+ ) {
+ CommonWeightLoader.applyEmbeddingWeights(
+ to: textModel.embedTokens, prefix: "embed_tokens", from: weights)
+ CommonWeightLoader.applyRMSNormWeights(
+ to: textModel.norm, prefix: "norm", from: weights)
+ for (index, layer) in textModel.layers.enumerated() {
+ applyFloatDecoderLayerWeights(
+ to: layer, prefix: "layers.\(index)", from: weights)
+ }
+ }
+
+ // MARK: - ASR-specific Weight Application Helpers
+
+ private static func applyQuantizedDecoderLayerWeights(
+ to layer: QuantizedTextDecoderLayer,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ // Self attention
+ CommonWeightLoader.applyQuantizedLinearWeights(to: layer.selfAttn.qProj, prefix: "\(prefix).self_attn.q_proj", from: weights)
+ CommonWeightLoader.applyQuantizedLinearWeights(to: layer.selfAttn.kProj, prefix: "\(prefix).self_attn.k_proj", from: weights)
+ CommonWeightLoader.applyQuantizedLinearWeights(to: layer.selfAttn.vProj, prefix: "\(prefix).self_attn.v_proj", from: weights)
+ CommonWeightLoader.applyQuantizedLinearWeights(to: layer.selfAttn.oProj, prefix: "\(prefix).self_attn.o_proj", from: weights)
+
+ // Q/K norms
+ CommonWeightLoader.applyRMSNormWeights(to: layer.selfAttn.qNorm, prefix: "\(prefix).self_attn.q_norm", from: weights)
+ CommonWeightLoader.applyRMSNormWeights(to: layer.selfAttn.kNorm, prefix: "\(prefix).self_attn.k_norm", from: weights)
+
+ // Layer norms
+ CommonWeightLoader.applyRMSNormWeights(to: layer.inputLayerNorm, prefix: "\(prefix).input_layernorm", from: weights)
+ CommonWeightLoader.applyRMSNormWeights(to: layer.postAttentionLayerNorm, prefix: "\(prefix).post_attention_layernorm", from: weights)
+
+ // MLP
+ CommonWeightLoader.applyQuantizedLinearWeights(to: layer.mlp.gateProj, prefix: "\(prefix).mlp.gate_proj", from: weights)
+ CommonWeightLoader.applyQuantizedLinearWeights(to: layer.mlp.upProj, prefix: "\(prefix).mlp.up_proj", from: weights)
+ CommonWeightLoader.applyQuantizedLinearWeights(to: layer.mlp.downProj, prefix: "\(prefix).mlp.down_proj", from: weights)
+ }
+
+ private static func applyFloatDecoderLayerWeights(
+ to layer: FloatTextDecoderLayer,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ CommonWeightLoader.applyLinearWeights(to: layer.selfAttn.qProj, prefix: "\(prefix).self_attn.q_proj", from: weights)
+ CommonWeightLoader.applyLinearWeights(to: layer.selfAttn.kProj, prefix: "\(prefix).self_attn.k_proj", from: weights)
+ CommonWeightLoader.applyLinearWeights(to: layer.selfAttn.vProj, prefix: "\(prefix).self_attn.v_proj", from: weights)
+ CommonWeightLoader.applyLinearWeights(to: layer.selfAttn.oProj, prefix: "\(prefix).self_attn.o_proj", from: weights)
+ CommonWeightLoader.applyRMSNormWeights(to: layer.selfAttn.qNorm, prefix: "\(prefix).self_attn.q_norm", from: weights)
+ CommonWeightLoader.applyRMSNormWeights(to: layer.selfAttn.kNorm, prefix: "\(prefix).self_attn.k_norm", from: weights)
+ CommonWeightLoader.applyRMSNormWeights(to: layer.inputLayerNorm, prefix: "\(prefix).input_layernorm", from: weights)
+ CommonWeightLoader.applyRMSNormWeights(to: layer.postAttentionLayerNorm, prefix: "\(prefix).post_attention_layernorm", from: weights)
+ CommonWeightLoader.applyLinearWeights(to: layer.mlp.gateProj, prefix: "\(prefix).mlp.gate_proj", from: weights)
+ CommonWeightLoader.applyLinearWeights(to: layer.mlp.upProj, prefix: "\(prefix).mlp.up_proj", from: weights)
+ CommonWeightLoader.applyLinearWeights(to: layer.mlp.downProj, prefix: "\(prefix).mlp.down_proj", from: weights)
+ }
+
+ // MARK: - Audio Encoder Weight Helpers
+
+ private static func applyConv2dWeights(
+ to conv: Conv2d,
+ prefix: String,
+ from weights: [String: MLXArray],
+ transposePyTorch: Bool = false
+ ) {
+ var params: [String: NestedItem] = [:]
+
+ if let weight = weights["\(prefix).weight"] {
+ if transposePyTorch {
+ // PyTorch Conv2d: [outC, inC, kH, kW] -> MLX Conv2d: [outC, kH, kW, inC]
+ params["weight"] = .value(weight.transposed(0, 2, 3, 1))
+ } else {
+ params["weight"] = .value(weight)
+ }
+ }
+ if let bias = weights["\(prefix).bias"] {
+ params["bias"] = .value(bias)
+ }
+
+ if !params.isEmpty {
+ conv.update(parameters: ModuleParameters(values: params))
+ }
+ }
+
+ private static func applyEncoderLayerWeights(
+ to layer: AudioEncoderLayer,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ // Self attention
+ CommonWeightLoader.applyLinearWeights(to: layer.selfAttn.qProj, prefix: "\(prefix).self_attn.q_proj", from: weights)
+ CommonWeightLoader.applyLinearWeights(to: layer.selfAttn.kProj, prefix: "\(prefix).self_attn.k_proj", from: weights)
+ CommonWeightLoader.applyLinearWeights(to: layer.selfAttn.vProj, prefix: "\(prefix).self_attn.v_proj", from: weights)
+ CommonWeightLoader.applyLinearWeights(to: layer.selfAttn.outProj, prefix: "\(prefix).self_attn.out_proj", from: weights)
+
+ // Layer norms
+ CommonWeightLoader.applyLayerNormWeights(to: layer.selfAttnLayerNorm, prefix: "\(prefix).self_attn_layer_norm", from: weights)
+ CommonWeightLoader.applyLayerNormWeights(to: layer.finalLayerNorm, prefix: "\(prefix).final_layer_norm", from: weights)
+
+ // FFN
+ CommonWeightLoader.applyLinearWeights(to: layer.fc1, prefix: "\(prefix).fc1", from: weights)
+ CommonWeightLoader.applyLinearWeights(to: layer.fc2, prefix: "\(prefix).fc2", from: weights)
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/ChatSampler.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/ChatSampler.swift
new file mode 100644
index 0000000..0df5e89
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/ChatSampler.swift
@@ -0,0 +1,107 @@
+import Foundation
+
+/// Standalone token sampler for text generation.
+///
+/// Supports temperature scaling, top-K filtering, top-P (nucleus) filtering,
+/// and repetition penalty.
+enum ChatSampler {
+
+ /// Sample a token from logits using the given sampling config.
+ ///
+ /// - Parameters:
+ /// - logits: Raw logits array of size vocab_size
+ /// - config: Sampling parameters (temperature, topK, topP, repetitionPenalty)
+ /// - previousTokens: Recently generated tokens for repetition penalty
+ /// - Returns: Sampled token index
+ static func sample(
+ logits: [Float],
+ config: ChatSamplingConfig,
+ previousTokens: [Int] = []
+ ) -> Int {
+ var logits = logits
+
+ // Repetition penalty
+ if config.repetitionPenalty > 1.0 {
+ let seen = Set(previousTokens.suffix(64))
+ for tokenId in seen {
+ if tokenId < logits.count {
+ if logits[tokenId] > 0 {
+ logits[tokenId] /= config.repetitionPenalty
+ } else {
+ logits[tokenId] *= config.repetitionPenalty
+ }
+ }
+ }
+ }
+
+ // Greedy (argmax) when temperature is 0
+ if config.temperature <= 0 {
+ var maxIdx = 0
+ var maxVal = logits[0]
+ for i in 1.. maxVal {
+ maxVal = logits[i]
+ maxIdx = i
+ }
+ }
+ return maxIdx
+ }
+
+ // Temperature scaling
+ if config.temperature != 1.0 {
+ for i in 0.. 0 && config.topK < probs.count {
+ let indexed = probs.enumerated().sorted { $0.element > $1.element }
+ let topK = Array(indexed.prefix(config.topK))
+ var filtered = [Float](repeating: 0, count: probs.count)
+ for (idx, prob) in topK {
+ filtered[idx] = prob
+ }
+ let filteredSum = filtered.reduce(0, +)
+ if filteredSum > 0 {
+ probs = filtered.map { $0 / filteredSum }
+ }
+ }
+
+ // Top-P (nucleus) filtering
+ if config.topP < 1.0 {
+ let indexed = probs.enumerated().sorted { $0.element > $1.element }
+ var cumProb: Float = 0
+ var mask = [Bool](repeating: false, count: probs.count)
+ for (idx, prob) in indexed {
+ cumProb += prob
+ mask[idx] = true
+ if cumProb >= config.topP { break }
+ }
+ for i in 0.. 0 {
+ probs = probs.map { $0 / filteredSum }
+ }
+ }
+
+ // Sample from distribution
+ let r = Float.random(in: 0..<1)
+ var cumulative: Float = 0
+ for (i, p) in probs.enumerated() {
+ cumulative += p
+ if cumulative >= r {
+ return i
+ }
+ }
+ return probs.count - 1
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/ChatTemplate.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/ChatTemplate.swift
new file mode 100644
index 0000000..4a790f9
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/ChatTemplate.swift
@@ -0,0 +1,104 @@
+import Foundation
+
+/// Chat message for Qwen3.5 models.
+public struct ChatMessage: Sendable {
+ public enum Role: String, Sendable {
+ case system
+ case user
+ case assistant
+ }
+
+ public let role: Role
+ public let content: String
+
+ public init(role: Role, content: String) {
+ self.role = role
+ self.content = content
+ }
+}
+
+/// Formats messages into Qwen3.5 chat template tokens.
+///
+/// ```
+/// <|im_start|>system
+/// {system_message}<|im_end|>
+/// <|im_start|>user
+/// {user_message}<|im_end|>
+/// <|im_start|>assistant
+/// ```
+enum ChatTemplate {
+ // Qwen3.5 special token IDs (248K vocab)
+ static let imStartId = 248045 // <|im_start|>
+ static let imEndId = 248046 // <|im_end|>
+ static let endOfTextId = 248044 // <|endoftext|>
+ static let thinkStartId = 248068 //
+ static let thinkEndId = 248069 //
+ static let newlineId = 198 // \n
+
+ /// Strip thinking block from generated tokens.
+ ///
+ /// Removes tokens from `` through `` (inclusive)
+ /// and any trailing newlines, returning only the response content.
+ static func stripThinking(from tokens: [Int]) -> [Int] {
+ let thinkTokens: Set = [thinkStartId, thinkEndId]
+ let newlines: Set = [newlineId, 271] // 198 = \n, 271 = \n\n
+
+ guard let startIdx = tokens.firstIndex(where: { thinkTokens.contains($0) && $0 == thinkStartId }) else {
+ // No — strip any leading + newlines
+ // (happens when non-thinking template causes model to echo end-think)
+ var i = 0
+ while i < tokens.count && (tokens[i] == thinkEndId || newlines.contains(tokens[i])) {
+ i += 1
+ }
+ return i > 0 ? Array(tokens[i...]) : tokens
+ }
+ if let endIdx = tokens[startIdx...].firstIndex(of: thinkEndId) {
+ var afterThink = endIdx + 1
+ while afterThink < tokens.count && newlines.contains(tokens[afterThink]) {
+ afterThink += 1
+ }
+ return Array(tokens[0.. [Int] {
+ var tokens: [Int] = []
+
+ for message in messages {
+ tokens.append(imStartId)
+ tokens.append(contentsOf: tokenizer.encode(message.role.rawValue))
+ tokens.append(newlineId)
+ tokens.append(contentsOf: tokenizer.encode(message.content))
+ tokens.append(imEndId)
+ tokens.append(newlineId)
+ }
+
+ if addGenerationPrompt {
+ tokens.append(imStartId)
+ tokens.append(contentsOf: tokenizer.encode("assistant"))
+ tokens.append(newlineId)
+
+ if !enableThinking {
+ let doubleNewline = tokenizer.encode("\n\n")
+ tokens.append(thinkStartId)
+ tokens.append(contentsOf: doubleNewline)
+ tokens.append(thinkEndId)
+ tokens.append(contentsOf: doubleNewline)
+ }
+ }
+
+ return tokens
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/ChatTokenizer.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/ChatTokenizer.swift
new file mode 100644
index 0000000..147c5c2
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/ChatTokenizer.swift
@@ -0,0 +1,272 @@
+import Foundation
+
+/// Tokenizer for Qwen3 chat models.
+///
+/// Loads vocabulary from HuggingFace tokenizer files and provides
+/// encode/decode functionality for chat text.
+public final class ChatTokenizer: @unchecked Sendable {
+ private var idToToken: [Int: String] = [:]
+ private var tokenToId: [String: Int] = [:]
+ private var bpeMerges: [(String, String)] = []
+ private var bpeMergeRanks: [String: Int] = [:]
+ private var addedTokens: [String: Int] = [:]
+
+ public var eosTokenId: Int = 248046 // <|im_end|>
+ public var vocabSize: Int { idToToken.count }
+
+ public init() {}
+
+ /// Load tokenizer from a directory.
+ ///
+ /// Supports two formats:
+ /// 1. `tokenizer.json` (HuggingFace format, preferred) — contains vocab, merges, and added tokens
+ /// 2. `vocab.json` + `merges.txt` (legacy) — separate files
+ public func load(from directory: URL) throws {
+ let tokenizerJsonURL = directory.appendingPathComponent("tokenizer.json")
+ let vocabURL = directory.appendingPathComponent("vocab.json")
+
+ if FileManager.default.fileExists(atPath: tokenizerJsonURL.path) {
+ try loadFromTokenizerJson(from: tokenizerJsonURL)
+ } else {
+ try loadVocab(from: vocabURL)
+
+ let mergesURL = directory.appendingPathComponent("merges.txt")
+ if FileManager.default.fileExists(atPath: mergesURL.path) {
+ try loadMerges(from: mergesURL)
+ }
+ }
+
+ let configURL = directory.appendingPathComponent("tokenizer_config.json")
+ if FileManager.default.fileExists(atPath: configURL.path) {
+ try loadAddedTokens(from: configURL)
+ }
+ }
+
+ /// Load from HuggingFace tokenizer.json (contains vocab + merges + added tokens).
+ private func loadFromTokenizerJson(from url: URL) throws {
+ let data = try Data(contentsOf: url)
+ guard let root = try JSONSerialization.jsonObject(with: data) as? [String: Any],
+ let model = root["model"] as? [String: Any],
+ let vocab = model["vocab"] as? [String: Int] else {
+ throw ChatModelError.tokenizerLoadFailed("Invalid tokenizer.json format")
+ }
+
+ tokenToId = vocab
+ idToToken = Dictionary(uniqueKeysWithValues: vocab.map { ($1, $0) })
+
+ // Load merges
+ if let merges = model["merges"] as? [[String]] {
+ for (i, pair) in merges.enumerated() {
+ guard pair.count == 2 else { continue }
+ bpeMerges.append((pair[0], pair[1]))
+ bpeMergeRanks["\(pair[0]) \(pair[1])"] = i
+ }
+ } else if let merges = model["merges"] as? [String] {
+ // Alternative format: each merge as "a b" string
+ for (i, merge) in merges.enumerated() {
+ let parts = merge.split(separator: " ", maxSplits: 1)
+ guard parts.count == 2 else { continue }
+ let pair = (String(parts[0]), String(parts[1]))
+ bpeMerges.append(pair)
+ bpeMergeRanks["\(pair.0) \(pair.1)"] = i
+ }
+ }
+
+ // Load added tokens
+ if let addedList = root["added_tokens"] as? [[String: Any]] {
+ for entry in addedList {
+ guard let content = entry["content"] as? String,
+ let id = entry["id"] as? Int else { continue }
+ addedTokens[content] = id
+ tokenToId[content] = id
+ idToToken[id] = content
+ }
+ }
+ }
+
+ private func loadVocab(from url: URL) throws {
+ let data = try Data(contentsOf: url)
+ guard let vocab = try JSONSerialization.jsonObject(with: data) as? [String: Int] else {
+ throw ChatModelError.tokenizerLoadFailed("Invalid vocab.json format")
+ }
+ tokenToId = vocab
+ idToToken = Dictionary(uniqueKeysWithValues: vocab.map { ($1, $0) })
+ }
+
+ private func loadMerges(from url: URL) throws {
+ let content = try String(contentsOf: url, encoding: .utf8)
+ let lines = content.components(separatedBy: "\n")
+ for (i, line) in lines.enumerated() {
+ if line.hasPrefix("#") || line.isEmpty { continue }
+ let parts = line.split(separator: " ", maxSplits: 1)
+ if parts.count == 2 {
+ let pair = (String(parts[0]), String(parts[1]))
+ bpeMerges.append(pair)
+ bpeMergeRanks["\(pair.0) \(pair.1)"] = i
+ }
+ }
+ }
+
+ private func loadAddedTokens(from url: URL) throws {
+ let data = try Data(contentsOf: url)
+ guard let config = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
+ return
+ }
+ if let added = config["added_tokens_decoder"] as? [String: Any] {
+ for (idStr, value) in added {
+ guard let id = Int(idStr),
+ let info = value as? [String: Any],
+ let content = info["content"] as? String else { continue }
+ addedTokens[content] = id
+ tokenToId[content] = id
+ idToToken[id] = content
+ }
+ }
+ }
+
+ // MARK: - Encode
+
+ /// Encode text to token IDs using BPE.
+ public func encode(_ text: String) -> [Int] {
+ if text.isEmpty { return [] }
+
+ // Check if it's a special/added token
+ if let id = addedTokens[text] ?? tokenToId[text] {
+ return [id]
+ }
+
+ // Simple BPE encoding
+ var words = tokenizeToWords(text)
+ var allTokens: [Int] = []
+
+ for word in words {
+ let wordTokens = bpeEncode(word)
+ allTokens.append(contentsOf: wordTokens)
+ }
+
+ return allTokens
+ }
+
+ /// Split text into BPE-ready words (GPT-2/Qwen style).
+ private func tokenizeToWords(_ text: String) -> [String] {
+ // Simplified: split on spaces, prefix non-first words with Ġ (space marker)
+ var words: [String] = []
+ var isFirst = true
+ for part in text.components(separatedBy: " ") {
+ if part.isEmpty { continue }
+ if isFirst {
+ words.append(part)
+ isFirst = false
+ } else {
+ words.append("Ġ" + part)
+ }
+ }
+ return words
+ }
+
+ /// BPE encode a single word.
+ private func bpeEncode(_ word: String) -> [Int] {
+ if let id = tokenToId[word] {
+ return [id]
+ }
+
+ var symbols = word.map { String($0) }
+ if symbols.isEmpty { return [] }
+
+ while symbols.count > 1 {
+ // Find best merge
+ var bestRank = Int.max
+ var bestIdx = -1
+ for i in 0..<(symbols.count - 1) {
+ let pair = "\(symbols[i]) \(symbols[i + 1])"
+ if let rank = bpeMergeRanks[pair], rank < bestRank {
+ bestRank = rank
+ bestIdx = i
+ }
+ }
+ if bestIdx < 0 { break }
+
+ // Apply merge
+ let merged = symbols[bestIdx] + symbols[bestIdx + 1]
+ symbols.replaceSubrange(bestIdx...bestIdx + 1, with: [merged])
+ }
+
+ // Look up token IDs
+ return symbols.compactMap { tokenToId[$0] }
+ }
+
+ // MARK: - Decode
+
+ /// Decode token IDs to text.
+ ///
+ /// Uses byte-level BPE decoding: token strings are mapped back to bytes
+ /// via the GPT-2 byte-to-unicode table, then assembled into UTF-8 text.
+ public func decode(_ tokenIds: [Int]) -> String {
+ let pieces = tokenIds.compactMap { idToToken[$0] }
+ let joined = pieces.joined()
+ return decodeBPEString(joined)
+ }
+
+ /// Decode a single token ID.
+ public func decodeToken(_ tokenId: Int) -> String? {
+ guard let piece = idToToken[tokenId] else { return nil }
+ return decodeBPEString(piece)
+ }
+
+ /// Convert a BPE token string to UTF-8 text.
+ ///
+ /// GPT-2/Qwen byte-level BPE represents each byte as a specific Unicode
+ /// character. This reverses that mapping and decodes the bytes as UTF-8.
+ private func decodeBPEString(_ bpeString: String) -> String {
+ var bytes: [UInt8] = []
+ for char in bpeString {
+ if let byte = Self.unicodeToByte[char] {
+ bytes.append(byte)
+ }
+ }
+ return String(bytes: bytes, encoding: .utf8) ?? bpeString
+ }
+
+ /// GPT-2 byte-to-unicode mapping (reversed for decoding).
+ ///
+ /// Maps Unicode characters back to the byte values they represent in
+ /// GPT-2/Qwen byte-level BPE vocabulary.
+ private static let unicodeToByte: [Character: UInt8] = {
+ // Build the standard GPT-2 bytes_to_unicode table
+ var byteToUnicode: [UInt8: Character] = [:]
+ var n = 0
+
+ // Printable ASCII + Latin supplement ranges that map to themselves
+ let ranges: [ClosedRange] = [
+ 33...126, // ! through ~
+ 161...172, // ¡ through ¬
+ 174...255, // ® through ÿ
+ ]
+ for range in ranges {
+ for b in range {
+ byteToUnicode[b] = Character(Unicode.Scalar(UInt32(b))!)
+ }
+ }
+
+ // Remaining bytes (0-32, 127-160, 173) map to 256+n
+ for b: UInt16 in 0...255 {
+ if byteToUnicode[UInt8(b)] == nil {
+ byteToUnicode[UInt8(b)] = Character(Unicode.Scalar(256 + UInt32(n))!)
+ n += 1
+ }
+ }
+
+ // Reverse the mapping: unicode char → byte value
+ var result: [Character: UInt8] = [:]
+ for (byte, char) in byteToUnicode {
+ result[char] = byte
+ }
+ return result
+ }()
+
+ /// Check if a token ID is a special token (should not appear in output).
+ public func isSpecialToken(_ tokenId: Int) -> Bool {
+ guard let token = idToToken[tokenId] else { return false }
+ return token.hasPrefix("<|") && token.hasSuffix("|>")
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/MLXGenerator.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/MLXGenerator.swift
new file mode 100644
index 0000000..a8d3fef
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/MLXGenerator.swift
@@ -0,0 +1,425 @@
+import Foundation
+import MLXCommon
+import MLX
+import MLXNN
+import AudioCommon
+import os.log
+
+private let log = OSLog(subsystem: "com.soniqo.qwen3chat", category: "MLX")
+
+// MARK: - MLX Generator for Qwen3.5 Chat
+
+/// MLX-based text generator for Qwen3.5-0.8B hybrid model.
+///
+/// Uses MLX for GPU inference
+/// on Apple Silicon GPUs. The hybrid DeltaNet + GatedAttention architecture
+/// requires managing two types of state:
+/// 1. DeltaNet recurrent states (carried across all tokens, O(1) per layer)
+/// 2. GatedAttention KV caches (grow with sequence length, only 6 layers)
+///
+/// Usage:
+/// ```swift
+/// let model = try await Qwen35MLXChat.fromPretrained()
+/// let response = try model.generate(messages: [
+/// ChatMessage(role: .user, content: "Hello!")
+/// ])
+/// ```
+public final class Qwen35MLXChat: @unchecked Sendable {
+ public static let defaultModelId = "aufklarer/Qwen3.5-0.8B-Chat-MLX"
+
+ public let config: Qwen3ChatConfig
+ public let tokenizer: ChatTokenizer
+ let model: Qwen35MLXModel
+ var state: Qwen35MLXModel.InferenceState
+ var _isLoaded = true
+
+ // MARK: - Metrics
+
+ /// Generation metrics for performance tracking.
+ public struct Metrics {
+ public var prefillTimeMs: Double = 0
+ public var prefillTokens: Int = 0
+ public var decodeTimeMs: Double = 0
+ public var decodeTokens: Int = 0
+
+ public var tokensPerSecond: Double {
+ guard decodeTimeMs > 0 else { return 0 }
+ return Double(decodeTokens) / (decodeTimeMs / 1000.0)
+ }
+
+ public var msPerToken: Double {
+ guard decodeTokens > 0 else { return 0 }
+ return decodeTimeMs / Double(decodeTokens)
+ }
+
+ public var prefillTokensPerSecond: Double {
+ guard prefillTimeMs > 0 else { return 0 }
+ return Double(prefillTokens) / (prefillTimeMs / 1000.0)
+ }
+ }
+
+ private(set) var metrics = Metrics()
+
+ /// Latest generation metrics.
+ public var lastMetrics: (tokensPerSec: Double, prefillMs: Double, decodeMs: Double, msPerToken: Double) {
+ (metrics.tokensPerSecond, metrics.prefillTimeMs, metrics.decodeTimeMs, metrics.msPerToken)
+ }
+
+ // MARK: - Init
+
+ private init(config: Qwen3ChatConfig, tokenizer: ChatTokenizer, model: Qwen35MLXModel) {
+ self.config = config
+ self.tokenizer = tokenizer
+ self.model = model
+ self.state = .initial(config: config)
+ }
+
+ // MARK: - Factory
+
+ /// Quantization variant.
+ public enum Quantization: String {
+ case int4
+ case int8
+ }
+
+ /// Load a pre-trained Qwen3.5 chat model from HuggingFace.
+ ///
+ /// Downloads quantized safetensors and tokenizer on first use.
+ /// Model is loaded into MLX for GPU inference on Apple Silicon.
+ ///
+ /// - Parameters:
+ /// - modelId: HuggingFace model ID (repo with int4/ and int8/ subdirs)
+ /// - quantization: INT4 (404 MB) or INT8 (763 MB)
+ /// - progressHandler: Optional callback for download/load progress
+ public static func fromPretrained(
+ modelId: String = defaultModelId,
+ quantization: Quantization = .int4,
+ cacheDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> Qwen35MLXChat {
+ let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
+ let variant = quantization.rawValue
+
+ // Download model files from variant subdirectory (int4/ or int8/)
+ progressHandler?(0.05, "Downloading \(variant) model...")
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: modelId,
+ to: cacheDir,
+ additionalFiles: [
+ "\(variant)/model.safetensors",
+ "\(variant)/config.json",
+ "\(variant)/tokenizer.json",
+ "\(variant)/tokenizer_config.json",
+ ],
+ offlineMode: offlineMode,
+ progressHandler: { progress in
+ progressHandler?(progress * 0.5, "Downloading...")
+ }
+ )
+
+ // Variant files are in a subdirectory
+ let variantDir = cacheDir.appendingPathComponent(variant)
+
+ // Load config
+ progressHandler?(0.5, "Loading config...")
+ let config: Qwen3ChatConfig
+ let configURL = variantDir.appendingPathComponent("config.json")
+ if FileManager.default.fileExists(atPath: configURL.path) {
+ config = try Qwen3ChatConfig.load(from: configURL)
+ } else {
+ config = .qwen35_08B
+ }
+
+ // Load tokenizer
+ progressHandler?(0.55, "Loading tokenizer...")
+ let tokenizer = ChatTokenizer()
+ try tokenizer.load(from: variantDir)
+
+ // Create model
+ progressHandler?(0.6, "Creating model...")
+ let model = Qwen35MLXModel(config: config)
+
+ // Load weights
+ progressHandler?(0.65, "Loading weights...")
+ try Qwen35WeightLoader.loadWeights(
+ into: model, from: variantDir,
+ progressHandler: { pct, msg in
+ progressHandler?(0.65 + pct * 0.3, msg)
+ })
+
+ progressHandler?(1.0, "Ready")
+ return Qwen35MLXChat(config: config, tokenizer: tokenizer, model: model)
+ }
+
+ /// Download tokenizer + weight files only — does not load MLX/Metal.
+ public static func downloadWeightsOnly(
+ modelId: String = defaultModelId,
+ quantization: Quantization = .int4,
+ cacheDir: URL? = nil,
+ registry: ModelRegistry = .huggingFace(),
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws {
+ let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
+ let variant = quantization.rawValue
+ progressHandler?(0.0, "Downloading \(variant) model...")
+ let files = [
+ "\(variant)/model.safetensors",
+ "\(variant)/config.json",
+ "\(variant)/tokenizer.json",
+ "\(variant)/tokenizer_config.json",
+ ]
+ switch registry {
+ case .huggingFace(let hubEndpoint):
+ try await HuggingFaceDownloader.downloadFiles(
+ modelId: modelId,
+ to: cacheDir,
+ files: files,
+ hubEndpoint: hubEndpoint,
+ progressHandler: { progress in
+ progressHandler?(progress, "Downloading...")
+ }
+ )
+ case .modelScope(let baseURL, let revision):
+ try await ModelScopeDownloader.downloadFiles(
+ modelId: modelId,
+ to: cacheDir,
+ files: files,
+ baseURL: baseURL,
+ revision: revision,
+ progressHandler: { progress in
+ progressHandler?(progress, "Downloading...")
+ }
+ )
+ }
+ progressHandler?(1.0, "Downloaded")
+ }
+
+ /// Load from a local directory (no HuggingFace download).
+ public static func fromLocal(
+ directory: URL,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> Qwen35MLXChat {
+ let config: Qwen3ChatConfig
+ let configURL = directory.appendingPathComponent("config.json")
+ if FileManager.default.fileExists(atPath: configURL.path) {
+ config = try Qwen3ChatConfig.load(from: configURL)
+ } else {
+ config = .qwen35_08B
+ }
+
+ let tokenizer = ChatTokenizer()
+ try tokenizer.load(from: directory)
+
+ let model = Qwen35MLXModel(config: config)
+ try Qwen35WeightLoader.loadWeights(
+ into: model, from: directory,
+ progressHandler: progressHandler)
+
+ return Qwen35MLXChat(config: config, tokenizer: tokenizer, model: model)
+ }
+
+ // MARK: - State Management
+
+ /// Reset all inference state for a new conversation.
+ public func resetState() {
+ state = .initial(config: config)
+ metrics = Metrics()
+ }
+
+ // MARK: - Generation
+
+ /// Generate a response from chat messages.
+ ///
+ /// Encodes the messages using the chat template, prefills the prompt,
+ /// then generates tokens autoregressively until EOS or max tokens.
+ public func generate(
+ messages: [ChatMessage],
+ sampling: ChatSamplingConfig = .default
+ ) throws -> String {
+ resetState()
+
+ let promptTokens = ChatTemplate.encode(
+ messages: messages,
+ tokenizer: tokenizer,
+ config: config,
+ enableThinking: false)
+
+ // Prefill
+ let prefillStart = CFAbsoluteTimeGetCurrent()
+ let promptArray = MLXArray(promptTokens.map { Int32($0) })
+ .expandedDimensions(axis: 0)
+ let (prefillLogits, prefillState) = model.forward(inputIds: promptArray, state: state)
+ eval(prefillLogits)
+ state = prefillState
+
+ let prefillMs = (CFAbsoluteTimeGetCurrent() - prefillStart) * 1000
+ metrics.prefillTimeMs = prefillMs
+ metrics.prefillTokens = promptTokens.count
+
+ // Extract last-position logits and sample first token
+ var logits = extractLastPositionLogits(prefillLogits)
+ var generatedTokens: [Int] = []
+ var inThinking = false
+ let thinkBudget = 100
+
+ // Decode loop
+ let decodeStart = CFAbsoluteTimeGetCurrent()
+
+ for _ in 0..<(sampling.maxTokens + thinkBudget) {
+ let nextToken = ChatSampler.sample(
+ logits: logits,
+ config: sampling,
+ previousTokens: promptTokens + generatedTokens)
+
+ if nextToken == config.eosTokenId { break }
+ if nextToken == ChatTemplate.imEndId { break }
+
+ generatedTokens.append(nextToken)
+
+ // Thinking token tracking (handle both Qwen3 and Qwen3.5 token IDs)
+ if nextToken == ChatTemplate.thinkStartId {
+ inThinking = true
+ } else if nextToken == ChatTemplate.thinkEndId {
+ inThinking = false
+ }
+
+ if inThinking && generatedTokens.count > thinkBudget {
+ let thinkEnd = ChatTemplate.thinkEndId
+ generatedTokens.append(thinkEnd)
+ let tokenArr = MLXArray([Int32(thinkEnd)])
+ .expandedDimensions(axis: 0)
+ let (stepLogits, newState) = model.forward(inputIds: tokenArr, state: state)
+ eval(stepLogits)
+ state = newState
+ logits = extractLastPositionLogits(stepLogits)
+ inThinking = false
+ continue
+ }
+
+ let thinkTokens: Set = [
+ ChatTemplate.thinkStartId, ChatTemplate.thinkEndId
+ ]
+ let responseCount = generatedTokens.filter { !thinkTokens.contains($0) }.count
+ if !inThinking && responseCount >= sampling.maxTokens { break }
+
+ // Decode one step
+ let tokenArr = MLXArray([Int32(nextToken)]).expandedDimensions(axis: 0)
+ let (stepLogits, newState) = model.forward(inputIds: tokenArr, state: state)
+ eval(stepLogits)
+ state = newState
+ logits = extractLastPositionLogits(stepLogits)
+ }
+
+ let decodeMs = (CFAbsoluteTimeGetCurrent() - decodeStart) * 1000
+ metrics.decodeTimeMs = decodeMs
+ metrics.decodeTokens = generatedTokens.count
+
+ var memInfo = mach_task_basic_info()
+ var memCount = mach_msg_type_number_t(MemoryLayout.size) / 4
+ _ = withUnsafeMutablePointer(to: &memInfo) {
+ $0.withMemoryRebound(to: integer_t.self, capacity: Int(memCount)) {
+ task_info(mach_task_self_, task_flavor_t(MACH_TASK_BASIC_INFO), $0, &memCount)
+ }
+ }
+ let memMB = Double(memInfo.resident_size) / 1024 / 1024
+ let tps = decodeMs > 0 ? Double(generatedTokens.count) / (decodeMs / 1000.0) : 0
+ os_log(.info, log: log,
+ "Generate done: prefill=%.0fms (%d tokens), decode=%.0fms (%d tokens, %.1f tok/s), memory=%.0f MB",
+ metrics.prefillTimeMs, promptTokens.count, decodeMs, generatedTokens.count, tps, memMB)
+
+ let responseTokens = ChatTemplate.stripThinking(from: generatedTokens)
+ return tokenizer.decode(responseTokens)
+ }
+
+ /// Generate a streaming response.
+ public func generateStream(
+ messages: [ChatMessage],
+ sampling: ChatSamplingConfig = .default
+ ) -> AsyncThrowingStream {
+ AsyncThrowingStream { continuation in
+ Task {
+ self.resetState()
+
+ let promptTokens = ChatTemplate.encode(
+ messages: messages,
+ tokenizer: self.tokenizer,
+ config: self.config,
+ enableThinking: false)
+
+ let promptArray = MLXArray(promptTokens.map { Int32($0) })
+ .expandedDimensions(axis: 0)
+ let (prefillLogits, prefillState) = self.model.forward(
+ inputIds: promptArray, state: self.state)
+ eval(prefillLogits)
+ self.state = prefillState
+
+ var logits = self.extractLastPositionLogits(prefillLogits)
+ var generatedTokens: [Int] = []
+ var inThinking = false
+
+ for _ in 0.. [Float] {
+ let t = logits.dim(1)
+ let lastPos = logits[0, t - 1].asType(.float32) // [vocabSize]
+ eval(lastPos)
+ // Bulk extract all floats at once — do NOT use per-element .item() (248K syncs)
+ let all: [Float] = lastPos.asArray(Float.self)
+ return Array(all.prefix(config.vocabSize))
+ }
+}
+
+// MARK: - Memory Management
+
+extension Qwen35MLXChat: ModelMemoryManageable {
+ public var isLoaded: Bool { _isLoaded }
+
+ public func unload() {
+ guard _isLoaded else { return }
+ model.clearParameters()
+ state = .initial(config: config)
+ _isLoaded = false
+ }
+
+ public var memoryFootprint: Int {
+ guard _isLoaded else { return 0 }
+ return model.parameterMemoryBytes()
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen35CoreMLChat.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen35CoreMLChat.swift
new file mode 100644
index 0000000..b9bccfc
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen35CoreMLChat.swift
@@ -0,0 +1,381 @@
+import CoreML
+import Foundation
+import AudioCommon
+import os.log
+
+private let log = OSLog(subsystem: "com.soniqo.qwen3chat", category: "CoreML")
+
+/// CoreML-based Qwen3.5-0.8B chat for iOS Neural Engine.
+///
+/// Uses two CoreML models:
+/// - `embedding.mlmodelc` — token ID → embedding vector
+/// - `decoder.mlmodelc` — autoregressive transformer with MLState
+///
+/// All DeltaNet recurrent states and GatedAttention KV caches are managed
+/// by CoreML's MLState API — no manual cache tracking needed.
+public final class Qwen35CoreMLChat: @unchecked Sendable {
+ public static let defaultModelId = "aufklarer/Qwen3.5-0.8B-Chat-CoreML"
+
+ private let embeddingModel: MLModel
+ private let decoderModel: MLModel
+ private var decoderState: MLState
+ public let config: Qwen3ChatConfig
+ public let tokenizer: ChatTokenizer
+ private var position: Int = 0
+ private let maxSeqLen: Int
+
+ /// Quantization variant. Only INT8 available (INT4 removed — CoreML dequantization issues).
+ public enum Quantization: String { case int8 }
+
+ // MARK: - Metrics
+
+ private var _prefillMs: Double = 0
+ private var _decodeMs: Double = 0
+ private var _decodeTokens: Int = 0
+
+ public var lastMetrics: (tokensPerSec: Double, prefillMs: Double, decodeMs: Double, msPerToken: Double) {
+ let tps = _decodeMs > 0 ? Double(_decodeTokens) / (_decodeMs / 1000.0) : 0
+ let mpt = _decodeTokens > 0 ? _decodeMs / Double(_decodeTokens) : 0
+ return (tps, _prefillMs, _decodeMs, mpt)
+ }
+
+ /// Current process memory in MB.
+ private static var memoryMB: Double {
+ var info = mach_task_basic_info()
+ var count = mach_msg_type_number_t(MemoryLayout.size) / 4
+ let result = withUnsafeMutablePointer(to: &info) {
+ $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) {
+ task_info(mach_task_self_, task_flavor_t(MACH_TASK_BASIC_INFO), $0, &count)
+ }
+ }
+ guard result == KERN_SUCCESS else { return 0 }
+ return Double(info.resident_size) / 1024 / 1024
+ }
+
+ // MARK: - Init
+
+ private init(embedding: MLModel, decoder: MLModel, state: MLState,
+ config: Qwen3ChatConfig, tokenizer: ChatTokenizer, maxSeqLen: Int) {
+ self.embeddingModel = embedding
+ self.decoderModel = decoder
+ self.decoderState = state
+ self.config = config
+ self.tokenizer = tokenizer
+ self.maxSeqLen = maxSeqLen
+ }
+
+ // MARK: - Factory
+
+ /// Load from HuggingFace.
+ public static func fromPretrained(
+ modelId: String = defaultModelId,
+ quantization: Quantization = .int8,
+ computeUnits: MLComputeUnits = .cpuAndNeuralEngine,
+ cacheDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> Qwen35CoreMLChat {
+ let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
+ let variant = quantization.rawValue
+
+ progressHandler?(0.05, "Downloading \(variant) model...")
+ // Fetch the pre-compiled ``.mlmodelc`` bundle only. On-device
+ // ``MLModel.compileModel`` drifts per runtime, and the legacy
+ // ``.mlpackage`` internals (``*.mlmodel`` / ``Manifest.json``) would
+ // force that code path. Users with a stale ``.mlpackage`` cache
+ // transparently re-download because ``.mlmodelc`` is missing.
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: modelId,
+ to: cacheDir,
+ additionalFiles: [
+ "\(variant)/*.json",
+ "\(variant)/embedding.mlmodelc/**",
+ "\(variant)/decoder.mlmodelc/**",
+ ],
+ offlineMode: offlineMode,
+ progressHandler: { p in progressHandler?(p * 0.5, "Downloading...") }
+ )
+
+ let variantDir = cacheDir.appendingPathComponent(variant)
+ return try await fromLocal(directory: variantDir, computeUnits: computeUnits,
+ progressHandler: progressHandler)
+ }
+
+ /// Load from a local directory.
+ public static func fromLocal(
+ directory: URL,
+ computeUnits: MLComputeUnits = .cpuAndNeuralEngine,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> Qwen35CoreMLChat {
+ progressHandler?(0.5, "Loading config...")
+
+ // Debug: list directory contents to diagnose missing file issues
+ os_log(.info, log: log, "Loading from directory: %{public}@", directory.path)
+ if let contents = try? FileManager.default.contentsOfDirectory(at: directory, includingPropertiesForKeys: nil) {
+ for item in contents {
+ os_log(.info, log: log, " %{public}@", item.lastPathComponent)
+ }
+ } else {
+ os_log(.error, log: log, "Cannot list directory: %{public}@", directory.path)
+ }
+
+ // Use built-in config — chat_config.json from CoreML conversion has a different schema
+ let config = Qwen3ChatConfig.qwen35_08B
+ os_log(.info, log: log, "Using built-in Qwen3.5-0.8B config")
+
+ progressHandler?(0.55, "Loading tokenizer...")
+ os_log(.info, log: log, "Loading tokenizer from: %{public}@", directory.resolvingSymlinksInPath().path)
+ let tokenizer = ChatTokenizer()
+ try tokenizer.load(from: directory.resolvingSymlinksInPath())
+
+ let memBefore = memoryMB
+ os_log(.info, log: log, "Loading CoreML models, memory before: %.0f MB", memBefore)
+
+ progressHandler?(0.6, "Compiling embedding...")
+ let embModel: MLModel
+ do {
+ embModel = try await loadModel(named: "embedding", from: directory, computeUnits: computeUnits)
+ os_log(.info, log: log, "Embedding loaded, memory: %.0f MB", memoryMB)
+ } catch {
+ os_log(.error, log: log, "Embedding load FAILED: %{public}@", error.localizedDescription)
+ throw error
+ }
+
+ progressHandler?(0.75, "Compiling decoder...")
+ let decModel: MLModel
+ do {
+ decModel = try await loadModel(named: "decoder", from: directory, computeUnits: computeUnits)
+ os_log(.info, log: log, "Decoder loaded, memory: %.0f MB", memoryMB)
+ } catch {
+ os_log(.error, log: log, "Decoder load FAILED: %{public}@", error.localizedDescription)
+ throw error
+ }
+
+ let state = decModel.makeState()
+ let maxSeq = config.maxSeqLen
+
+ os_log(.info, log: log, "Model ready, total memory: %.0f MB (delta: +%.0f MB)",
+ memoryMB, memoryMB - memBefore)
+ progressHandler?(1.0, "Ready")
+ return Qwen35CoreMLChat(
+ embedding: embModel, decoder: decModel, state: state,
+ config: config, tokenizer: tokenizer, maxSeqLen: maxSeq)
+ }
+
+ // MARK: - Model Loading Helpers
+
+ private static func loadModel(
+ named name: String, from dir: URL, computeUnits: MLComputeUnits
+ ) async throws -> MLModel {
+ let compiledURL = dir.appendingPathComponent("\(name).mlmodelc")
+ guard FileManager.default.fileExists(atPath: compiledURL.path) else {
+ os_log(.error, log: log,
+ "Model not found: %{public}@.mlmodelc in %{public}@",
+ name, dir.path)
+ throw ChatModelError.modelNotFound(dir)
+ }
+
+ let mlConfig = MLModelConfiguration()
+ mlConfig.computeUnits = CoreMLComputeUnitsResolver.resolved(default: computeUnits)
+ return try await MLModel.load(contentsOf: compiledURL, configuration: mlConfig)
+ }
+
+ // MARK: - Generation
+
+ /// Reset state for a new conversation.
+ public func resetState() {
+ decoderState = decoderModel.makeState()
+ position = 0
+ _prefillMs = 0; _decodeMs = 0; _decodeTokens = 0
+ }
+
+ /// Generate a response from chat messages.
+ public func generate(
+ messages: [ChatMessage],
+ sampling: ChatSamplingConfig = .default
+ ) throws -> String {
+ resetState()
+
+ let promptTokens = ChatTemplate.encode(
+ messages: messages, tokenizer: tokenizer,
+ config: config, enableThinking: false)
+
+ // Prefill: feed all prompt tokens one at a time
+ let prefillStart = CFAbsoluteTimeGetCurrent()
+ var lastLogits: [Float] = []
+
+ for token in promptTokens {
+ lastLogits = try forwardStep(tokenId: token)
+ }
+ _prefillMs = (CFAbsoluteTimeGetCurrent() - prefillStart) * 1000
+
+ // Decode loop
+ let decodeStart = CFAbsoluteTimeGetCurrent()
+ var generatedTokens: [Int] = []
+
+ for _ in 0.. 0 ? Double(_decodeTokens) / (_decodeMs / 1000.0) : 0
+ os_log(.info, log: log,
+ "Generate done: prefill=%.0fms (%d tokens), decode=%.0fms (%d tokens, %.1f tok/s), memory=%.0f MB",
+ _prefillMs, promptTokens.count, _decodeMs, _decodeTokens, tps, Self.memoryMB)
+
+ let responseTokens = ChatTemplate.stripThinking(from: generatedTokens)
+ let responseText = tokenizer.decode(responseTokens)
+ os_log(.info, log: log, "Response (%d tokens → %d after strip): '%{public}@'",
+ generatedTokens.count, responseTokens.count, responseText)
+ return responseText
+ }
+
+ /// Generate a streaming response.
+ public func generateStream(
+ messages: [ChatMessage],
+ sampling: ChatSamplingConfig = .default
+ ) -> AsyncThrowingStream {
+ AsyncThrowingStream { continuation in
+ Task {
+ do {
+ self.resetState()
+ let promptTokens = ChatTemplate.encode(
+ messages: messages, tokenizer: self.tokenizer,
+ config: self.config, enableThinking: false)
+
+ var lastLogits: [Float] = []
+ for token in promptTokens {
+ lastLogits = try self.forwardStep(tokenId: token)
+ }
+
+ var generatedTokens: [Int] = []
+ var inThinking = false
+ let thinkBudget = 100
+ let thinkTokens: Set = [
+ ChatTemplate.thinkStartId, ChatTemplate.thinkEndId
+ ]
+
+ for _ in 0..<(sampling.maxTokens + thinkBudget) {
+ let nextToken = ChatSampler.sample(
+ logits: lastLogits, config: sampling,
+ previousTokens: promptTokens + generatedTokens)
+
+ if nextToken == self.config.eosTokenId { break }
+ if nextToken == ChatTemplate.imEndId { break }
+
+ generatedTokens.append(nextToken)
+
+ if nextToken == ChatTemplate.thinkStartId { inThinking = true }
+ else if nextToken == ChatTemplate.thinkEndId { inThinking = false }
+ else if !inThinking,
+ let text = self.tokenizer.decodeToken(nextToken),
+ !self.tokenizer.isSpecialToken(nextToken) {
+ continuation.yield(text)
+ }
+
+ // Force-end thinking if budget exceeded
+ if inThinking && generatedTokens.count > thinkBudget {
+ generatedTokens.append(ChatTemplate.thinkEndId)
+ lastLogits = try self.forwardStep(tokenId: ChatTemplate.thinkEndId)
+ inThinking = false
+ continue
+ }
+
+ // Only count non-thinking tokens against maxTokens
+ let responseCount = generatedTokens.filter { !thinkTokens.contains($0) }.count
+ if !inThinking && responseCount >= sampling.maxTokens { break }
+
+ lastLogits = try self.forwardStep(tokenId: nextToken)
+ }
+ continuation.finish()
+ } catch {
+ continuation.finish(throwing: error)
+ }
+ }
+ }
+ }
+
+ // MARK: - Single Step
+
+ /// Run one decoder step: token ID → logits.
+ private func forwardStep(tokenId: Int) throws -> [Float] {
+ // Embedding lookup
+ let tokenInput = try MLMultiArray(shape: [1, 1], dataType: .int32)
+ tokenInput[0] = NSNumber(value: Int32(tokenId))
+
+ let embFeatures = try MLDictionaryFeatureProvider(dictionary: [
+ "token_id": MLFeatureValue(multiArray: tokenInput)
+ ])
+ let embResult = try embeddingModel.prediction(from: embFeatures)
+ guard let embedding = embResult.featureValue(for: "embedding")?.multiArrayValue else {
+ throw ChatModelError.inferenceFailed("Embedding output missing")
+ }
+
+ // Build attention mask: 0 for positions ≤ current, -FLT_MAX for future
+ let mask = try MLMultiArray(shape: [1, 1, 1, maxSeqLen as NSNumber], dataType: .float32)
+ let maskPtr = mask.dataPointer.bindMemory(to: Float.self, capacity: maxSeqLen)
+ for i in 0.. State {
+ State(
+ s: MLXArray.zeros([batchSize, numHeads, headDim, headDim], dtype: dtype),
+ convState: MLXArray.zeros([batchSize, qkvDim, convKernel - 1], dtype: dtype)
+ )
+ }
+ }
+
+ /// Forward pass processing a sequence of tokens.
+ ///
+ /// Implements the gated delta rule recurrence (reference: mlx-lm/gated_delta.py):
+ /// 1. Decay: S = g * S
+ /// 2. Error: kv_mem = (S * k).sum(-1); delta = (v - kv_mem) * beta
+ /// 3. Update: S = S + k * delta
+ /// 4. Output: y = (S * q).sum(-1)
+ ///
+ /// - Parameters:
+ /// - x: Input hidden states [B, T, hiddenSize]
+ /// - state: Previous recurrent state (nil for first call)
+ /// - Returns: (output [B, T, hiddenSize], updated state)
+ public func callAsFunction(_ x: MLXArray, state: State? = nil) -> (MLXArray, State) {
+ let b = x.dim(0)
+ let t = x.dim(1)
+
+ // Project inputs (separate projections matching HuggingFace weight format)
+ let qkvRaw = inProjQKV(x) // [B, T, 3*H*D=6144]
+ let zRaw = inProjZ(x) // [B, T, 2*hiddenSize=2048]
+ let bRaw = inProjB(x) // [B, T, H=16]
+ let aRaw = inProjA(x) // [B, T, H=16]
+
+
+ // Causal conv1d on QKV only (not Z, B, A)
+ let prevConvState: MLXArray
+ if let s = state {
+ prevConvState = s.convState
+ } else {
+ prevConvState = MLXArray.zeros([b, qkvDim, convKernel - 1], dtype: x.dtype)
+ }
+
+ let qkvTransposed = qkvRaw.transposed(0, 2, 1) // [B, C, T]
+ let padded = concatenated([prevConvState, qkvTransposed], axis: 2) // [B, C, T+K-1]
+
+ // Save new conv state (last K-1 columns)
+ let totalLen = padded.dim(2)
+ let newConvState = padded[0..., 0..., (totalLen - convKernel + 1)...]
+
+ // Apply depthwise causal conv1d + SiLU
+ let qkvConv = depthwiseConv1dCausal(padded, outputLen: t)
+ let qkvActivated = silu(qkvConv.transposed(0, 2, 1)) // [B, T, C]
+
+ // Split into Q, K, V — each [B, T, H, D]
+ let hd = numHeads * headDim
+ var q = qkvActivated[0..., 0..., .. MLXArray {
+ let a = aRaw + dtBias.reshaped(1, 1, numHeads)
+ let dt = softplus(a)
+ let negExpA = -exp(aLog.asType(.float32)).reshaped(1, 1, numHeads)
+ return exp(negExpA * dt.asType(.float32)).asType(aRaw.dtype)
+ }
+
+ /// RMS normalization without learnable weight (used for Q/K normalization).
+ private func rmsNormNoWeight(_ x: MLXArray) -> MLXArray {
+ let meanSq = (x * x).mean(axis: -1, keepDims: true)
+ return x * rsqrt(meanSq + MLXArray(Float(1e-6)))
+ }
+
+ // MARK: - Depthwise Conv1d
+
+ /// Depthwise causal conv1d via unfolding + element-wise multiply + sum.
+ /// - Parameter input: [B, C, T+K-1] (pre-padded with conv state)
+ /// - Parameter outputLen: number of output time steps T
+ /// - Returns: [B, C, T]
+ private func depthwiseConv1dCausal(_ input: MLXArray, outputLen: Int) -> MLXArray {
+ let c = input.dim(1)
+ let k = convKernel
+
+ // Unfold: gather windows of size K for each output position
+ var windows: [MLXArray] = []
+ windows.reserveCapacity(outputLen)
+ for t in 0.. squeeze axis 2 -> [C, K] -> [1, C, 1, K]
+ let kernelBcast = convWeight.squeezed(axis: 2).reshaped(1, c, 1, k)
+
+ return (unfolded * kernelBcast).sum(axis: -1) // [B, C, T]
+ }
+}
+
+// MARK: - Softplus
+
+private func softplus(_ x: MLXArray) -> MLXArray {
+ // Numerically stable: for large x, softplus(x) ~ x
+ MLX.where(x .> MLXArray(Float(20.0)), x, log(1 + exp(x)))
+}
+
+// MARK: - GatedAttention (Full Attention)
+
+/// GatedAttention layer for Qwen3.5 hybrid model.
+///
+/// Standard multi-head attention with:
+/// - GQA: 8 query heads, 2 KV heads, head_dim=256
+/// - Partial RoPE: only first 25% of head_dim (64 dims) get rotary encoding
+/// - QK norm: RMSNorm applied per-head to Q and K before RoPE
+/// - Gated output: q_proj produces [Q; gate], both of dim numQHeads*headDim.
+/// After attention, output is element-wise multiplied with silu(gate),
+/// then projected through o_proj.
+///
+/// Weight shapes:
+/// - q_proj: [4096, 1024] = [2 * numQHeads * headDim, hiddenSize] (Q + gate)
+/// - k_proj: [512, 1024] = [numKVHeads * headDim, hiddenSize]
+/// - v_proj: [512, 1024] = [numKVHeads * headDim, hiddenSize]
+/// - o_proj: [1024, 2048] = [hiddenSize, numQHeads * headDim]
+/// - q_norm: [256] = [headDim]
+/// - k_norm: [256] = [headDim]
+public final class GatedAttentionLayer: Module {
+ let numQHeads: Int
+ let numKVHeads: Int
+ let headDim: Int
+ let hiddenSize: Int
+ let scale: Float
+ let ropeDims: Int // partial RoPE dimensions
+
+ @ModuleInfo(key: "q_proj") var qProj: QuantizedLinear
+ @ModuleInfo(key: "k_proj") var kProj: QuantizedLinear
+ @ModuleInfo(key: "v_proj") var vProj: QuantizedLinear
+ @ModuleInfo(key: "o_proj") var oProj: QuantizedLinear
+
+ @ModuleInfo(key: "q_norm") var qNorm: RMSNorm
+ @ModuleInfo(key: "k_norm") var kNorm: RMSNorm
+
+ let rope: MLXNN.RoPE
+
+ public init(config: Qwen3ChatConfig) {
+ self.numQHeads = config.numAttentionHeads // 8
+ self.numKVHeads = config.numKeyValueHeads // 2
+ self.headDim = config.headDim // 256
+ self.hiddenSize = config.hiddenSize // 1024
+ self.scale = 1.0 / sqrt(Float(headDim))
+
+ let factor = config.partialRotaryFactor ?? 0.25
+ self.ropeDims = Int(Double(headDim) * factor) // 64
+
+ let groupSize = 64
+ let bits = 4
+ let qDim = numQHeads * headDim // 2048
+
+ // q_proj outputs 2 * qDim (Q + gate)
+ self._qProj = ModuleInfo(wrappedValue: QuantizedLinear(
+ hiddenSize, 2 * qDim, bias: false,
+ groupSize: groupSize, bits: bits))
+ self._kProj = ModuleInfo(wrappedValue: QuantizedLinear(
+ hiddenSize, numKVHeads * headDim, bias: false,
+ groupSize: groupSize, bits: bits))
+ self._vProj = ModuleInfo(wrappedValue: QuantizedLinear(
+ hiddenSize, numKVHeads * headDim, bias: false,
+ groupSize: groupSize, bits: bits))
+ // o_proj: qDim -> hiddenSize (after gating reduces 2*qDim to qDim)
+ self._oProj = ModuleInfo(wrappedValue: QuantizedLinear(
+ qDim, hiddenSize, bias: false,
+ groupSize: groupSize, bits: bits))
+
+ self._qNorm = ModuleInfo(wrappedValue: RMSNorm(dimensions: headDim, eps: Float(config.rmsNormEps)))
+ self._kNorm = ModuleInfo(wrappedValue: RMSNorm(dimensions: headDim, eps: Float(config.rmsNormEps)))
+
+ // Partial RoPE: only rotates first `ropeDims` of each head
+ self.rope = MLXNN.RoPE(
+ dimensions: ropeDims,
+ traditional: false,
+ base: Float(config.ropeTheta))
+
+ super.init()
+ }
+
+ /// Forward pass.
+ ///
+ /// - Parameters:
+ /// - hiddenStates: [B, T, hiddenSize]
+ /// - cache: Optional (keys, values) from previous steps, each [B, H_kv, S, D]
+ /// - offset: RoPE position offset (used when cache is nil, e.g. first call)
+ /// - Returns: (output [B, T, hiddenSize], updated KV cache)
+ public func callAsFunction(
+ _ hiddenStates: MLXArray,
+ cache: (MLXArray, MLXArray)? = nil,
+ offset: Int = 0
+ ) -> (MLXArray, (MLXArray, MLXArray)) {
+ let b = hiddenStates.dim(0)
+ let seqLen = hiddenStates.dim(1)
+ let qDim = numQHeads * headDim
+
+ // Q projection: [B, T, 2*qDim] → reshape to [B, T, H, 2*D] → split Q/gate INTERLEAVED per head
+ // CRITICAL: Must reshape BEFORE split (per Python reference).
+ // Interleaved format: for each head, first D dims are Q, next D are gate.
+ let qProjOut = qProj(hiddenStates) // [B, T, 4096 = 2*numQHeads*headDim]
+ let qProjReshaped = qProjOut.reshaped(b, seqLen, numQHeads, 2 * headDim)
+ let qgSplit = qProjReshaped.split(parts: 2, axis: -1)
+ var queries = qgSplit[0] // [B, T, H, D=256]
+ let gateSignal = qgSplit[1].reshaped(b, seqLen, qDim) // [B, T, 2048]
+
+ var keys = kProj(hiddenStates) // [B, T, numKVHeads * headDim]
+ var values = vProj(hiddenStates) // [B, T, numKVHeads * headDim]
+
+ // Reshape K/V to multi-head: [B, T, H, D]
+ keys = keys.reshaped(b, seqLen, numKVHeads, headDim)
+ values = values.reshaped(b, seqLen, numKVHeads, headDim)
+
+ // QK norm (per-head)
+ queries = qNorm(queries)
+ keys = kNorm(keys)
+
+ // Transpose to [B, H, T, D]
+ queries = queries.transposed(0, 2, 1, 3)
+ keys = keys.transposed(0, 2, 1, 3)
+ values = values.transposed(0, 2, 1, 3)
+
+ // Partial RoPE: MLXNN.RoPE with dimensions=ropeDims only rotates first ropeDims
+ let ropeOffset = cache?.0.dim(2) ?? offset
+ queries = rope(queries, offset: ropeOffset)
+ keys = rope(keys, offset: ropeOffset)
+
+ // Update KV cache
+ var cachedKeys = keys
+ var cachedValues = values
+ if let (prevK, prevV) = cache {
+ cachedKeys = concatenated([prevK, keys], axis: 2)
+ cachedValues = concatenated([prevV, values], axis: 2)
+ }
+
+ // Causal mask
+ let mask: MLXFast.ScaledDotProductAttentionMaskMode
+ if seqLen <= 1 && (cache != nil || offset > 0) {
+ mask = .none
+ } else {
+ let kvLen = cachedKeys.dim(2)
+ let pastLen = kvLen - seqLen
+ let causal = MLXArray.tri(seqLen, m: kvLen, k: pastLen, type: Float.self) - 1
+ let additiveMask = causal * Float.greatestFiniteMagnitude // 0 for attended, -FLT_MAX for masked
+ mask = .array(additiveMask.reshaped(1, 1, seqLen, kvLen).asType(queries.dtype))
+ }
+
+ // SDPA (handles GQA natively)
+ let attnOut = SDPA.attendAndMerge(
+ qHeads: queries, kHeads: cachedKeys, vHeads: cachedValues,
+ scale: scale, mask: mask)
+
+ // Gated output: attn_out * sigmoid(gate), then o_proj
+ // Reference: self.o_proj(output * mx.sigmoid(gate))
+ let gated = attnOut * sigmoid(gateSignal) // [B, T, qDim=2048]
+ let output = oProj(gated) // [B, T, hiddenSize=1024]
+
+ return (output, (cachedKeys, cachedValues))
+ }
+}
+
+// MARK: - Qwen3.5 Transformer Layer
+
+/// A single transformer layer in the Qwen3.5 hybrid model.
+///
+/// Either a DeltaNet (linear_attention) or GatedAttention (full_attention) layer,
+/// both sharing the same pre-norm structure and SwiGLU MLP.
+///
+/// The attention submodule is stored as the base `Module` type and registered
+/// under the key `"self_attn"` via `@ModuleInfo`. This ensures the MLX Module
+/// system discovers it for parameter traversal (`eval`, `clearParameters`, etc.)
+/// and weight loading maps to the correct key path.
+public final class Qwen35TransformerLayer: Module {
+ public let layerType: String
+
+ @ModuleInfo(key: "input_layernorm") var inputLayerNorm: RMSNorm
+ @ModuleInfo(key: "post_attention_layernorm") var postAttentionLayerNorm: RMSNorm
+ @ModuleInfo var mlp: Qwen35MLP
+
+ /// The attention submodule — either DeltaNetLayer or GatedAttentionLayer.
+ /// Key is "linear_attn" for DeltaNet, "self_attn" for GatedAttention (HuggingFace convention).
+ @ModuleInfo var attn: Module
+
+ public init(config: Qwen3ChatConfig, layerType: String) {
+ self.layerType = layerType
+
+ self._inputLayerNorm = ModuleInfo(wrappedValue: RMSNorm(dimensions: config.hiddenSize, eps: Float(config.rmsNormEps)))
+ self._postAttentionLayerNorm = ModuleInfo(wrappedValue: RMSNorm(dimensions: config.hiddenSize, eps: Float(config.rmsNormEps)))
+ self._mlp = ModuleInfo(
+ wrappedValue: Qwen35MLP(config: config))
+
+ if layerType == "linear_attention" {
+ self._attn = ModuleInfo(
+ wrappedValue: DeltaNetLayer(config: config),
+ key: "linear_attn")
+ } else {
+ self._attn = ModuleInfo(
+ wrappedValue: GatedAttentionLayer(config: config),
+ key: "self_attn")
+ }
+
+ super.init()
+ }
+
+ /// Access the DeltaNet submodule (only valid for linear_attention layers).
+ public var deltaNet: DeltaNetLayer? { attn as? DeltaNetLayer }
+
+ /// Access the GatedAttention submodule (only valid for full_attention layers).
+ public var gatedAttn: GatedAttentionLayer? { attn as? GatedAttentionLayer }
+
+ /// Forward for DeltaNet (linear attention) layer.
+ public func forwardDeltaNet(
+ _ x: MLXArray,
+ state: DeltaNetLayer.State?
+ ) -> (MLXArray, DeltaNetLayer.State) {
+ guard let dn = deltaNet else {
+ fatalError("forwardDeltaNet called on full_attention layer")
+ }
+ let normed = inputLayerNorm(x)
+ let (attnOut, newState) = dn(normed, state: state)
+ var h = x + attnOut
+ h = h + mlp(postAttentionLayerNorm(h))
+ return (h, newState)
+ }
+
+ /// Forward for GatedAttention (full attention) layer.
+ public func forwardGatedAttention(
+ _ x: MLXArray,
+ cache: (MLXArray, MLXArray)?,
+ offset: Int
+ ) -> (MLXArray, (MLXArray, MLXArray)) {
+ guard let ga = gatedAttn else {
+ fatalError("forwardGatedAttention called on linear_attention layer")
+ }
+ let normed = inputLayerNorm(x)
+ let (attnOut, newCache) = ga(normed, cache: cache, offset: offset)
+ var h = x + attnOut
+
+ h = h + mlp(postAttentionLayerNorm(h))
+ return (h, newCache)
+ }
+}
+
+// MARK: - SwiGLU MLP
+
+/// SwiGLU MLP for Qwen3.5 (quantized INT4).
+public final class Qwen35MLP: Module {
+ @ModuleInfo(key: "gate_proj") var gateProj: QuantizedLinear
+ @ModuleInfo(key: "up_proj") var upProj: QuantizedLinear
+ @ModuleInfo(key: "down_proj") var downProj: QuantizedLinear
+
+ public init(config: Qwen3ChatConfig) {
+ let hs = config.hiddenSize
+ let is_ = config.intermediateSize
+ let gs = 64, bits = 4
+
+ self._gateProj = ModuleInfo(
+ wrappedValue: QuantizedLinear(hs, is_, bias: false, groupSize: gs, bits: bits),
+ key: "gate_proj")
+ self._upProj = ModuleInfo(
+ wrappedValue: QuantizedLinear(hs, is_, bias: false, groupSize: gs, bits: bits),
+ key: "up_proj")
+ self._downProj = ModuleInfo(
+ wrappedValue: QuantizedLinear(is_, hs, bias: false, groupSize: gs, bits: bits),
+ key: "down_proj")
+
+ super.init()
+ }
+
+ public func callAsFunction(_ x: MLXArray) -> MLXArray {
+ downProj(silu(gateProj(x)) * upProj(x))
+ }
+}
+
+// MARK: - Qwen3.5 Full Model
+
+/// Qwen3.5-0.8B hybrid transformer with DeltaNet linear attention and GatedAttention.
+///
+/// Architecture: 24 layers in pattern [3x DeltaNet, 1x GatedAttention] x 6.
+/// - DeltaNet layers (18 of 24): O(1) memory per step via recurrent state, no KV cache.
+/// - GatedAttention layers (6 of 24): standard SDPA with KV cache, partial RoPE (25%).
+/// - Tied embeddings: lm_head reuses embed_tokens weights (PreQuantizedEmbedding.asLinear).
+///
+/// This gives a favorable memory/compute tradeoff: recurrent DeltaNet layers handle
+/// most computation with fixed memory, while sparse full attention layers provide
+/// global context at every 4th layer.
+public final class Qwen35MLXModel: Module {
+ public let config: Qwen3ChatConfig
+ public let layerTypes: [String]
+ public let fullAttentionIndices: [Int]
+
+ @ModuleInfo(key: "embed_tokens") var embedTokens: PreQuantizedEmbedding
+ @ModuleInfo var layers: [Qwen35TransformerLayer]
+ @ModuleInfo var norm: RMSNorm
+
+ public init(config: Qwen3ChatConfig) {
+ self.config = config
+
+ let types = config.layerTypes ?? Array(
+ repeating: "full_attention", count: config.numHiddenLayers)
+ self.layerTypes = types
+ self.fullAttentionIndices = types.enumerated().compactMap {
+ $0.element == "full_attention" ? $0.offset : nil
+ }
+
+ self._embedTokens = ModuleInfo(wrappedValue: PreQuantizedEmbedding(
+ embeddingCount: config.vocabSize,
+ dimensions: config.hiddenSize,
+ groupSize: 64, bits: 4))
+
+ self._layers = ModuleInfo(
+ wrappedValue: types.map { Qwen35TransformerLayer(config: config, layerType: $0) })
+
+ self._norm = ModuleInfo(
+ wrappedValue: RMSNorm(dimensions: config.hiddenSize, eps: Float(config.rmsNormEps)))
+
+ super.init()
+ }
+
+ // MARK: - Inference State
+
+ /// Combined inference state: DeltaNet recurrent states + GatedAttention KV caches.
+ public struct InferenceState {
+ /// Per-layer DeltaNet state (nil for full_attention layers).
+ public var deltaNetStates: [DeltaNetLayer.State?]
+ /// Per-layer KV cache (nil for linear_attention layers, and for full_attention
+ /// layers before any tokens have been processed).
+ public var kvCaches: [(MLXArray, MLXArray)?]
+ /// Current sequence position (for RoPE offset in GatedAttention layers).
+ public var position: Int
+
+ public static func initial(config: Qwen3ChatConfig, batchSize: Int = 1) -> InferenceState {
+ let types = config.layerTypes ?? Array(
+ repeating: "full_attention", count: config.numHiddenLayers)
+ let numHeads = config.linearNumKeyHeads ?? 16
+ let headDim = config.linearKeyHeadDim ?? 128
+ let qkvDim = 3 * numHeads * headDim
+ let convKernel = config.linearConvKernelDim ?? 4
+
+ return InferenceState(
+ deltaNetStates: types.map { type in
+ type == "linear_attention"
+ ? DeltaNetLayer.State.initial(
+ batchSize: batchSize, numHeads: numHeads, headDim: headDim,
+ qkvDim: qkvDim, convKernel: convKernel)
+ : nil
+ },
+ kvCaches: types.map { _ in nil },
+ position: 0
+ )
+ }
+ }
+
+ // MARK: - Forward Pass
+
+ /// Forward pass through the full model.
+ ///
+ /// - Parameters:
+ /// - inputIds: Token IDs [B, T]
+ /// - state: Inference state
+ /// - Returns: (logits [B, T, vocabSize], updated state)
+ public func forward(
+ inputIds: MLXArray,
+ state: InferenceState
+ ) -> (MLXArray, InferenceState) {
+ let seqLen = inputIds.dim(1)
+ var hidden = embedTokens(inputIds) // [B, T, hiddenSize]
+
+ var newDeltaStates = state.deltaNetStates
+ var newKVCaches = state.kvCaches
+
+ for (i, layer) in layers.enumerated() {
+ if layerTypes[i] == "linear_attention" {
+ let (h, newState) = layer.forwardDeltaNet(hidden, state: state.deltaNetStates[i])
+ hidden = h
+ newDeltaStates[i] = newState
+ } else {
+ let (h, newCache) = layer.forwardGatedAttention(
+ hidden, cache: state.kvCaches[i], offset: state.position)
+ hidden = h
+ newKVCaches[i] = newCache
+ }
+ }
+
+ hidden = norm(hidden)
+
+ // Tied LM head
+ let logits = embedTokens.asLinear(hidden)
+
+ let newState = InferenceState(
+ deltaNetStates: newDeltaStates,
+ kvCaches: newKVCaches,
+ position: state.position + seqLen)
+
+ return (logits, newState)
+ }
+
+ // MARK: - Text Generation
+
+ /// Generate text tokens autoregressively.
+ ///
+ /// - Parameters:
+ /// - promptIds: Prompt token IDs
+ /// - sampling: Sampling configuration
+ /// - Returns: Generated token IDs (excluding prompt)
+ public func generate(
+ promptIds: [Int],
+ sampling: ChatSamplingConfig = .default
+ ) -> [Int] {
+ var state = InferenceState.initial(config: config)
+
+ // Prefill
+ let prompt = MLXArray(promptIds.map { Int32($0) }).expandedDimensions(axis: 0)
+ let (prefillLogits, prefillState) = forward(inputIds: prompt, state: state)
+ state = prefillState
+ eval(prefillLogits)
+
+ // Sample first token
+ var token = sampleFromLogits(prefillLogits, at: promptIds.count - 1,
+ config: sampling, history: promptIds)
+ if token == config.eosTokenId { return [] }
+
+ var generated = [token]
+
+ // Decode loop
+ for _ in 1.. Int {
+ let posLogits = logits[0, position] // [vocabSize]
+ let f32 = posLogits.asType(.float32)
+ eval(f32)
+ let count = self.config.vocabSize
+ let floats: [Float] = f32.asArray(Float.self)
+ return ChatSampler.sample(logits: Array(floats.prefix(count)), config: config, previousTokens: history)
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen35PipelineLLM.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen35PipelineLLM.swift
new file mode 100644
index 0000000..d64b1f8
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen35PipelineLLM.swift
@@ -0,0 +1,79 @@
+import Foundation
+import AudioCommon
+
+/// Common interface for Qwen3.5 chat backends (MLX or CoreML).
+public protocol Qwen35ChatBackend: AnyObject {
+ var tokenizer: ChatTokenizer { get }
+ var config: Qwen3ChatConfig { get }
+ func generateStream(messages: [ChatMessage], sampling: ChatSamplingConfig)
+ -> AsyncThrowingStream
+ func resetState()
+}
+
+extension Qwen35MLXChat: Qwen35ChatBackend {}
+extension Qwen35CoreMLChat: Qwen35ChatBackend {}
+
+/// Bridges any Qwen3.5 backend to VoicePipeline's PipelineLLM protocol.
+public final class Qwen35PipelineLLM: PipelineLLM {
+ private let model: any Qwen35ChatBackend
+ private let systemPrompt: String
+ private let sampling: ChatSamplingConfig
+ private var cancelled = false
+
+ public var onToken: ((String) -> Void)?
+
+ public init(
+ model: any Qwen35ChatBackend,
+ systemPrompt: String = "Your name is Tama. Give short direct answers. Do not explain your reasoning.",
+ sampling: ChatSamplingConfig = .default
+ ) {
+ self.model = model
+ self.systemPrompt = systemPrompt
+ self.sampling = sampling
+ }
+
+ public func chat(
+ messages: [(role: MessageRole, content: String)],
+ onToken: @escaping (String, Bool) -> Void
+ ) {
+ cancelled = false
+
+ let chatMessages = messages.compactMap { msg -> ChatMessage? in
+ switch msg.role {
+ case .system: return ChatMessage(role: .system, content: msg.content)
+ case .user: return ChatMessage(role: .user, content: msg.content)
+ case .assistant: return ChatMessage(role: .assistant, content: msg.content)
+ default: return nil
+ }
+ }
+
+ var fullMessages = [ChatMessage(role: .system, content: systemPrompt)]
+ fullMessages.append(contentsOf: chatMessages)
+
+ let stream = model.generateStream(messages: fullMessages, sampling: sampling)
+ let semaphore = DispatchSemaphore(value: 0)
+ var fullResponse = ""
+
+ Task {
+ do {
+ for try await chunk in stream {
+ guard !self.cancelled else { break }
+ fullResponse += chunk
+ self.onToken?(chunk)
+ onToken(chunk, false)
+ }
+ } catch { }
+
+ if !fullResponse.isEmpty {
+ onToken("", true)
+ }
+ semaphore.signal()
+ }
+
+ semaphore.wait()
+ }
+
+ public func cancel() {
+ cancelled = true
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen35WeightLoading.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen35WeightLoading.swift
new file mode 100644
index 0000000..16bc565
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen35WeightLoading.swift
@@ -0,0 +1,226 @@
+import Foundation
+import MLXCommon
+import MLX
+import MLXNN
+
+// MARK: - Weight Loading for Qwen3.5-0.8B MLX Model
+
+/// Loads quantized safetensors weights into the Qwen3.5 MLX model.
+///
+/// Expected weight key structure (HuggingFace / mlx-community format):
+///
+/// Keys may have `model.` or `language_model.model.` prefix — both are stripped.
+///
+/// - `embed_tokens.*` -> embed_tokens (PreQuantizedEmbedding)
+/// - `layers.{i}.linear_attn.*` -> DeltaNet (linear_attention layers)
+/// - `layers.{i}.self_attn.*` -> GatedAttention (full_attention layers)
+/// - `layers.{i}.mlp.*` -> SwiGLU MLP
+/// - `layers.{i}.input_layernorm.*`
+/// - `layers.{i}.post_attention_layernorm.*`
+/// - `norm.*` -> final RMSNorm
+///
+/// DeltaNet (linear_attention) weights under `linear_attn.`:
+/// - `in_proj_qkv.{weight,scales,biases}`: quantized [6144, 1024]
+/// - `in_proj_z.{weight,scales,biases}`: quantized [2048, 1024]
+/// - `in_proj_b.{weight,scales,biases}`: quantized [16, 1024]
+/// - `in_proj_a.{weight,scales,biases}`: quantized [16, 1024]
+/// - `conv1d.weight`: [6144, 1, 4]
+/// - `dt_bias`: [16]
+/// - `A_log`: [16]
+/// - `norm.weight`: [128]
+/// - `out_proj.{weight,scales,biases}`: quantized [1024, 2048]
+///
+/// GatedAttention (full_attention) weights under `self_attn.`:
+/// - `q_proj.{weight,scales,biases}`: quantized [4096, 1024]
+/// - `k_proj.{weight,scales,biases}`: quantized [512, 1024]
+/// - `v_proj.{weight,scales,biases}`: quantized [512, 1024]
+/// - `o_proj.{weight,scales,biases}`: quantized [1024, 2048]
+/// - `q_norm.weight`: [256]
+/// - `k_norm.weight`: [256]
+///
+/// MLP weights under `mlp.`:
+/// - `gate_proj.{weight,scales,biases}`: quantized [3584, 1024]
+/// - `up_proj.{weight,scales,biases}`: quantized [3584, 1024]
+/// - `down_proj.{weight,scales,biases}`: quantized [1024, 3584]
+public enum Qwen35WeightLoader {
+
+ /// Load weights from a directory containing safetensors files.
+ ///
+ /// - Parameters:
+ /// - model: The Qwen3.5 MLX model to load weights into
+ /// - directory: Directory containing safetensors files
+ /// - progressHandler: Optional progress callback
+ public static func loadWeights(
+ into model: Qwen35MLXModel,
+ from directory: URL,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) throws {
+ progressHandler?(0.05, "Loading weight files...")
+
+ // Load all safetensors files from the directory
+ let allWeights = try CommonWeightLoader.loadAllSafetensors(from: directory)
+ progressHandler?(0.3, "Loaded \(allWeights.count) tensors")
+
+ // Strip prefix from keys. Handles two formats:
+ // - Our format: "model.layers.0.*"
+ // - mlx-community VLM: "language_model.model.layers.0.*" (also has vision_tower.* which we skip)
+ var modelWeights: [String: MLXArray] = [:]
+ for (key, value) in allWeights {
+ if key.hasPrefix("language_model.model.") {
+ modelWeights[String(key.dropFirst("language_model.model.".count))] = value
+ } else if key.hasPrefix("model.") {
+ modelWeights[String(key.dropFirst("model.".count))] = value
+ } else if key.hasPrefix("lm_head.") || key.hasPrefix("vision_tower.") {
+ // Skip — lm_head is tied to embed_tokens, vision_tower not needed
+ continue
+ } else {
+ modelWeights[key] = value
+ }
+ }
+
+ progressHandler?(0.4, "Applying embedding weights...")
+
+ // Load embed_tokens (PreQuantizedEmbedding)
+ CommonWeightLoader.applyQuantizedEmbeddingWeights(
+ to: model.embedTokens,
+ prefix: "embed_tokens",
+ from: modelWeights)
+
+ // Load final norm
+ CommonWeightLoader.applyRMSNormWeights(
+ to: model.norm, prefix: "norm", from: modelWeights)
+
+ progressHandler?(0.5, "Loading transformer layers...")
+
+ // Load each layer
+ let numLayers = model.config.numHiddenLayers
+ for i in 0..] = [:]
+ if let w = weights["\(prefix).conv1d.weight"] {
+ rawParams["convWeight"] = .value(w)
+ }
+ if let dtb = weights["\(prefix).dt_bias"] {
+ rawParams["dtBias"] = .value(dtb)
+ }
+ if let alog = weights["\(prefix).A_log"] {
+ rawParams["aLog"] = .value(alog)
+ }
+ if !rawParams.isEmpty {
+ layer.update(parameters: ModuleParameters(values: rawParams))
+ }
+
+ // Per-head norm
+ CommonWeightLoader.applyRMSNormWeights(
+ to: layer.norm, prefix: "\(prefix).norm", from: weights)
+
+ // Output projection (quantized)
+ CommonWeightLoader.applyQuantizedLinearWeights(
+ to: layer.outProj, prefix: "\(prefix).out_proj", from: weights)
+ }
+
+ // MARK: - GatedAttention Weight Loading
+
+ /// Apply weights to a GatedAttention (full attention) layer.
+ ///
+ /// All projections are quantized (INT4 with group_size=64).
+ private static func applyGatedAttentionWeights(
+ to layer: GatedAttentionLayer,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ CommonWeightLoader.applyQuantizedLinearWeights(
+ to: layer.qProj, prefix: "\(prefix).q_proj", from: weights)
+ CommonWeightLoader.applyQuantizedLinearWeights(
+ to: layer.kProj, prefix: "\(prefix).k_proj", from: weights)
+ CommonWeightLoader.applyQuantizedLinearWeights(
+ to: layer.vProj, prefix: "\(prefix).v_proj", from: weights)
+ CommonWeightLoader.applyQuantizedLinearWeights(
+ to: layer.oProj, prefix: "\(prefix).o_proj", from: weights)
+
+ CommonWeightLoader.applyRMSNormWeights(
+ to: layer.qNorm, prefix: "\(prefix).q_norm", from: weights)
+ CommonWeightLoader.applyRMSNormWeights(
+ to: layer.kNorm, prefix: "\(prefix).k_norm", from: weights)
+ }
+
+ // MARK: - MLP Weight Loading
+
+ /// Apply quantized MLP weights (SwiGLU: gate_proj, up_proj, down_proj).
+ private static func applyQuantizedMLPWeights(
+ to mlp: Qwen35MLP,
+ prefix: String,
+ from weights: [String: MLXArray]
+ ) {
+ CommonWeightLoader.applyQuantizedLinearWeights(
+ to: mlp.gateProj, prefix: "\(prefix).gate_proj", from: weights)
+ CommonWeightLoader.applyQuantizedLinearWeights(
+ to: mlp.upProj, prefix: "\(prefix).up_proj", from: weights)
+ CommonWeightLoader.applyQuantizedLinearWeights(
+ to: mlp.downProj, prefix: "\(prefix).down_proj", from: weights)
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen3ChatConfig.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen3ChatConfig.swift
new file mode 100644
index 0000000..93cd466
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen3ChatConfig.swift
@@ -0,0 +1,146 @@
+import Foundation
+
+/// Model architecture type.
+public enum ChatModelArch: String, Codable, Sendable {
+ /// Qwen3.5 hybrid (DeltaNet linear attention + GatedAttention)
+ case qwen35 = "qwen3_5_text"
+}
+
+/// Configuration for Qwen3.5 chat model.
+public struct Qwen3ChatConfig: Codable, Sendable {
+ public let hiddenSize: Int
+ public let numHiddenLayers: Int
+ public let numAttentionHeads: Int
+ public let numKeyValueHeads: Int
+ public let headDim: Int
+ public let intermediateSize: Int
+ public let vocabSize: Int
+ public let maxSeqLen: Int
+ public let ropeTheta: Double
+ public let rmsNormEps: Double
+ public let eosTokenId: Int
+ public let padTokenId: Int
+ public let quantization: String
+
+ // Qwen3.5-specific fields
+ public let modelType: ChatModelArch?
+ /// Per-layer type: "linear_attention" (DeltaNet) or "full_attention" (GatedAttention)
+ public let layerTypes: [String]?
+ /// How often a full_attention layer appears (e.g., 4 = every 4th layer)
+ public let fullAttentionInterval: Int?
+ /// DeltaNet linear attention head config
+ public let linearNumKeyHeads: Int?
+ public let linearKeyHeadDim: Int?
+ public let linearNumValueHeads: Int?
+ public let linearValueHeadDim: Int?
+ /// Causal conv1d kernel size for DeltaNet
+ public let linearConvKernelDim: Int?
+ /// Partial RoPE factor for GatedAttention (e.g., 0.25)
+ public let partialRotaryFactor: Double?
+ /// Whether embeddings are tied (lm_head = embed_tokens)
+ public let tieWordEmbeddings: Bool?
+
+ enum CodingKeys: String, CodingKey {
+ case hiddenSize = "hidden_size"
+ case numHiddenLayers = "num_hidden_layers"
+ case numAttentionHeads = "num_attention_heads"
+ case numKeyValueHeads = "num_key_value_heads"
+ case headDim = "head_dim"
+ case intermediateSize = "intermediate_size"
+ case vocabSize = "vocab_size"
+ case maxSeqLen = "max_seq_len"
+ case ropeTheta = "rope_theta"
+ case rmsNormEps = "rms_norm_eps"
+ case eosTokenId = "eos_token_id"
+ case padTokenId = "pad_token_id"
+ case quantization
+ case modelType = "model_type"
+ case layerTypes = "layer_types"
+ case fullAttentionInterval = "full_attention_interval"
+ case linearNumKeyHeads = "linear_num_key_heads"
+ case linearKeyHeadDim = "linear_key_head_dim"
+ case linearNumValueHeads = "linear_num_value_heads"
+ case linearValueHeadDim = "linear_value_head_dim"
+ case linearConvKernelDim = "linear_conv_kernel_dim"
+ case partialRotaryFactor = "partial_rotary_factor"
+ case tieWordEmbeddings = "tie_word_embeddings"
+ }
+
+ /// Whether this is a Qwen3.5 hybrid model.
+ public var isQwen35: Bool {
+ modelType == .qwen35 || layerTypes != nil
+ }
+
+ /// Number of full-attention layers (that need KV cache).
+ public var numFullAttentionLayers: Int {
+ guard let types = layerTypes else { return numHiddenLayers }
+ return types.filter { $0 == "full_attention" }.count
+ }
+
+ /// Default config for Qwen3.5-0.8B.
+ public static let qwen35_08B = Qwen3ChatConfig(
+ hiddenSize: 1024,
+ numHiddenLayers: 24,
+ numAttentionHeads: 8,
+ numKeyValueHeads: 2,
+ headDim: 256,
+ intermediateSize: 3584,
+ vocabSize: 248320,
+ maxSeqLen: 2048,
+ ropeTheta: 10_000_000.0,
+ rmsNormEps: 1e-6,
+ eosTokenId: 248046, // <|im_end|> — stops generation at end of assistant turn
+ padTokenId: 248044, // <|endoftext|>
+ quantization: "int4",
+ modelType: .qwen35,
+ layerTypes: [
+ "linear_attention", "linear_attention", "linear_attention", "full_attention",
+ "linear_attention", "linear_attention", "linear_attention", "full_attention",
+ "linear_attention", "linear_attention", "linear_attention", "full_attention",
+ "linear_attention", "linear_attention", "linear_attention", "full_attention",
+ "linear_attention", "linear_attention", "linear_attention", "full_attention",
+ "linear_attention", "linear_attention", "linear_attention", "full_attention",
+ ],
+ fullAttentionInterval: 4,
+ linearNumKeyHeads: 16,
+ linearKeyHeadDim: 128,
+ linearNumValueHeads: 16,
+ linearValueHeadDim: 128,
+ linearConvKernelDim: 4,
+ partialRotaryFactor: 0.25,
+ tieWordEmbeddings: true
+ )
+
+ /// Load config from a JSON file.
+ public static func load(from url: URL) throws -> Qwen3ChatConfig {
+ let data = try Data(contentsOf: url)
+ return try JSONDecoder().decode(Qwen3ChatConfig.self, from: data)
+ }
+}
+
+/// Sampling parameters for text generation.
+public struct ChatSamplingConfig: Sendable {
+ public var temperature: Float
+ public var topK: Int
+ public var topP: Float
+ public var maxTokens: Int
+ public var repetitionPenalty: Float
+
+ public init(
+ temperature: Float = 0.7,
+ topK: Int = 50,
+ topP: Float = 0.9,
+ maxTokens: Int = 256,
+ repetitionPenalty: Float = 1.1
+ ) {
+ self.temperature = temperature
+ self.topK = topK
+ self.topP = topP
+ self.maxTokens = maxTokens
+ self.repetitionPenalty = repetitionPenalty
+ }
+
+ public static let `default` = ChatSamplingConfig()
+ public static let creative = ChatSamplingConfig(temperature: 0.9, topP: 0.95)
+ public static let precise = ChatSamplingConfig(temperature: 0.3, topK: 20, topP: 0.8)
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen3ChatError.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen3ChatError.swift
new file mode 100644
index 0000000..4a725e6
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen3ChatError.swift
@@ -0,0 +1,25 @@
+import Foundation
+
+/// Errors for Qwen3.5 chat model operations.
+public enum ChatModelError: LocalizedError {
+ case modelLoadFailed(String)
+ case tokenizerLoadFailed(String)
+ case inferenceFailed(String)
+ case configNotFound(URL)
+ case modelNotFound(URL)
+
+ public var errorDescription: String? {
+ switch self {
+ case .modelLoadFailed(let reason):
+ "Failed to load chat model: \(reason)"
+ case .tokenizerLoadFailed(let reason):
+ "Failed to load tokenizer: \(reason)"
+ case .inferenceFailed(let reason):
+ "Inference failed: \(reason)"
+ case .configNotFound(let url):
+ "Config not found at \(url.path)"
+ case .modelNotFound(let url):
+ "Model not found at \(url.path)"
+ }
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/BiLSTM.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/BiLSTM.swift
new file mode 100644
index 0000000..d95ac26
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/BiLSTM.swift
@@ -0,0 +1,100 @@
+import Foundation
+import MLX
+import MLXNN
+
+/// A single LSTM layer with updatable parameters.
+///
+/// Uses `var` properties so they can be loaded via `update(parameters:)`.
+/// Parameter keys match MLX convention: `Wx` (input→hidden), `Wh` (hidden→hidden), `bias`.
+class LSTMLayer: Module {
+ let hiddenSize: Int
+
+ @ParameterInfo(key: "Wx") var wx: MLXArray
+ @ParameterInfo(key: "Wh") var wh: MLXArray
+ var bias: MLXArray
+
+ init(inputSize: Int, hiddenSize: Int) {
+ self.hiddenSize = hiddenSize
+ let scale = 1.0 / Foundation.sqrt(Float(hiddenSize))
+ self._wx.wrappedValue = MLXRandom.uniform(low: -scale, high: scale, [4 * hiddenSize, inputSize])
+ self._wh.wrappedValue = MLXRandom.uniform(low: -scale, high: scale, [4 * hiddenSize, hiddenSize])
+ self.bias = MLXRandom.uniform(low: -scale, high: scale, [4 * hiddenSize])
+ }
+
+ /// Process a sequence.
+ /// - Parameter x: `[batch, seq_len, input_size]`
+ /// - Returns: `[batch, seq_len, hidden_size]`
+ func callAsFunction(_ x: MLXArray) -> MLXArray {
+ // Project all timesteps at once: [B, L, 4H]
+ let projected = addMM(bias, x, wx.T)
+
+ let seqLen = x.dim(-2)
+ var hidden: MLXArray? = nil
+ var cell: MLXArray? = nil
+ var allHidden = [MLXArray]()
+
+ for t in 0 ..< seqLen {
+ var ifgo = projected[.ellipsis, t, 0...]
+ if let h = hidden {
+ ifgo = ifgo + matmul(h, wh.T)
+ }
+
+ let pieces = split(ifgo, parts: 4, axis: -1)
+ let i = sigmoid(pieces[0])
+ let f = sigmoid(pieces[1])
+ let g = tanh(pieces[2])
+ let o = sigmoid(pieces[3])
+
+ if let c = cell {
+ cell = f * c + i * g
+ } else {
+ cell = i * g
+ }
+ hidden = o * tanh(cell!)
+
+ allHidden.append(hidden!)
+ }
+
+ return stacked(allHidden, axis: -2)
+ }
+}
+
+/// Container for LSTM layers (enables module parameter tree: `layers.0`, `layers.1`, etc.)
+class LSTMStack: Module {
+ let layers: [LSTMLayer]
+
+ init(layers: [LSTMLayer]) {
+ self.layers = layers
+ }
+}
+
+/// Run bidirectional LSTM across multiple layers.
+///
+/// For each layer, runs forward LSTM on the sequence and backward LSTM on the
+/// reversed sequence, then concatenates outputs along the feature dimension.
+///
+/// - Parameters:
+/// - x: `[batch, seq_len, features]`
+/// - fwd: forward LSTM stack
+/// - bwd: backward LSTM stack
+/// - Returns: `[batch, seq_len, 2 * hidden_size]`
+func runBiLSTM(_ x: MLXArray, fwd: LSTMStack, bwd: LSTMStack) -> MLXArray {
+ var input = x
+
+ for i in 0 ..< fwd.layers.count {
+ // Forward direction
+ let fwdOut = fwd.layers[i](input)
+
+ // Backward direction: reverse → process → reverse back
+ let seqLen = input.dim(-2)
+ let indices = MLXArray(Array((0 ..< seqLen).reversed()))
+ let reversed = input.take(indices, axis: -2)
+ let bwdOutRev = bwd.layers[i](reversed)
+ let bwdOut = bwdOutRev.take(indices, axis: -2)
+
+ // Concatenate along feature dimension
+ input = concatenated([fwdOut, bwdOut], axis: -1)
+ }
+
+ return input
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/Configuration.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/Configuration.swift
new file mode 100644
index 0000000..6f430d7
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/Configuration.swift
@@ -0,0 +1,92 @@
+import Foundation
+
+/// Model configuration for pyannote PyanNet segmentation.
+public struct SegmentationConfig: Sendable {
+ /// Audio sample rate in Hz
+ public let sampleRate: Int
+
+ /// SincNet filter counts per layer
+ public let sincnetFilters: [Int]
+ /// SincNet kernel sizes per layer
+ public let sincnetKernelSizes: [Int]
+ /// SincNet strides per layer
+ public let sincnetStrides: [Int]
+ /// SincNet max-pool kernel sizes per layer
+ public let sincnetPoolSizes: [Int]
+
+ /// LSTM hidden size (per direction; output is 2x for bidirectional)
+ public let lstmHiddenSize: Int
+ /// Number of LSTM layers
+ public let lstmNumLayers: Int
+
+ /// Linear layer hidden size
+ public let linearHiddenSize: Int
+ /// Number of linear layers (before classifier)
+ public let linearNumLayers: Int
+
+ /// Number of output classes (powerset)
+ public let numClasses: Int
+
+ /// Default configuration for pyannote/segmentation-3.0
+ public static let `default` = SegmentationConfig(
+ sampleRate: 16000,
+ sincnetFilters: [80, 60, 60],
+ sincnetKernelSizes: [251, 5, 5],
+ sincnetStrides: [10, 1, 1],
+ sincnetPoolSizes: [3, 3, 3],
+ lstmHiddenSize: 128,
+ lstmNumLayers: 4,
+ linearHiddenSize: 128,
+ linearNumLayers: 2,
+ numClasses: 7
+ )
+}
+
+/// VAD pipeline configuration with default hysteresis thresholds.
+public struct VADConfig: Sendable {
+ /// Onset threshold (speech starts when probability exceeds this)
+ public var onset: Float
+ /// Offset threshold (speech ends when probability drops below this)
+ public var offset: Float
+ /// Minimum speech duration in seconds
+ public var minSpeechDuration: Float
+ /// Minimum silence duration in seconds
+ public var minSilenceDuration: Float
+ /// Analysis window duration in seconds
+ public var windowDuration: Float
+ /// Step ratio for sliding window (fraction of window)
+ public var stepRatio: Float
+
+ public init(
+ onset: Float, offset: Float,
+ minSpeechDuration: Float, minSilenceDuration: Float,
+ windowDuration: Float, stepRatio: Float
+ ) {
+ self.onset = onset
+ self.offset = offset
+ self.minSpeechDuration = minSpeechDuration
+ self.minSilenceDuration = minSilenceDuration
+ self.windowDuration = windowDuration
+ self.stepRatio = stepRatio
+ }
+
+ /// Default pyannote VAD thresholds
+ public static let `default` = VADConfig(
+ onset: 0.767,
+ offset: 0.377,
+ minSpeechDuration: 0.136,
+ minSilenceDuration: 0.067,
+ windowDuration: 10.0,
+ stepRatio: 0.1
+ )
+
+ /// Default Silero VAD thresholds (streaming-optimized)
+ public static let sileroDefault = VADConfig(
+ onset: 0.5,
+ offset: 0.35,
+ minSpeechDuration: 0.25,
+ minSilenceDuration: 0.1,
+ windowDuration: 0.032,
+ stepRatio: 1.0
+ )
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/CoreMLSileroInference.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/CoreMLSileroInference.swift
new file mode 100644
index 0000000..6740659
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/CoreMLSileroInference.swift
@@ -0,0 +1,64 @@
+#if canImport(CoreML)
+import AudioCommon
+import CoreML
+import Foundation
+
+extension SileroVADModel {
+
+ /// Run CoreML inference for one 576-sample chunk (64 context + 512 new).
+ ///
+ /// Creates input MLMultiArrays in float16, runs prediction, and updates
+ /// the internal LSTM h/c state for the next chunk.
+ ///
+ /// - Parameter fullSamples: 576 Float32 samples (context prepended)
+ /// - Returns: speech probability in `[0, 1]`
+ func processChunkCoreML(_ fullSamples: [Float]) throws -> Float {
+ guard let model = coremlModel else {
+ throw AudioModelError.inferenceFailed(
+ operation: "VAD", reason: "CoreML model not loaded")
+ }
+
+ // Create audio input: [1, 1, 576] float16
+ let audioArray = try MLMultiArray(shape: [1, 1, 576], dataType: .float16)
+ let audioPtr = audioArray.dataPointer.assumingMemoryBound(to: Float16.self)
+ for i in 0..<576 {
+ audioPtr[i] = Float16(fullSamples[i])
+ }
+
+ // Initialize h/c to zeros on first call
+ if coremlH == nil {
+ coremlH = try MLMultiArray(shape: [1, 1, 128], dataType: .float16)
+ coremlC = try MLMultiArray(shape: [1, 1, 128], dataType: .float16)
+ zeroFillFloat16(coremlH!)
+ zeroFillFloat16(coremlC!)
+ }
+
+ let input = try MLDictionaryFeatureProvider(dictionary: [
+ "audio": MLFeatureValue(multiArray: audioArray),
+ "h": MLFeatureValue(multiArray: coremlH!),
+ "c": MLFeatureValue(multiArray: coremlC!),
+ ])
+
+ let result = try model.prediction(from: input)
+
+ // Update LSTM state
+ coremlH = result.featureValue(for: "h_out")!.multiArrayValue!
+ coremlC = result.featureValue(for: "c_out")!.multiArrayValue!
+
+ // Extract probability scalar
+ let probArray = result.featureValue(for: "probability")!.multiArrayValue!
+ let probPtr = probArray.dataPointer.assumingMemoryBound(to: Float16.self)
+ return Float(probPtr[0])
+ }
+
+ /// Zero-fill a float16 MLMultiArray.
+ private func zeroFillFloat16(_ array: MLMultiArray) {
+ let ptr = UnsafeMutableBufferPointer(
+ start: array.dataPointer.assumingMemoryBound(to: Float16.self),
+ count: array.count)
+ for i in 0.. [Float] {
+ guard let model = coremlModel else {
+ throw AudioModelError.inferenceFailed(
+ operation: "SpeakerEmbedding", reason: "CoreML model not loaded")
+ }
+
+ // Find nearest enumerated length >= nFrames
+ let targetLength = Self.enumeratedMelLengths.first { $0 >= nFrames }
+ ?? Self.enumeratedMelLengths.last!
+
+ // Create input: [1, targetLength, 80] float16
+ // The CoreML model internally permutes (T,80) → (80,T) to match
+ // the trained weight orientation (freq as height, time as width).
+ let melArray = try MLMultiArray(
+ shape: [1, targetLength as NSNumber, 80],
+ dataType: .float16
+ )
+ let melPtr = melArray.dataPointer.assumingMemoryBound(to: Float16.self)
+
+ // Fill with mel data (row-major: frame-major, 80 mels per frame)
+ let copyCount = min(nFrames, targetLength) * 80
+ for i in 0.. 1e-10 {
+ for i in 0..<256 { embedding[i] /= norm }
+ }
+
+ return embedding
+ }
+}
+#endif
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/DERScoring.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/DERScoring.swift
new file mode 100644
index 0000000..baa36eb
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/DERScoring.swift
@@ -0,0 +1,408 @@
+import Foundation
+import AudioCommon
+
+// MARK: - RTTM Format
+
+/// RTTM (Rich Transcription Time Marked) segment for standard diarization evaluation.
+public struct RTTMSegment: Sendable {
+ public let filename: String
+ public let startTime: Float
+ public let duration: Float
+ public let speakerLabel: String
+
+ public init(filename: String, startTime: Float, duration: Float, speakerLabel: String) {
+ self.filename = filename
+ self.startTime = startTime
+ self.duration = duration
+ self.speakerLabel = speakerLabel
+ }
+
+ /// Format as standard RTTM line: `SPEAKER 1 `
+ public var rttmLine: String {
+ let s = String(format: "%.3f", startTime)
+ let d = String(format: "%.3f", duration)
+ return "SPEAKER \(filename) 1 \(s) \(d) \(speakerLabel) "
+ }
+}
+
+/// Convert diarization result to RTTM format.
+public func toRTTM(segments: [DiarizedSegment], filename: String) -> [RTTMSegment] {
+ segments.map { seg in
+ RTTMSegment(
+ filename: filename,
+ startTime: seg.startTime,
+ duration: seg.duration,
+ speakerLabel: "speaker_\(seg.speakerId)"
+ )
+ }
+}
+
+/// Write RTTM segments to string.
+public func formatRTTM(_ rttmSegments: [RTTMSegment]) -> String {
+ rttmSegments.map(\.rttmLine).joined(separator: "\n")
+}
+
+// MARK: - DER Computation
+
+/// Diarization Error Rate result.
+public struct DERResult: Sendable {
+ /// Total scored speech duration in seconds
+ public let totalSpeech: Float
+ /// False alarm duration (non-speech classified as speech)
+ public let falseAlarm: Float
+ /// Missed speech duration
+ public let missedSpeech: Float
+ /// Speaker confusion duration (wrong speaker assigned)
+ public let confusion: Float
+
+ /// Diarization Error Rate = (FA + Miss + Confusion) / TotalSpeech
+ public var der: Float {
+ guard totalSpeech > 0 else { return 0 }
+ return (falseAlarm + missedSpeech + confusion) / totalSpeech
+ }
+
+ /// Diarization Error Rate as percentage
+ public var derPercent: Float { der * 100 }
+}
+
+/// Compute Diarization Error Rate between reference and hypothesis.
+///
+/// Uses frame-level scoring with configurable resolution and collar.
+/// Collar applies forgiveness around reference segment boundaries.
+///
+/// - Parameters:
+/// - reference: reference (ground truth) segments
+/// - hypothesis: hypothesis (system output) segments
+/// - collar: forgiveness collar in seconds around boundaries (default 0.25s)
+/// - resolution: scoring resolution in seconds (default 0.01s = 10ms)
+/// - Returns: DER breakdown
+public func computeDER(
+ reference: [DiarizedSegment],
+ hypothesis: [DiarizedSegment],
+ collar: Float = 0.25,
+ resolution: Float = 0.01
+) -> DERResult {
+ guard !reference.isEmpty else {
+ let hTotal = hypothesis.reduce(Float(0)) { $0 + $1.duration }
+ return DERResult(totalSpeech: 0, falseAlarm: hTotal, missedSpeech: 0, confusion: 0)
+ }
+
+ // Find time range
+ let allSegments: [DiarizedSegment] = reference + hypothesis
+ let maxTime = allSegments.map(\.endTime).max()!
+ let numFrames = Int(ceil(maxTime / resolution))
+ guard numFrames > 0 else {
+ return DERResult(totalSpeech: 0, falseAlarm: 0, missedSpeech: 0, confusion: 0)
+ }
+
+ // Build collar mask: frames near reference boundaries are excluded from scoring
+ var collarMask = [Bool](repeating: false, count: numFrames)
+ if collar > 0 {
+ for seg in reference {
+ let startFrame = Int(seg.startTime / resolution)
+ let endFrame = Int(seg.endTime / resolution)
+ let collarFrames = Int(collar / resolution)
+
+ for f in max(0, startFrame - collarFrames)..8 speakers, falls back to greedy matching.
+public func computeDERWithOptimalMapping(
+ reference: [DiarizedSegment],
+ hypothesis: [DiarizedSegment],
+ collar: Float = 0.25,
+ resolution: Float = 0.01
+) -> DERResult {
+ let refSpeakers = Set(reference.map(\.speakerId)).sorted()
+ let hypSpeakers = Set(hypothesis.map(\.speakerId)).sorted()
+
+ guard !refSpeakers.isEmpty, !hypSpeakers.isEmpty else {
+ return computeDER(reference: reference, hypothesis: hypothesis,
+ collar: collar, resolution: resolution)
+ }
+
+ // For small speaker counts, try all permutations
+ if hypSpeakers.count <= 8 {
+ return bruteForceOptimalMapping(
+ reference: reference, hypothesis: hypothesis,
+ refSpeakers: refSpeakers, hypSpeakers: hypSpeakers,
+ collar: collar, resolution: resolution
+ )
+ }
+
+ // For large speaker counts, use greedy matching
+ return greedyOptimalMapping(
+ reference: reference, hypothesis: hypothesis,
+ refSpeakers: refSpeakers, hypSpeakers: hypSpeakers,
+ collar: collar, resolution: resolution
+ )
+}
+
+// MARK: - RTTM Parsing
+
+/// Parse RTTM file content into DiarizedSegments.
+public func parseRTTM(_ content: String) -> [DiarizedSegment] {
+ var segments = [DiarizedSegment]()
+ var speakerMap = [String: Int]()
+ var nextId = 0
+
+ for line in content.split(separator: "\n") {
+ let trimmed = line.trimmingCharacters(in: .whitespaces)
+ guard !trimmed.isEmpty, trimmed.hasPrefix("SPEAKER") else { continue }
+
+ let parts = trimmed.split(whereSeparator: \.isWhitespace).map(String.init)
+ guard parts.count >= 8 else { continue }
+
+ guard let start = Float(parts[3]),
+ let dur = Float(parts[4]) else { continue }
+
+ let speaker = parts[7]
+ if speakerMap[speaker] == nil {
+ speakerMap[speaker] = nextId
+ nextId += 1
+ }
+
+ segments.append(DiarizedSegment(
+ startTime: start,
+ endTime: start + dur,
+ speakerId: speakerMap[speaker]!
+ ))
+ }
+
+ return segments.sorted { $0.startTime < $1.startTime }
+}
+
+// MARK: - Internals
+
+private func buildFrameSpeakers(
+ segments: [DiarizedSegment],
+ numFrames: Int,
+ resolution: Float
+) -> [[Int]] {
+ var result = [[Int]](repeating: [], count: numFrames)
+
+ for seg in segments {
+ let startFrame = max(0, Int(seg.startTime / resolution))
+ let endFrame = min(numFrames, Int(seg.endTime / resolution))
+
+ for f in startFrame.. Int {
+ var matched = 0
+ for rSpk in ref {
+ if hyp.contains(rSpk) {
+ matched += 1
+ }
+ }
+ return matched
+}
+
+private func bruteForceOptimalMapping(
+ reference: [DiarizedSegment],
+ hypothesis: [DiarizedSegment],
+ refSpeakers: [Int],
+ hypSpeakers: [Int],
+ collar: Float,
+ resolution: Float
+) -> DERResult {
+ let permutations = generatePermutations(Array(0.. DERResult {
+ // Compute overlap matrix between ref and hyp speakers
+ let allSegs: [DiarizedSegment] = reference + hypothesis
+ let maxTime = allSegs.map(\.endTime).max()!
+ let numFrames = Int(ceil(maxTime / resolution))
+
+ let refFrames = buildFrameSpeakers(segments: reference, numFrames: numFrames, resolution: resolution)
+ let hypFrames = buildFrameSpeakers(segments: hypothesis, numFrames: numFrames, resolution: resolution)
+
+ // Overlap[r][h] = number of frames where ref speaker r and hyp speaker h both active
+ var overlap = [[Int]](repeating: [Int](repeating: 0, count: hypSpeakers.count), count: refSpeakers.count)
+ let refIndex = Dictionary(uniqueKeysWithValues: refSpeakers.enumerated().map { ($1, $0) })
+ let hypIndex = Dictionary(uniqueKeysWithValues: hypSpeakers.enumerated().map { ($1, $0) })
+
+ for f in 0..()
+ var usedHyp = Set()
+
+ for _ in 0.. bestOverlap {
+ bestOverlap = overlap[r][h]
+ bestR = r
+ bestH = h
+ }
+ }
+ }
+ guard bestR >= 0 else { break }
+ mapping[hypSpeakers[bestH]] = refSpeakers[bestR]
+ usedRef.insert(bestR)
+ usedHyp.insert(bestH)
+ }
+
+ let remapped = hypothesis.map { seg in
+ DiarizedSegment(
+ startTime: seg.startTime,
+ endTime: seg.endTime,
+ speakerId: mapping[seg.speakerId] ?? (1000 + seg.speakerId)
+ )
+ }
+
+ return computeDER(reference: reference, hypothesis: remapped,
+ collar: collar, resolution: resolution)
+}
+
+private func generatePermutations(_ elements: [Int]) -> [[Int]] {
+ if elements.count <= 1 { return [elements] }
+
+ var result = [[Int]]()
+ // Generate permutations of size elements.count from range 0.. [DiarizedSegment] {
+ guard !segments.isEmpty else { return [] }
+
+ var bySpeaker = [Int: [DiarizedSegment]]()
+ for seg in segments {
+ bySpeaker[seg.speakerId, default: []].append(seg)
+ }
+
+ var merged = [DiarizedSegment]()
+ for (spk, spkSegs) in bySpeaker {
+ let sorted = spkSegs.sorted { $0.startTime < $1.startTime }
+ var current = sorted[0]
+
+ for i in 1.. [DiarizedSegment] {
+ let usedIds = Set(segments.map(\.speakerId)).sorted()
+ let idMap = Dictionary(uniqueKeysWithValues: usedIds.enumerated().map { ($1, $0) })
+ return segments.map {
+ DiarizedSegment(
+ startTime: $0.startTime,
+ endTime: $0.endTime,
+ speakerId: idMap[$0.speakerId] ?? $0.speakerId
+ )
+ }
+ }
+
+ /// Resample audio via AVAudioConverter (delegates to AudioFileLoader).
+ static func resample(_ audio: [Float], from sourceSR: Int, to targetSR: Int) -> [Float] {
+ AudioFileLoader.resample(audio, from: sourceSR, to: targetSR)
+ }
+
+ // MARK: - Constrained Agglomerative Clustering
+
+ /// Item for constrained agglomerative clustering.
+ struct ClusterItem {
+ let windowIndex: Int
+ let localSpeakerId: Int
+ let embedding: [Float]
+ }
+
+ /// Constrained agglomerative clustering with centroid linkage and cosine distance.
+ ///
+ /// Items from the same window can never be merged (same-window constraint).
+ /// Merges closest unconstrained pair until distance exceeds threshold.
+ ///
+ /// - Parameters:
+ /// - items: per-window per-speaker embeddings
+ /// - threshold: cosine distance threshold (0–2). Pairs with distance >= threshold are not merged.
+ /// - Returns: cluster assignment for each item, and cluster centroids
+ static func constrainedAgglomerativeClustering(
+ items: [ClusterItem],
+ threshold: Float
+ ) -> (clusterAssignment: [Int], centroids: [[Float]]) {
+ guard !items.isEmpty else { return ([], []) }
+ if items.count == 1 {
+ return ([0], [items[0].embedding])
+ }
+
+ let n = items.count
+ let dim = items[0].embedding.count
+
+ // Each item starts as its own cluster
+ var clusterOf = Array(0.. 1 {
+ // Find closest unconstrained pair
+ var bestDist: Float = Float.greatestFiniteMagnitude
+ var bestI = -1, bestJ = -1
+
+ let activeList = active.sorted()
+ for ai in 0..= 0 else { break }
+
+ // Merge bestJ into bestI
+ let sizeI = clusterMembers[bestI].count
+ let sizeJ = clusterMembers[bestJ].count
+ let totalSize = Float(sizeI + sizeJ)
+
+ // Weighted average centroid
+ var newCentroid = [Float](repeating: 0, count: dim)
+ for d in 0.. Float {
+ let n = min(a.count, b.count)
+ guard n > 0 else { return 2.0 }
+
+ var dot: Float = 0, normA: Float = 0, normB: Float = 0
+ for i in 0.. 1e-10 else { return 2.0 }
+ return 1.0 - dot / denom
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/DiarizationPipeline.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/DiarizationPipeline.swift
new file mode 100644
index 0000000..b8a2779
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/DiarizationPipeline.swift
@@ -0,0 +1,570 @@
+import Foundation
+import MLXCommon
+import MLX
+import AudioCommon
+
+// MARK: - Configuration
+
+/// Configuration for speaker diarization thresholds.
+///
+/// Shared by all diarization engines (Pyannote and Sortformer).
+public struct DiarizationConfig: Sendable {
+ /// Onset threshold for speaker activity
+ public var onset: Float
+ /// Offset threshold for speaker activity
+ public var offset: Float
+ /// Minimum speech segment duration in seconds
+ public var minSpeechDuration: Float
+ /// Minimum silence duration between segments in seconds
+ public var minSilenceDuration: Float
+ /// Cosine distance threshold for merging speaker clusters (0.0-2.0).
+ /// Lower = more merges (fewer speakers). Default 0.715.
+ public var clusteringThreshold: Float
+
+ public init(
+ onset: Float = 0.5,
+ offset: Float = 0.3,
+ minSpeechDuration: Float = 0.3,
+ minSilenceDuration: Float = 0.15,
+ clusteringThreshold: Float = 0.715
+ ) {
+ self.onset = onset
+ self.offset = offset
+ self.minSpeechDuration = minSpeechDuration
+ self.minSilenceDuration = minSilenceDuration
+ self.clusteringThreshold = clusteringThreshold
+ }
+
+ public static let `default` = DiarizationConfig()
+}
+
+// MARK: - Result
+
+/// Result of speaker diarization.
+public struct DiarizationResult: Sendable {
+ /// Diarized speech segments with speaker IDs
+ public let segments: [DiarizedSegment]
+ /// Number of distinct speakers found
+ public let numSpeakers: Int
+ /// Centroid embedding for each speaker (speaker ID → 256-dim embedding)
+ public let speakerEmbeddings: [[Float]]
+
+ public init(segments: [DiarizedSegment], numSpeakers: Int, speakerEmbeddings: [[Float]]) {
+ self.segments = segments
+ self.numSpeakers = numSpeakers
+ self.speakerEmbeddings = speakerEmbeddings
+ }
+}
+
+// MARK: - Pipeline
+
+/// Pyannote-based speaker diarization: segmentation + per-window embedding + constrained clustering.
+///
+/// - Warning: This class is not thread-safe. Create separate instances for concurrent use.
+///
+/// Pipeline (with optional VAD pre-filter):
+/// 0. **VAD Pre-filter** (optional): Silero VAD masks non-speech regions → reduces false alarms
+/// 1. **Segmentation**: Pyannote on 10s sliding windows (50% overlap) → per-speaker probability tracks
+/// 2. **Per-window Embedding**: WeSpeaker 256-dim embedding per local speaker from non-overlapping speech
+/// 3. **Constrained Clustering**: Agglomerative clustering with same-window constraint → global speaker IDs
+///
+/// ```swift
+/// let pipeline = try await PyannoteDiarizationPipeline.fromPretrained(useVADFilter: true)
+/// let result = pipeline.diarize(audio: samples, sampleRate: 16000)
+/// for seg in result.segments {
+/// print("Speaker \(seg.speakerId): [\(seg.startTime)s - \(seg.endTime)s]")
+/// }
+/// ```
+public final class PyannoteDiarizationPipeline {
+
+ /// Pyannote segmentation model
+ let segmentationModel: SegmentationModel
+
+ /// Segmentation config
+ let segConfig: SegmentationConfig
+
+ /// WeSpeaker embedding model
+ public let embeddingModel: WeSpeakerModel
+
+ /// Optional Silero VAD for pre-filtering non-speech
+ let vadModel: SileroVADModel?
+
+ init(
+ segmentationModel: SegmentationModel,
+ segConfig: SegmentationConfig,
+ embeddingModel: WeSpeakerModel,
+ vadModel: SileroVADModel? = nil
+ ) {
+ self.segmentationModel = segmentationModel
+ self.segConfig = segConfig
+ self.embeddingModel = embeddingModel
+ self.vadModel = vadModel
+ }
+
+ /// Load pre-trained models for diarization.
+ ///
+ /// Downloads both the pyannote segmentation model and WeSpeaker embedding model.
+ ///
+ /// - Parameters:
+ /// - segModelId: HuggingFace model ID for segmentation
+ /// - embModelId: HuggingFace model ID for speaker embeddings (auto-selected by engine if nil)
+ /// - embeddingEngine: inference backend for speaker embeddings (`.mlx` or `.coreml`)
+ /// - progressHandler: callback for download progress
+ /// - Returns: ready-to-use diarization pipeline
+ public static func fromPretrained(
+ segModelId: String = PyannoteVADModel.defaultModelId,
+ embModelId: String? = nil,
+ embeddingEngine: WeSpeakerEngine = .mlx,
+ useVADFilter: Bool = false,
+ cacheBaseDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> PyannoteDiarizationPipeline {
+ progressHandler?(0.0, "Downloading segmentation model...")
+
+ // Load segmentation model
+ let segCacheDir = try HuggingFaceDownloader.getCacheDirectory(for: segModelId, basePath: cacheBaseDir)
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: segModelId,
+ to: segCacheDir,
+ offlineMode: offlineMode,
+ progressHandler: { progress in
+ progressHandler?(progress * 0.3, "Downloading segmentation weights...")
+ }
+ )
+
+ let segConfig = SegmentationConfig.default
+ let segModel = SegmentationModel(config: segConfig)
+ try SegmentationWeightLoader.loadWeights(model: segModel, from: segCacheDir)
+
+ progressHandler?(0.3, "Downloading speaker embedding model...")
+
+ // Load embedding model
+ let resolvedEmbModelId = embModelId ?? (embeddingEngine == .coreml ? WeSpeakerModel.defaultCoreMLModelId : WeSpeakerModel.defaultModelId)
+ let embCacheDir: URL? = if let cacheBaseDir { try HuggingFaceDownloader.getCacheDirectory(for: resolvedEmbModelId, basePath: cacheBaseDir) } else { nil }
+ let embModel = try await WeSpeakerModel.fromPretrained(
+ modelId: embModelId,
+ engine: embeddingEngine,
+ cacheDir: embCacheDir,
+ offlineMode: offlineMode,
+ progressHandler: { progress, status in
+ progressHandler?(0.3 + progress * 0.4, status)
+ }
+ )
+
+ // Optionally load Silero VAD for pre-filtering
+ var vadModel: SileroVADModel? = nil
+ if useVADFilter {
+ progressHandler?(0.7, "Downloading VAD filter model...")
+ let vadCacheDir: URL? = if let cacheBaseDir { try HuggingFaceDownloader.getCacheDirectory(for: SileroVADModel.defaultModelId, basePath: cacheBaseDir) } else { nil }
+ vadModel = try await SileroVADModel.fromPretrained(
+ engine: .mlx,
+ cacheDir: vadCacheDir,
+ offlineMode: offlineMode,
+ progressHandler: { progress, status in
+ progressHandler?(0.7 + progress * 0.25, status)
+ }
+ )
+ }
+
+ progressHandler?(1.0, "Ready")
+
+ return PyannoteDiarizationPipeline(
+ segmentationModel: segModel,
+ segConfig: segConfig,
+ embeddingModel: embModel,
+ vadModel: vadModel
+ )
+ }
+
+ /// Run speaker diarization on audio.
+ ///
+ /// - Parameters:
+ /// - audio: PCM Float32 audio samples
+ /// - sampleRate: sample rate of the input audio
+ /// - config: diarization configuration
+ /// - Returns: diarization result with speaker-labeled segments
+ public func diarize(
+ audio: [Float],
+ sampleRate: Int,
+ config: DiarizationConfig = .default
+ ) -> DiarizationResult {
+ diarize(audio: audio, sampleRate: sampleRate, config: config, progressHandler: nil)
+ }
+
+ /// Diarize audio with progress reporting and optional cancellation.
+ ///
+ /// Same as `diarize(audio:sampleRate:config:)` but reports progress during
+ /// the two most expensive stages (segmentation and embedding extraction).
+ /// The handler returns a `Bool`: `true` to continue, `false` to cancel.
+ /// When cancelled, an empty `DiarizationResult` is returned immediately.
+ ///
+ /// - Parameters:
+ /// - audio: PCM Float32 audio samples
+ /// - sampleRate: sample rate of the input audio
+ /// - config: diarization configuration
+ /// - progressHandler: called with (progress 0.0–1.0, stage description);
+ /// return `true` to continue or `false` to cancel
+ /// - Returns: diarization result with speaker-labeled segments
+ public func diarize(
+ audio: [Float],
+ sampleRate: Int,
+ config: DiarizationConfig = .default,
+ progressHandler: ((Float, String) -> Bool)?
+ ) -> DiarizationResult {
+ let emptyResult = DiarizationResult(segments: [], numSpeakers: 0, speakerEmbeddings: [])
+ let samples = DiarizationHelpers.resample(audio, from: sampleRate, to: segConfig.sampleRate)
+
+ // Stage 0 (optional): VAD pre-filter — mask non-speech to reduce false alarms
+ let speechMask: [SpeechSegment]?
+ if let vadModel {
+ if progressHandler?(0, "VAD pre-filtering") == false {
+ return emptyResult
+ }
+ speechMask = vadModel.detectSpeech(
+ audio: samples, sampleRate: segConfig.sampleRate)
+ } else {
+ speechMask = nil
+ }
+
+ if let speechMask, speechMask.isEmpty {
+ return emptyResult
+ }
+
+ // Run embedding-clustered diarization pipeline
+ return runEmbeddingClusteredDiarization(
+ samples: samples, config: config, speechMask: speechMask,
+ progressHandler: progressHandler)
+ }
+
+ /// Extract segments of a target speaker from audio.
+ ///
+ /// Given a reference embedding (from `WeSpeakerModel.embed()`), finds the
+ /// speaker with highest cosine similarity and returns their segments.
+ ///
+ /// - Parameters:
+ /// - audio: PCM Float32 audio samples
+ /// - sampleRate: sample rate of the input audio
+ /// - targetEmbedding: 256-dim reference embedding of the target speaker
+ /// - config: diarization configuration
+ /// - Returns: speech segments belonging to the target speaker
+ public func extractSpeaker(
+ audio: [Float],
+ sampleRate: Int,
+ targetEmbedding: [Float],
+ config: DiarizationConfig = .default
+ ) -> [SpeechSegment] {
+ let result = diarize(audio: audio, sampleRate: sampleRate, config: config)
+
+ guard result.numSpeakers > 0 else { return [] }
+
+ // Find speaker with highest cosine similarity to target
+ var bestSpeaker = 0
+ var bestSimilarity: Float = -1
+
+ for (i, centroid) in result.speakerEmbeddings.enumerated() {
+ let sim = WeSpeakerModel.cosineSimilarity(centroid, targetEmbedding)
+ if sim > bestSimilarity {
+ bestSimilarity = sim
+ bestSpeaker = i
+ }
+ }
+
+ return result.segments
+ .filter { $0.speakerId == bestSpeaker }
+ .map { SpeechSegment(startTime: $0.startTime, endTime: $0.endTime) }
+ }
+
+ // MARK: - Embedding-Clustered Diarization
+
+ /// Per-window raw probability tracks (3 speakers × nFrames).
+ private struct WindowProbs {
+ let startSample: Int
+ let endSample: Int
+ /// Speaker probability tracks [3][nFrames]
+ let tracks: [[Float]]
+ }
+
+ /// Per-window per-speaker embedding for clustering.
+ private struct WindowSpeakerEmbedding {
+ let windowIndex: Int
+ let localSpeakerId: Int
+ let embedding: [Float]
+ }
+
+ /// Run diarization using per-window speaker embeddings + constrained agglomerative clustering.
+ ///
+ /// 1. Segment all windows → per-speaker probability tracks
+ /// 2. Extract per-window per-speaker embeddings from non-overlapping speech
+ /// 3. Constrained clustering (same-window items never merge) → global speaker IDs
+ /// 4. Map cluster IDs back to binarized segments
+ private func runEmbeddingClusteredDiarization(
+ samples: [Float],
+ config: DiarizationConfig,
+ speechMask: [SpeechSegment]?,
+ progressHandler: ((Float, String) -> Bool)? = nil
+ ) -> DiarizationResult {
+ let windowDuration: Float = 10.0
+ let sampleRate = segConfig.sampleRate
+ let windowSamples = Int(windowDuration * Float(sampleRate))
+ let framesPerChunk = 589
+ let frameDuration = windowDuration / Float(framesPerChunk)
+ let stepSamples = windowSamples / 2 // 50% overlap
+
+ let numSamples = samples.count
+ guard numSamples > 0 else {
+ return DiarizationResult(segments: [], numSpeakers: 0, speakerEmbeddings: [])
+ }
+
+ // Generate window positions with 50% overlap
+ var positions = [(start: Int, end: Int)]()
+ if numSamples <= windowSamples {
+ positions.append((0, numSamples))
+ } else {
+ var start = 0
+ while start + windowSamples <= numSamples {
+ positions.append((start, start + windowSamples))
+ start += stepSamples
+ }
+ if positions.isEmpty || positions.last!.end < numSamples {
+ positions.append((numSamples - windowSamples, numSamples))
+ }
+ }
+
+ // Step 1: Run segmentation on all windows, collect probability tracks
+ // Progress: both steps iterate over all windows, so total = 2 * windowCount
+ let totalUnits = positions.count * 2
+ var completedUnits = 0
+ var windowProbs = [WindowProbs]()
+
+ let emptyResult = DiarizationResult(segments: [], numSpeakers: 0, speakerEmbeddings: [])
+
+ for (posIdx, (start, end)) in positions.enumerated() {
+ completedUnits += 1
+ if progressHandler?(Float(completedUnits) / Float(totalUnits), "Segmenting \(posIdx + 1)/\(positions.count)") == false {
+ return emptyResult
+ }
+
+ var window = Array(samples[start..= config.offset {
+ otherActive = true
+ break
+ }
+ }
+ if otherActive { continue }
+
+ // Extract audio samples for this frame
+ let frameStartSample = windowStartSample + Int(Float(frame) * frameDuration * Float(sampleRate))
+ let frameEndSample = min(
+ windowStartSample + Int(Float(frame + 1) * frameDuration * Float(sampleRate)),
+ samples.count
+ )
+ if frameEndSample > frameStartSample {
+ spkAudio.append(contentsOf: samples[frameStartSample..= minEmbeddingSamples else { continue }
+
+ let embedding = embeddingModel.embed(audio: spkAudio, sampleRate: sampleRate)
+ windowEmbeddings.append(WindowSpeakerEmbedding(
+ windowIndex: wIdx, localSpeakerId: localSpk, embedding: embedding))
+ }
+ }
+
+ // Handle edge case: no embeddings could be extracted
+ guard !windowEmbeddings.isEmpty else {
+ return DiarizationResult(segments: [], numSpeakers: 0, speakerEmbeddings: [])
+ }
+
+ // Step 3: Constrained agglomerative clustering
+ let clusterItems = windowEmbeddings.map {
+ DiarizationHelpers.ClusterItem(
+ windowIndex: $0.windowIndex,
+ localSpeakerId: $0.localSpeakerId,
+ embedding: $0.embedding)
+ }
+
+ let (clusterAssignment, centroids) = DiarizationHelpers.constrainedAgglomerativeClustering(
+ items: clusterItems, threshold: config.clusteringThreshold)
+
+ // Build mapping: (windowIndex, localSpeakerId) → global cluster ID
+ var localToGlobal = [Int: [Int: Int]]() // windowIndex → (localSpeakerId → globalId)
+ for (i, we) in windowEmbeddings.enumerated() {
+ localToGlobal[we.windowIndex, default: [:]][we.localSpeakerId] = clusterAssignment[i]
+ }
+
+ // Step 4: Build segments with global speaker IDs
+ var diarizedSegments = [DiarizedSegment]()
+
+ for (w, wp) in windowProbs.enumerated() {
+ let windowStartTime = Float(wp.startSample) / Float(sampleRate)
+ let windowEndTime = Float(wp.endSample) / Float(sampleRate)
+
+ // Center zone ownership (same as before)
+ let prevEnd = w > 0 ?
+ Float(positions[w - 1].end) / Float(sampleRate) : 0
+ let nextStart = w + 1 < positions.count ?
+ Float(positions[w + 1].start) / Float(sampleRate) : Float(numSamples) / Float(sampleRate)
+ let ownStart = w > 0 ? (windowStartTime + prevEnd) / 2 : 0
+ let ownEnd = w + 1 < positions.count ?
+ (windowEndTime + nextStart) / 2 : Float(numSamples) / Float(sampleRate)
+
+ for localSpk in 0..<3 {
+ guard let globalSpk = localToGlobal[w]?[localSpk] else {
+ continue // No embedding for this speaker — skip (insufficient audio)
+ }
+
+ let probs = wp.tracks[localSpk]
+ let segments = PowersetDecoder.binarize(
+ probs: probs, onset: config.onset,
+ offset: config.offset, frameDuration: frameDuration)
+
+ for seg in segments {
+ let absStart = windowStartTime + seg.startTime
+ let absEnd = min(windowStartTime + seg.endTime, windowEndTime)
+
+ // Clip to center zone
+ let clippedStart = max(absStart, ownStart)
+ let clippedEnd = min(absEnd, ownEnd)
+ guard clippedEnd - clippedStart >= config.minSpeechDuration else { continue }
+
+ // Apply VAD mask if present
+ if let speechMask {
+ if let trimmed = trimToSpeechMask(
+ start: clippedStart, end: clippedEnd,
+ speechRegions: speechMask,
+ minDuration: config.minSpeechDuration
+ ) {
+ diarizedSegments.append(DiarizedSegment(
+ startTime: trimmed.startTime,
+ endTime: trimmed.endTime,
+ speakerId: globalSpk
+ ))
+ }
+ } else {
+ diarizedSegments.append(DiarizedSegment(
+ startTime: clippedStart,
+ endTime: clippedEnd,
+ speakerId: globalSpk
+ ))
+ }
+ }
+ }
+ }
+
+ diarizedSegments.sort { $0.startTime < $1.startTime }
+
+ // Compact speaker IDs and merge
+ diarizedSegments = DiarizationHelpers.compactSpeakerIds(diarizedSegments)
+ let merged = DiarizationHelpers.mergeSegments(
+ diarizedSegments, minSilence: config.minSilenceDuration)
+ let numSpeakers = Set(merged.map(\.speakerId)).count
+
+ // Re-compact centroids to match compacted speaker IDs
+ let finalCentroids: [[Float]]
+ if numSpeakers <= centroids.count {
+ finalCentroids = Array(centroids.prefix(numSpeakers))
+ } else {
+ // Pad with zero embeddings for speakers that had no embedding
+ var padded = centroids
+ while padded.count < numSpeakers {
+ padded.append([Float](repeating: 0, count: 256))
+ }
+ finalCentroids = padded
+ }
+
+ return DiarizationResult(
+ segments: merged,
+ numSpeakers: numSpeakers,
+ speakerEmbeddings: finalCentroids
+ )
+ }
+
+ /// Trim a segment to intersect with speech regions.
+ private func trimToSpeechMask(
+ start: Float, end: Float,
+ speechRegions: [SpeechSegment],
+ minDuration: Float
+ ) -> (startTime: Float, endTime: Float)? {
+ let segDuration = end - start
+ guard segDuration > 0 else { return nil }
+
+ var totalOverlap: Float = 0
+ var trimStart: Float = end
+ var trimEnd: Float = start
+
+ for vad in speechRegions {
+ let oStart = max(start, vad.startTime)
+ let oEnd = min(end, vad.endTime)
+ if oStart < oEnd {
+ totalOverlap += oEnd - oStart
+ trimStart = min(trimStart, oStart)
+ trimEnd = max(trimEnd, oEnd)
+ }
+ }
+
+ guard totalOverlap / segDuration >= 0.5,
+ trimEnd - trimStart >= minDuration else { return nil }
+ return (trimStart, trimEnd)
+ }
+
+}
+
+/// Backwards-compatible alias for the renamed pipeline.
+public typealias DiarizationPipeline = PyannoteDiarizationPipeline
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/FireRedVAD.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/FireRedVAD.swift
new file mode 100644
index 0000000..f66c3dc
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/FireRedVAD.swift
@@ -0,0 +1,505 @@
+import Foundation
+import Accelerate
+import AudioCommon
+
+#if canImport(CoreML)
+import CoreML
+#endif
+
+/// Voice Activity Detection using FireRedVAD (DFSMN, CoreML).
+///
+/// A lightweight 588K-param model using DFSMN blocks (depthwise Conv1d)
+/// for temporal context. Runs on Neural Engine + CPU via CoreML.
+///
+/// - Warning: This class is not thread-safe. Create separate instances for concurrent use.
+///
+/// ```swift
+/// let vad = try await FireRedVADModel.fromPretrained()
+/// let segments = vad.detectSpeech(audio: samples, sampleRate: 16000)
+/// ```
+public final class FireRedVADModel {
+
+ /// Default HuggingFace model ID
+ public static let defaultModelId = "aufklarer/FireRedVAD-CoreML"
+
+ /// Whether the model weights are loaded and ready for inference.
+ var _isLoaded = true
+
+ #if canImport(CoreML)
+ /// CoreML compiled model
+ private let coremlModel: MLModel
+ #endif
+
+ /// Feature extractor: 80-dim log Mel fbank (Kaldi-compatible)
+ private let featureExtractor: KaldiFbankExtractor
+
+ /// Post-processing config
+ public var speechThreshold: Float = 0.4
+ public var smoothWindowSize: Int = 5
+ public var minSpeechDuration: Float = 0.2
+ public var minSilenceDuration: Float = 0.2
+
+ #if canImport(CoreML)
+ init(coremlModel: MLModel) {
+ self.coremlModel = coremlModel
+ self.featureExtractor = KaldiFbankExtractor()
+ }
+ #endif
+
+ // MARK: - Model Loading
+
+ /// Load FireRedVAD from HuggingFace.
+ public static func fromPretrained(
+ modelId: String = defaultModelId,
+ cacheDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> FireRedVADModel {
+ #if canImport(CoreML)
+ progressHandler?(0.0, "Downloading model...")
+
+ let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
+
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: modelId,
+ to: cacheDir,
+ additionalFiles: [
+ "fireredvad.mlmodelc/**",
+ "config.json",
+ "cmvn.json",
+ ],
+ offlineMode: offlineMode,
+ progressHandler: { progress in
+ progressHandler?(progress * 0.8, "Downloading model...")
+ }
+ )
+
+ progressHandler?(0.8, "Loading CoreML model...")
+
+ let modelURL = cacheDir.appendingPathComponent(
+ "fireredvad.mlmodelc", isDirectory: true)
+ guard FileManager.default.fileExists(atPath: modelURL.path) else {
+ throw AudioModelError.modelLoadFailed(
+ modelId: modelId,
+ reason: "CoreML model not found at \(modelURL.path)")
+ }
+
+ let mlConfig = MLModelConfiguration()
+ mlConfig.computeUnits = CoreMLComputeUnitsResolver.resolved(default: .cpuAndNeuralEngine)
+
+ let model = try MLModel(contentsOf: modelURL, configuration: mlConfig)
+
+ progressHandler?(1.0, "Ready")
+ return FireRedVADModel(coremlModel: model)
+ #else
+ throw AudioModelError.invalidConfiguration(
+ model: "FireRedVAD", reason: "CoreML not available on this platform")
+ #endif
+ }
+
+ /// Load FireRedVAD from a local directory containing fireredvad.mlmodelc.
+ public static func fromLocal(path: String) throws -> FireRedVADModel {
+ #if canImport(CoreML)
+ let modelURL = URL(fileURLWithPath: path)
+ .appendingPathComponent("fireredvad.mlmodelc", isDirectory: true)
+ guard FileManager.default.fileExists(atPath: modelURL.path) else {
+ throw AudioModelError.modelLoadFailed(
+ modelId: path,
+ reason: "CoreML model not found at \(modelURL.path)")
+ }
+
+ let mlConfig = MLModelConfiguration()
+ mlConfig.computeUnits = CoreMLComputeUnitsResolver.resolved(default: .cpuAndNeuralEngine)
+
+ let model = try MLModel(contentsOf: modelURL, configuration: mlConfig)
+ return FireRedVADModel(coremlModel: model)
+ #else
+ throw AudioModelError.invalidConfiguration(
+ model: "FireRedVAD", reason: "CoreML not available on this platform")
+ #endif
+ }
+
+ // MARK: - Inference
+
+ /// Detect speech from pre-computed features (for testing with reference features).
+ public func detectSpeechFromFeatures(_ features: [Float]) -> [SpeechSegment] {
+ let numFrames = features.count / 80
+ guard numFrames > 0 else { return [] }
+
+ #if canImport(CoreML)
+ let maxFrames = 6000
+ let probs: [Float]
+ if numFrames <= maxFrames {
+ probs = runCoreML(features: features, numFrames: numFrames)
+ } else {
+ var allProbs = [Float]()
+ var offset = 0
+ while offset < numFrames {
+ let chunkFrames = min(maxFrames, numFrames - offset)
+ let chunkStart = offset * 80
+ let chunkEnd = chunkStart + chunkFrames * 80
+ let chunkFeatures = Array(features[chunkStart.. [SpeechSegment] {
+ // Resample if needed
+ let samples: [Float]
+ if sampleRate != 16000 {
+ samples = AudioFileLoader.resample(audio, from: sampleRate, to: 16000)
+ } else {
+ samples = audio
+ }
+
+ // Extract features
+ let features = featureExtractor.extract(samples)
+ guard features.count > 0 else { return [] }
+
+ let numFrames = features.count / 80
+
+ // Run CoreML inference (chunk if > 6000 frames / 60s)
+ #if canImport(CoreML)
+ let maxFrames = 6000
+ let probs: [Float]
+ if numFrames <= maxFrames {
+ probs = runCoreML(features: features, numFrames: numFrames)
+ } else {
+ // Process in chunks
+ var allProbs = [Float]()
+ var offset = 0
+ while offset < numFrames {
+ let chunkFrames = min(maxFrames, numFrames - offset)
+ let chunkStart = offset * 80
+ let chunkEnd = chunkStart + chunkFrames * 80
+ let chunkFeatures = Array(features[chunkStart.. [Float] {
+ // Create MLMultiArray [1, T, 80]
+ guard let input = try? MLMultiArray(
+ shape: [1, NSNumber(value: numFrames), 80],
+ dataType: .float32
+ ) else { return [] }
+
+ let ptr = input.dataPointer.assumingMemoryBound(to: Float.self)
+ features.withUnsafeBufferPointer { src in
+ ptr.update(from: src.baseAddress!, count: min(features.count, numFrames * 80))
+ }
+
+ guard let prediction = try? coremlModel.prediction(
+ from: FireRedVADInput(features: input)
+ ) else { return [] }
+
+ guard let outputArray = prediction.featureValue(
+ for: "probabilities"
+ )?.multiArrayValue else { return [] }
+
+ // Extract probabilities — output is [1, T, 1]
+ var probs = [Float](repeating: 0, count: numFrames)
+ for i in 0.. [Float] {
+ guard windowSize > 1, probs.count > windowSize else { return probs }
+ var smoothed = [Float](repeating: 0, count: probs.count)
+ let half = windowSize / 2
+
+ for i in 0.. [SpeechSegment] {
+ // Binary decisions
+ let decisions = probs.map { $0 >= threshold }
+
+ // Find contiguous speech regions
+ var segments = [SpeechSegment]()
+ var speechStart: Int?
+
+ for i in 0...decisions.count {
+ let isSpeech = i < decisions.count ? decisions[i] : false
+ if isSpeech && speechStart == nil {
+ speechStart = i
+ } else if !isSpeech, let start = speechStart {
+ let startTime = Float(start) * frameShift
+ let endTime = Float(i) * frameShift
+ if endTime - startTime >= minSpeechDuration {
+ segments.append(SpeechSegment(
+ startTime: startTime, endTime: endTime))
+ }
+ speechStart = nil
+ }
+ }
+
+ // Merge segments with short silence gaps
+ guard segments.count > 1 else { return segments }
+ var merged = [segments[0]]
+ for i in 1.. { ["features"] }
+
+ func featureValue(for featureName: String) -> MLFeatureValue? {
+ if featureName == "features" {
+ return MLFeatureValue(multiArray: features)
+ }
+ return nil
+ }
+}
+#endif
+
+// MARK: - Kaldi-compatible Fbank Extractor
+
+/// 80-dim log Mel filterbank extractor matching Kaldi's default settings.
+///
+/// Configuration: 16kHz, 25ms window, 10ms hop, 80 mel bins, snip_edges=true.
+/// Uses Povey window (Hann-like) and log energy.
+final class KaldiFbankExtractor {
+ let sampleRate: Int = 16000
+ let frameLength: Int = 400 // 25ms at 16kHz
+ let frameShift: Int = 160 // 10ms at 16kHz
+ let nMels: Int = 80
+ let nFFT: Int = 512
+ let preemphCoeff: Float = 0.97
+
+ private let window: [Float]
+ private let nBins: Int // nFFT/2 + 1 = 257
+ /// Pre-computed DFT basis: cos[k][n] and sin[k][n] for k=0..nBins-1, n=0..nFFT-1
+ /// Using matrix multiply instead of FFT for exact numerical match with numpy/torch.
+ private let dftCos: [Float] // [nBins * nFFT]
+ private let dftSin: [Float] // [nBins * nFFT]
+ private let melFilterbank: [Float] // [nMels * nBins]
+
+ init() {
+ // Povey window (Hann raised to 0.85 power)
+ var w = [Float](repeating: 0, count: 400)
+ for i in 0..<400 {
+ let hann = 0.5 - 0.5 * cos(2.0 * Float.pi * Float(i) / Float(400))
+ w[i] = pow(hann, 0.85)
+ }
+ self.window = w
+ self.nBins = 257
+
+ // Pre-compute DFT basis vectors
+ // X[k] = sum_n x[n] * exp(-j*2*pi*k*n/N)
+ // = sum_n x[n]*cos(2*pi*k*n/N) - j*sum_n x[n]*sin(2*pi*k*n/N)
+ var cosB = [Float](repeating: 0, count: 257 * 512)
+ var sinB = [Float](repeating: 0, count: 257 * 512)
+ for k in 0..<257 {
+ for n in 0..<512 {
+ let angle = 2.0 * Float.pi * Float(k) * Float(n) / 512.0
+ cosB[k * 512 + n] = cos(angle)
+ sinB[k * 512 + n] = sin(angle)
+ }
+ }
+ self.dftCos = cosB
+ self.dftSin = sinB
+
+ self.melFilterbank = KaldiFbankExtractor.buildMelFilterbank(
+ nMels: 80, nBins: 257, sampleRate: 16000, nFFT: 512)
+ }
+
+ /// Extract 80-dim log Mel features from audio samples.
+ /// Returns flat array [T * 80] where T is the number of frames.
+ ///
+ /// Uses pre-computed DFT basis matrices with Accelerate `vDSP_mmul` for
+ /// exact numerical match with `numpy.fft.rfft` while maintaining high speed.
+ func extract(_ samples: [Float]) -> [Float] {
+ let numFrames = max(0, (samples.count - frameLength) / frameShift + 1)
+ guard numFrames > 0 else { return [] }
+
+ var allFeatures = [Float]()
+ allFeatures.reserveCapacity(numFrames * nMels)
+
+ var padded = [Float](repeating: 0, count: nFFT)
+ var realParts = [Float](repeating: 0, count: nBins)
+ var imagParts = [Float](repeating: 0, count: nBins)
+ var powerSpec = [Float](repeating: 0, count: nBins)
+ var melEnergies = [Float](repeating: 0, count: nMels)
+
+ for frame in 0.. [Float] {
+ func hzToMel(_ hz: Float) -> Float {
+ return 1127.0 * log(1.0 + hz / 700.0)
+ }
+ func melToHz(_ mel: Float) -> Float {
+ return 700.0 * (exp(mel / 1127.0) - 1.0)
+ }
+
+ let fMin: Float = 20.0
+ let fMax = Float(sampleRate) / 2.0
+ let melMin = hzToMel(fMin)
+ let melMax = hzToMel(fMax)
+
+ // Mel center frequencies in Hz
+ var centerFreqs = [Float](repeating: 0, count: nMels + 2)
+ for i in 0..<(nMels + 2) {
+ let mel = melMin + Float(i) * (melMax - melMin) / Float(nMels + 1)
+ centerFreqs[i] = melToHz(mel)
+ }
+
+ // Build triangular filters using Hz-domain weights (Kaldi convention)
+ var filterbank = [Float](repeating: 0, count: nMels * nBins)
+ let binToHz = Float(sampleRate) / Float(nFFT)
+
+ for m in 0..= leftHz && freqHz <= centerHz && centerHz > leftHz {
+ filterbank[m * nBins + b] = (freqHz - leftHz) / (centerHz - leftHz)
+ } else if freqHz > centerHz && freqHz <= rightHz && rightHz > centerHz {
+ filterbank[m * nBins + b] = (rightHz - freqHz) / (rightHz - centerHz)
+ }
+ }
+ }
+
+ return filterbank
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/MelFeatureExtractor.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/MelFeatureExtractor.swift
new file mode 100644
index 0000000..2bcc58b
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/MelFeatureExtractor.swift
@@ -0,0 +1,238 @@
+import Foundation
+import Accelerate
+import MLX
+
+/// 80-dim log-mel feature extractor for WeSpeaker (speaker embeddings).
+///
+/// Uses vDSP for FFT and mel filterbank computation.
+/// Parameters: nFFT=400, hop=160, 80 mel bins, 16kHz.
+/// Matches Kaldi FBank defaults used by WeSpeaker: pre-emphasis=0.97,
+/// HTK mel scale, Povey window, fMin=20Hz.
+class MelFeatureExtractor {
+ let sampleRate: Int = 16000
+ let nFFT: Int = 400
+ let hopLength: Int = 160
+ let nMels: Int = 80
+ let preEmphasis: Float = 0.97
+
+ private let paddedFFT: Int = 512
+ private let log2PaddedFFT: vDSP_Length = 9
+ private var fftSetup: FFTSetup
+ private var window: [Float]
+ private var melFilterbank: [Float] // [nMels, nBins]
+
+ init() {
+ // Hamming window: matches pyannote/wespeaker inference pipeline
+ // (pyannote uses window_type='hamming', NOT Kaldi default 'povey')
+ window = [Float](repeating: 0, count: 400)
+ for i in 0..<400 {
+ window[i] = 0.54 - 0.46 * cos(2.0 * Float.pi * Float(i) / Float(399))
+ }
+
+ guard let setup = vDSP_create_fftsetup(9, FFTRadix(kFFTRadix2)) else {
+ fatalError("Failed to create vDSP FFT setup")
+ }
+ fftSetup = setup
+
+ melFilterbank = []
+ setupMelFilterbank()
+ }
+
+ deinit {
+ vDSP_destroy_fftsetup(fftSetup)
+ }
+
+ private func setupMelFilterbank() {
+ let fMin: Float = 20.0
+ let fMax: Float = Float(sampleRate) / 2.0
+
+ // HTK mel scale (Kaldi default): mel = 2595 * log10(1 + hz/700)
+ func hzToMel(_ hz: Float) -> Float {
+ return 2595.0 * log10(1.0 + hz / 700.0)
+ }
+
+ func melToHz(_ mel: Float) -> Float {
+ return 700.0 * (pow(10.0, mel / 2595.0) - 1.0)
+ }
+
+ let nBins = paddedFFT / 2 + 1 // 257
+
+ var fftFreqs = [Float](repeating: 0, count: nBins)
+ for i in 0.. (melSpec: [Float], nFrames: Int) {
+ let nBins = paddedFFT / 2 + 1
+ let halfPadded = paddedFFT / 2
+
+ // Pre-emphasis: y[n] = x[n] - coeff * x[n-1]
+ var emphasized = [Float](repeating: 0, count: audio.count)
+ if !audio.isEmpty {
+ emphasized[0] = audio[0]
+ for i in 1.. 0 {
+ // Compute per-bin mean
+ var binMeans = [Float](repeating: 0, count: nMels)
+ for frame in 0.. MLXArray {
+ let (melSpec, nFrames) = extractRaw(audio)
+ return MLXArray(melSpec, [nFrames, nMels])
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/PowersetDecoder.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/PowersetDecoder.swift
new file mode 100644
index 0000000..d849835
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/PowersetDecoder.swift
@@ -0,0 +1,73 @@
+import MLX
+
+/// Decodes pyannote's 7-class powerset output to per-speaker probabilities.
+///
+/// The 7 powerset classes represent all possible speaker combinations
+/// for up to 3 speakers:
+/// - 0: non-speech
+/// - 1: speaker 1 alone
+/// - 2: speaker 2 alone
+/// - 3: speaker 3 alone
+/// - 4: speakers 1+2 overlap
+/// - 5: speakers 1+3 overlap
+/// - 6: speakers 2+3 overlap
+enum PowersetDecoder {
+
+ /// Convert 7-class powerset posteriors to per-speaker probabilities.
+ ///
+ /// Each speaker's probability is the sum of all classes where that speaker
+ /// is active (alone or in overlap).
+ ///
+ /// - Parameter posteriors: `[batch, frames, 7]` softmax probabilities
+ /// - Returns: `[batch, frames, 3]` per-speaker probabilities
+ static func speakerProbabilities(from posteriors: MLXArray) -> MLXArray {
+ // spk1: alone(1) + with_spk2(4) + with_spk3(5)
+ let spk1 = posteriors[0..., 0..., 1] + posteriors[0..., 0..., 4] + posteriors[0..., 0..., 5]
+ // spk2: alone(2) + with_spk1(4) + with_spk3(6)
+ let spk2 = posteriors[0..., 0..., 2] + posteriors[0..., 0..., 4] + posteriors[0..., 0..., 6]
+ // spk3: alone(3) + with_spk1(5) + with_spk2(6)
+ let spk3 = posteriors[0..., 0..., 3] + posteriors[0..., 0..., 5] + posteriors[0..., 0..., 6]
+ return stacked([spk1, spk2, spk3], axis: -1)
+ }
+
+ /// Apply hysteresis binarization to per-speaker probabilities.
+ ///
+ /// Expects probabilities in `[0, 1]` range (post-softmax or post-sigmoid).
+ /// If values outside this range are passed, apply sigmoid first.
+ ///
+ /// - Parameters:
+ /// - probs: per-frame probabilities for one speaker `[frames]`, values in [0, 1]
+ /// - onset: threshold to start a segment
+ /// - offset: threshold to end a segment
+ /// - frameDuration: duration of one frame in seconds
+ /// - Returns: array of (startTime, endTime) tuples
+ static func binarize(
+ probs: [Float],
+ onset: Float,
+ offset: Float,
+ frameDuration: Float
+ ) -> [(startTime: Float, endTime: Float)] {
+ var segments = [(startTime: Float, endTime: Float)]()
+ var inSpeech = false
+ var speechStart: Float = 0
+
+ for (i, prob) in probs.enumerated() {
+ let time = Float(i) * frameDuration
+
+ if !inSpeech && prob >= onset {
+ inSpeech = true
+ speechStart = time
+ } else if inSpeech && prob < offset {
+ inSpeech = false
+ segments.append((speechStart, time))
+ }
+ }
+
+ if inSpeech {
+ let endTime = Float(probs.count) * frameDuration
+ segments.append((speechStart, endTime))
+ }
+
+ return segments
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/PyannoteVAD+Memory.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/PyannoteVAD+Memory.swift
new file mode 100644
index 0000000..606fd1b
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/PyannoteVAD+Memory.swift
@@ -0,0 +1,16 @@
+import AudioCommon
+
+extension PyannoteVADModel: ModelMemoryManageable {
+ public var isLoaded: Bool { _isLoaded }
+
+ public func unload() {
+ guard _isLoaded else { return }
+ model.clearParameters()
+ _isLoaded = false
+ }
+
+ public var memoryFootprint: Int {
+ guard _isLoaded else { return 0 }
+ return model.parameterMemoryBytes()
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/Segmentation.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/Segmentation.swift
new file mode 100644
index 0000000..f5976d5
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/Segmentation.swift
@@ -0,0 +1,97 @@
+import MLX
+import MLXNN
+
+/// PyanNet segmentation model: SincNet → BiLSTM → Linear → Classifier.
+///
+/// Produces per-frame powerset class probabilities for up to 3 speakers.
+/// Output classes (7): non-speech, spk1, spk2, spk3, spk1+2, spk1+3, spk2+3
+///
+/// Weight keys at model level:
+/// ```
+/// sincnet.{conv, norm, wav_norm}.*
+/// lstm_fwd.layers.{i}.{Wx, Wh, bias}
+/// lstm_bwd.layers.{i}.{Wx, Wh, bias}
+/// linear.{0,1}.{weight, bias}
+/// classifier.{weight, bias}
+/// ```
+class SegmentationModel: Module {
+ let config: SegmentationConfig
+
+ let sincnet: SincNet
+
+ /// Forward and backward LSTM stacks (at top level for weight loading)
+ @ModuleInfo(key: "lstm_fwd") var lstmFwd: LSTMStack
+ @ModuleInfo(key: "lstm_bwd") var lstmBwd: LSTMStack
+
+ let linear: [Linear]
+ let classifier: Linear
+
+ init(config: SegmentationConfig = .default) {
+ self.config = config
+
+ self.sincnet = SincNet(config: config)
+
+ // Build LSTM stacks
+ let sincnetOutputDim = config.sincnetFilters.last! // 60
+ var fwdLayers = [LSTMLayer]()
+ var bwdLayers = [LSTMLayer]()
+ for i in 0 ..< config.lstmNumLayers {
+ let inSize = (i == 0) ? sincnetOutputDim : config.lstmHiddenSize * 2
+ fwdLayers.append(LSTMLayer(inputSize: inSize, hiddenSize: config.lstmHiddenSize))
+ bwdLayers.append(LSTMLayer(inputSize: inSize, hiddenSize: config.lstmHiddenSize))
+ }
+
+ // Linear layers with LeakyReLU
+ let lstmOutputDim = config.lstmHiddenSize * 2 // 256 (bidirectional)
+ var layers = [Linear]()
+ for i in 0 ..< config.linearNumLayers {
+ let inDim = (i == 0) ? lstmOutputDim : config.linearHiddenSize
+ layers.append(Linear(inDim, config.linearHiddenSize))
+ }
+ self.linear = layers
+
+ self.classifier = Linear(config.linearHiddenSize, config.numClasses)
+
+ // Set @ModuleInfo properties after all stored properties
+ self._lstmFwd.wrappedValue = LSTMStack(layers: fwdLayers)
+ self._lstmBwd.wrappedValue = LSTMStack(layers: bwdLayers)
+ }
+
+ /// Run segmentation on audio.
+ /// - Parameter waveform: `[batch, 1, samples]` (mono, 16kHz)
+ /// - Returns: `[batch, num_frames, num_classes]` class probabilities
+ func callAsFunction(_ waveform: MLXArray) -> MLXArray {
+ // SincNet: [B, 1, T] → [B, 60, ~293]
+ var x = sincnet(waveform)
+
+ // Transpose for LSTM: [B, 60, L] → [B, L, 60]
+ x = x.transposed(0, 2, 1)
+
+ // BiLSTM: [B, L, 60] → [B, L, 256]
+ x = runBiLSTM(x, fwd: lstmFwd, bwd: lstmBwd)
+
+ // Linear layers with LeakyReLU
+ for layer in linear {
+ x = layer(x)
+ x = leakyRelu(x)
+ }
+
+ // Classifier + softmax: [B, L, 7]
+ x = classifier(x)
+ x = softmax(x, axis: -1)
+
+ return x
+ }
+
+ /// Extract speech probability from powerset output.
+ ///
+ /// Classes: [non-speech, spk1, spk2, spk3, spk1+2, spk1+3, spk2+3]
+ /// Speech probability = 1 - P(non-speech).
+ ///
+ /// - Parameter posteriors: `[batch, frames, 7]`
+ /// - Returns: `[batch, frames]` speech probability per frame
+ static func speechProbability(from posteriors: MLXArray) -> MLXArray {
+ let nonSpeech = posteriors[0..., 0..., 0]
+ return 1.0 - nonSpeech
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroModel.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroModel.swift
new file mode 100644
index 0000000..ad607b4
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroModel.swift
@@ -0,0 +1,186 @@
+import Foundation
+import MLX
+import MLXNN
+
+/// Silero VAD v5 neural network: STFT → Encoder → LSTM → Decoder.
+///
+/// Processes 512-sample audio chunks (32ms @ 16kHz) with 64 samples of context
+/// prepended from the previous chunk. The STFT uses a pre-computed DFT basis
+/// stored as Conv1d weights.
+///
+/// Architecture:
+/// ```
+/// Input: 576 samples (64 context + 512 new)
+/// → ReflectionPad(right=64) → 640 samples
+/// → STFT Conv1d(1→258, k=256, s=128) → 4 frames × 258
+/// → Magnitude: √(real² + imag²) → 4 × 129
+/// → Encoder: 4× Conv1d+ReLU → 1 × 128
+/// → LSTM(128→128, 1 layer) → hidden state [B, 128]
+/// → ReLU → Conv1d(128→1, k=1) → Sigmoid → probability
+/// ```
+///
+/// Weight keys:
+/// ```
+/// stft.weight [258, 256, 1]
+/// encoder.{0-3}.weight Conv1d weights
+/// encoder.{0-3}.bias Conv1d biases
+/// lstm.{Wx, Wh, bias} LSTM parameters
+/// decoder.{weight, bias} Final Conv1d
+/// ```
+class SileroVADNetwork: Module {
+
+ // STFT: pre-computed DFT basis as Conv1d (no bias)
+ @ModuleInfo(key: "stft") var stft: Conv1d
+
+ // Encoder: 4 Conv1d + ReLU
+ let encoder: [Conv1d]
+
+ // LSTM: 128→128, 1 layer (reuses LSTMLayer for parameter structure)
+ @ModuleInfo(key: "lstm") var lstm: LSTMLayer
+
+ // Decoder: Conv1d(128→1, k=1)
+ @ModuleInfo(key: "decoder") var decoder: Conv1d
+
+ override init() {
+ // STFT: filter_length=256, hop_length=128, 1→258 channels (129 real + 129 imag)
+ self._stft.wrappedValue = Conv1d(
+ inputChannels: 1, outputChannels: 258, kernelSize: 256,
+ stride: 128, bias: false)
+
+ // Encoder
+ self.encoder = [
+ Conv1d(inputChannels: 129, outputChannels: 128, kernelSize: 3, stride: 1, padding: 1),
+ Conv1d(inputChannels: 128, outputChannels: 64, kernelSize: 3, stride: 2, padding: 1),
+ Conv1d(inputChannels: 64, outputChannels: 64, kernelSize: 3, stride: 2, padding: 1),
+ Conv1d(inputChannels: 64, outputChannels: 128, kernelSize: 3, stride: 1, padding: 1),
+ ]
+
+ // LSTM: input_size=128, hidden_size=128
+ self._lstm.wrappedValue = LSTMLayer(inputSize: 128, hiddenSize: 128)
+
+ // Decoder: 128→1 with kernel=1
+ self._decoder.wrappedValue = Conv1d(
+ inputChannels: 128, outputChannels: 1, kernelSize: 1)
+ }
+
+ /// Forward pass for a single chunk.
+ ///
+ /// - Parameters:
+ /// - samples: `[B, T]` raw audio (576 samples: 64 context + 512 new)
+ /// - h: LSTM hidden state `[1, B, 128]` or nil for initial state
+ /// - c: LSTM cell state `[1, B, 128]` or nil for initial state
+ /// - Returns: `(probability [B], new_h [1, B, 128], new_c [1, B, 128])`
+ func forward(_ samples: MLXArray, h: MLXArray?, c: MLXArray?) -> (MLXArray, MLXArray, MLXArray) {
+ // [B, T] → [B, T, 1] for channels-last Conv1d
+ var x = samples.expandedDimensions(axis: -1)
+
+ // Reflection padding: 64 samples on the RIGHT side only
+ // (matching Silero's pad(input, [0, 64], "reflect"))
+ x = reflectionPadRight(x, padding: 64)
+
+ // STFT via Conv1d: [B, 640, 1] → [B, 4, 258]
+ x = stft(x)
+
+ // Split real/imaginary and compute magnitude
+ let real = x[0..., 0..., ..<129]
+ let imag = x[0..., 0..., 129...]
+ x = sqrt(real * real + imag * imag) // [B, 4, 129]
+
+ // Encoder: 4× Conv1d + ReLU → [B, 1, 128]
+ for conv in encoder {
+ x = relu(conv(x))
+ }
+
+ // LSTM with explicit h/c state
+ // Encoder output is [B, 1, 128] — single timestep
+ let (newH, newC) = lstmForward(x, h: h, c: c)
+
+ // Decoder: use LSTM hidden state h, not full sequence output
+ // h: [1, B, 128] → [B, 128] → [B, 1, 128] for Conv1d
+ let hForDecoder = newH.squeezed(axis: 0).expandedDimensions(axis: 1)
+
+ // ReLU → Conv1d(128→1, k=1) → Sigmoid
+ let prob = sigmoid(decoder(relu(hForDecoder))) // [B, 1, 1]
+
+ return (prob.squeezed(axes: [1, 2]), newH, newC)
+ }
+
+ /// Run LSTM with explicit hidden/cell state for streaming.
+ ///
+ /// Accesses LSTMLayer's parameters directly (Wx, Wh, bias) rather than
+ /// calling its `callAsFunction`, which doesn't support stateful operation.
+ ///
+ /// Returns (new_h [1, B, H], new_c [1, B, H])
+ private func lstmForward(
+ _ x: MLXArray, h: MLXArray?, c: MLXArray?
+ ) -> (MLXArray, MLXArray) {
+ // Project all timesteps: [B, T, 4*H]
+ let projected = addMM(lstm.bias, x, lstm.wx.T)
+ let seqLen = x.dim(-2)
+
+ // Squeeze state from [1, B, H] to [B, H]
+ var hidden = h?.squeezed(axis: 0)
+ var cell = c?.squeezed(axis: 0)
+
+ for t in 0 ..< seqLen {
+ var ifgo = projected[0..., t, 0...]
+ if let h = hidden {
+ ifgo = ifgo + matmul(h, lstm.wh.T)
+ }
+
+ let pieces = split(ifgo, parts: 4, axis: -1)
+ let i = sigmoid(pieces[0])
+ let f = sigmoid(pieces[1])
+ let g = tanh(pieces[2])
+ let o = sigmoid(pieces[3])
+
+ if let c = cell {
+ cell = f * c + i * g
+ } else {
+ cell = i * g
+ }
+ hidden = o * tanh(cell!)
+ }
+
+ let newH = hidden!.expandedDimensions(axis: 0) // [1, B, H]
+ let newC = cell!.expandedDimensions(axis: 0) // [1, B, H]
+
+ return (newH, newC)
+ }
+}
+
+// MARK: - Reflection Padding
+
+/// Right-only reflection padding for 1D data in channels-last format `[B, T, C]`.
+///
+/// Matches Silero's `F.pad(input, [0, 64], mode='reflect')`.
+/// For input `[a, b, c, d, e]` with padding=2: `[a, b, c, d, e, d, c]`
+func reflectionPadRight(_ x: MLXArray, padding: Int) -> MLXArray {
+ let T = x.dim(1)
+ guard padding > 0, T > padding else { return x }
+
+ // Right: reflect indices [T-2, T-3, ..., T-1-padding]
+ let rightIndices = MLXArray(Array(stride(from: T - 2, through: T - 1 - padding, by: -1)))
+ let rightPad = x.take(rightIndices, axis: 1)
+
+ return concatenated([x, rightPad], axis: 1)
+}
+
+/// Symmetric reflection padding for 1D data in channels-last format `[B, T, C]`.
+///
+/// Pads the time dimension by reflecting values at both boundaries.
+/// For input `[a, b, c, d, e]` with padding=2: `[c, b, a, b, c, d, e, d, c]`
+func reflectionPad1d(_ x: MLXArray, padding: Int) -> MLXArray {
+ let T = x.dim(1)
+ guard padding > 0, T > padding else { return x }
+
+ // Left: reflect indices [padding, padding-1, ..., 1]
+ let leftIndices = MLXArray(Array(stride(from: padding, through: 1, by: -1)))
+ let leftPad = x.take(leftIndices, axis: 1)
+
+ // Right: reflect indices [T-2, T-3, ..., T-1-padding]
+ let rightIndices = MLXArray(Array(stride(from: T - 2, through: T - 1 - padding, by: -1)))
+ let rightPad = x.take(rightIndices, axis: 1)
+
+ return concatenated([leftPad, x, rightPad], axis: 1)
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroVAD+Memory.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroVAD+Memory.swift
new file mode 100644
index 0000000..1e56ef5
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroVAD+Memory.swift
@@ -0,0 +1,19 @@
+import AudioCommon
+
+extension SileroVADModel: ModelMemoryManageable {
+ public var isLoaded: Bool { _isLoaded }
+
+ public func unload() {
+ guard _isLoaded else { return }
+ network?.clearParameters()
+ #if canImport(CoreML)
+ coremlModel = nil
+ #endif
+ _isLoaded = false
+ }
+
+ public var memoryFootprint: Int {
+ guard _isLoaded else { return 0 }
+ return network?.parameterMemoryBytes() ?? 0
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroVAD.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroVAD.swift
new file mode 100644
index 0000000..ca03eb2
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroVAD.swift
@@ -0,0 +1,321 @@
+import Foundation
+import MLXCommon
+import MLX
+import AudioCommon
+
+#if canImport(CoreML)
+import CoreML
+#endif
+
+/// Inference engine for Silero VAD.
+public enum SileroVADEngine: String, Sendable {
+ /// MLX backend — runs on GPU via Metal shaders.
+ case mlx
+ /// CoreML backend — runs on Neural Engine + CPU, freeing the GPU.
+ case coreml
+}
+
+/// Streaming Voice Activity Detection using Silero VAD v5.
+///
+/// A lightweight (~260K params) VAD model that processes 512-sample chunks
+/// (32ms @ 16kHz) with sub-millisecond latency. Carries LSTM state across
+/// chunks for streaming operation.
+///
+/// Supports two backends:
+/// - `.mlx`: GPU-based inference via MLX (default)
+/// - `.coreml`: Neural Engine inference via CoreML (lower power, frees GPU)
+///
+/// - Warning: This class is not thread-safe. Create separate instances for concurrent use.
+///
+/// ```swift
+/// let vad = try await SileroVADModel.fromPretrained(engine: .coreml)
+///
+/// // Streaming: process one chunk at a time
+/// let prob = vad.processChunk(samples512) // → 0.0...1.0
+///
+/// // Batch: detect all speech segments
+/// let segments = vad.detectSpeech(audio: samples, sampleRate: 16000)
+/// ```
+public final class SileroVADModel {
+
+ /// The inference engine in use.
+ public let engine: SileroVADEngine
+
+ /// Whether the model weights are loaded and ready for inference.
+ var _isLoaded = true
+
+ /// The MLX neural network (nil when using CoreML engine).
+ let network: SileroVADNetwork?
+
+ // MARK: - MLX State
+
+ /// LSTM hidden state (carried across chunks) — MLX engine
+ private var h: MLXArray?
+ /// LSTM cell state (carried across chunks) — MLX engine
+ private var c: MLXArray?
+
+ // MARK: - CoreML State
+
+ #if canImport(CoreML)
+ /// CoreML compiled model (nil when using MLX engine).
+ var coremlModel: MLModel?
+ /// CoreML LSTM hidden state
+ var coremlH: MLMultiArray?
+ /// CoreML LSTM cell state
+ var coremlC: MLMultiArray?
+ #endif
+
+ /// Context buffer: last 64 samples from previous chunk
+ private var context: [Float]
+
+ /// Default HuggingFace model ID (MLX weights)
+ public static let defaultModelId = "aufklarer/Silero-VAD-v5-MLX"
+
+ /// Default HuggingFace model ID (CoreML weights)
+ public static let defaultCoreMLModelId = "aufklarer/Silero-VAD-v5-CoreML"
+
+ /// Number of audio samples per chunk (32ms @ 16kHz)
+ public static let chunkSize = 512
+
+ /// Number of context samples prepended from previous chunk
+ public static let contextSize = 64
+
+ /// Expected input sample rate
+ public static let sampleRate = 16000
+
+ init(network: SileroVADNetwork) {
+ self.engine = .mlx
+ self.network = network
+ self.context = [Float](repeating: 0, count: Self.contextSize)
+ }
+
+ #if canImport(CoreML)
+ init(coremlModel: MLModel) {
+ self.engine = .coreml
+ self.network = nil
+ self.coremlModel = coremlModel
+ self.context = [Float](repeating: 0, count: Self.contextSize)
+ }
+ #endif
+
+ /// Process a single 512-sample audio chunk and return speech probability.
+ ///
+ /// Maintains internal LSTM state across calls. Call `resetState()` between
+ /// different audio streams.
+ ///
+ /// - Parameter samples: exactly 512 PCM Float32 samples at 16kHz
+ /// - Returns: speech probability in `[0, 1]`
+ public func processChunk(_ samples: [Float]) -> Float {
+ precondition(samples.count == Self.chunkSize,
+ "Chunk must be \(Self.chunkSize) samples, got \(samples.count)")
+
+ // Prepend 64-sample context from previous chunk
+ let fullSamples = context + samples
+
+ // Save last 64 samples as context for next chunk
+ context = Array(samples.suffix(Self.contextSize))
+
+ switch engine {
+ case .mlx:
+ return processChunkMLX(fullSamples)
+ case .coreml:
+ #if canImport(CoreML)
+ return (try? processChunkCoreML(fullSamples)) ?? 0.0
+ #else
+ fatalError("CoreML not available on this platform")
+ #endif
+ }
+ }
+
+ /// MLX inference path.
+ private func processChunkMLX(_ fullSamples: [Float]) -> Float {
+ guard let network else { fatalError("MLX network not loaded") }
+
+ let input = MLXArray(fullSamples).reshaped(1, fullSamples.count)
+ let (prob, newH, newC) = network.forward(input, h: h, c: c)
+
+ h = newH
+ c = newC
+
+ eval(prob)
+ return prob.item(Float.self)
+ }
+
+ /// Reset LSTM and context state.
+ ///
+ /// Call this between processing different audio streams to prevent
+ /// state leakage.
+ public func resetState() {
+ h = nil
+ c = nil
+ #if canImport(CoreML)
+ coremlH = nil
+ coremlC = nil
+ #endif
+ context = [Float](repeating: 0, count: Self.contextSize)
+ }
+
+ /// Detect speech segments in complete audio (batch mode).
+ ///
+ /// Processes the entire audio in 512-sample chunks, collects per-chunk
+ /// probabilities, then applies hysteresis thresholding and duration filtering.
+ ///
+ /// - Parameters:
+ /// - audio: PCM Float32 audio samples
+ /// - sampleRate: sample rate of input audio (resampled to 16kHz if needed)
+ /// - config: VAD configuration (defaults to Silero-tuned thresholds)
+ /// - Returns: array of speech segments with start/end times in seconds
+ public func detectSpeech(
+ audio: [Float],
+ sampleRate: Int,
+ config: VADConfig = .sileroDefault
+ ) -> [SpeechSegment] {
+ let samples: [Float]
+ if sampleRate != Self.sampleRate {
+ samples = AudioFileLoader.resample(audio, from: sampleRate, to: Self.sampleRate)
+ } else {
+ samples = audio
+ }
+
+ resetState()
+
+ // Collect per-chunk probabilities
+ var probs = [Float]()
+ var offset = 0
+
+ while offset + Self.chunkSize <= samples.count {
+ let chunk = Array(samples[offset ..< (offset + Self.chunkSize)])
+ probs.append(processChunk(chunk))
+ offset += Self.chunkSize
+ }
+
+ // Handle remaining samples with zero-padding
+ if offset < samples.count {
+ var lastChunk = Array(samples[offset...])
+ lastChunk.append(contentsOf: [Float](repeating: 0, count: Self.chunkSize - lastChunk.count))
+ probs.append(processChunk(lastChunk))
+ }
+
+ guard !probs.isEmpty else { return [] }
+
+ // Use VADPipeline for hysteresis binarization.
+ // Set windowDuration = probs.count * chunkDuration so frameDuration = 0.032s
+ let chunkDuration: Float = Float(Self.chunkSize) / Float(Self.sampleRate)
+ let batchPipeline = VADPipeline(
+ config: VADConfig(
+ onset: config.onset,
+ offset: config.offset,
+ minSpeechDuration: config.minSpeechDuration,
+ minSilenceDuration: config.minSilenceDuration,
+ windowDuration: Float(probs.count) * chunkDuration,
+ stepRatio: 1.0
+ ),
+ sampleRate: Self.sampleRate,
+ framesPerChunk: probs.count
+ )
+
+ return batchPipeline.binarize(probs: probs)
+ }
+
+ /// Load a pre-trained Silero VAD model from HuggingFace.
+ ///
+ /// Downloads model weights on first use, then caches locally.
+ ///
+ /// - Parameters:
+ /// - modelId: HuggingFace model ID (auto-selected by engine if not specified)
+ /// - engine: inference backend (`.mlx` or `.coreml`)
+ /// - progressHandler: callback for download progress
+ /// - Returns: ready-to-use VAD model
+ public static func fromPretrained(
+ modelId: String? = nil,
+ engine: SileroVADEngine = .mlx,
+ cacheDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> SileroVADModel {
+ let resolvedModelId = modelId ?? (engine == .coreml ? defaultCoreMLModelId : defaultModelId)
+
+ progressHandler?(0.0, "Downloading model...")
+
+ let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: resolvedModelId)
+
+ switch engine {
+ case .mlx:
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: resolvedModelId,
+ to: cacheDir,
+ offlineMode: offlineMode,
+ progressHandler: { progress in
+ progressHandler?(progress * 0.8, "Downloading weights...")
+ }
+ )
+
+ progressHandler?(0.8, "Loading model...")
+
+ let network = SileroVADNetwork()
+ try SileroWeightLoader.loadWeights(model: network, from: cacheDir)
+
+ progressHandler?(1.0, "Ready")
+ return SileroVADModel(network: network)
+
+ case .coreml:
+ #if canImport(CoreML)
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: resolvedModelId,
+ to: cacheDir,
+ additionalFiles: ["silero_vad.mlmodelc/**", "config.json"],
+ offlineMode: offlineMode,
+ progressHandler: { progress in
+ progressHandler?(progress * 0.8, "Downloading CoreML model...")
+ }
+ )
+
+ progressHandler?(0.8, "Loading CoreML model...")
+
+ let modelURL = cacheDir.appendingPathComponent("silero_vad.mlmodelc", isDirectory: true)
+ guard FileManager.default.fileExists(atPath: modelURL.path) else {
+ throw AudioModelError.modelLoadFailed(
+ modelId: resolvedModelId,
+ reason: "CoreML model not found at \(modelURL.path)")
+ }
+
+ let mlConfig = MLModelConfiguration()
+ mlConfig.computeUnits = CoreMLComputeUnitsResolver.resolved(default: .cpuAndNeuralEngine)
+
+ let model: MLModel
+ do {
+ model = try MLModel(contentsOf: modelURL, configuration: mlConfig)
+ } catch {
+ throw AudioModelError.modelLoadFailed(
+ modelId: resolvedModelId,
+ reason: "Failed to load CoreML model",
+ underlying: error)
+ }
+
+ progressHandler?(1.0, "Ready")
+ return SileroVADModel(coremlModel: model)
+ #else
+ throw AudioModelError.invalidConfiguration(
+ model: "SileroVAD", reason: "CoreML not available on this platform")
+ #endif
+ }
+ }
+
+}
+
+// MARK: - VoiceActivityDetectionModel
+
+extension SileroVADModel: VoiceActivityDetectionModel {
+ public var inputSampleRate: Int { Self.sampleRate }
+
+ /// Protocol conformance — uses default Silero config.
+ public func detectSpeech(audio: [Float], sampleRate: Int) -> [SpeechSegment] {
+ detectSpeech(audio: audio, sampleRate: sampleRate, config: .sileroDefault)
+ }
+}
+
+// MARK: - StreamingVADProvider
+
+extension SileroVADModel: StreamingVADProvider {
+ public var chunkSize: Int { Self.chunkSize }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroWeightLoading.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroWeightLoading.swift
new file mode 100644
index 0000000..cc4657a
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroWeightLoading.swift
@@ -0,0 +1,36 @@
+import Foundation
+import MLXCommon
+import MLX
+import MLXNN
+import AudioCommon
+
+/// Weight loading for the Silero VAD v5 model.
+///
+/// Loads from safetensors files produced by `scripts/convert_silero_vad.py`.
+/// The conversion script transposes Conv1d weights and sums LSTM biases,
+/// so loading is a straightforward parameter tree update.
+enum SileroWeightLoader {
+
+ /// Load weights from a directory containing model.safetensors.
+ static func loadWeights(
+ model: SileroVADNetwork,
+ from directory: URL
+ ) throws {
+ let weightsURL = directory.appendingPathComponent("model.safetensors")
+
+ guard FileManager.default.fileExists(atPath: weightsURL.path) else {
+ throw WeightLoadingError.noWeightsFound(directory)
+ }
+
+ let weights = try MLX.loadArrays(url: weightsURL)
+
+ // Build nested parameter tree from flat keys
+ let parameters = ModuleParameters.unflattened(weights)
+
+ // Apply to model
+ try model.update(parameters: parameters, verify: .noUnusedKeys)
+
+ // Materialize all parameters
+ MLX.eval(model.parameters())
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SincNet.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SincNet.swift
new file mode 100644
index 0000000..6418315
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SincNet.swift
@@ -0,0 +1,129 @@
+import MLX
+import MLXNN
+
+/// SincNet feature extractor: 3 conv+pool+norm+activation layers.
+///
+/// The first conv layer uses pre-computed sinc bandpass filters (computed during
+/// weight conversion). At runtime, all three layers are standard Conv1d.
+///
+/// All data flows in MLX channels-last format: `[batch, length, channels]`.
+///
+/// Architecture:
+/// InstanceNorm(1) → Conv1d(1,80,k=251,s=10) → |·| → MaxPool(3,3) → InstanceNorm(80) → LeakyReLU
+/// → Conv1d(80,60,k=5) → MaxPool(3,3) → InstanceNorm(60) → LeakyReLU
+/// → Conv1d(60,60,k=5) → MaxPool(3,3) → InstanceNorm(60) → LeakyReLU
+class SincNet: Module {
+ /// Input waveform normalization
+ @ModuleInfo(key: "wav_norm") var wavNorm: InstanceNorm
+
+ /// Three conv layers (first is pre-computed sinc filterbank)
+ let conv: [Conv1d]
+
+ /// Instance norm after each conv+pool
+ let norm: [InstanceNorm]
+
+ init(config: SegmentationConfig) {
+ let filters = config.sincnetFilters
+ let kernels = config.sincnetKernelSizes
+ let strides = config.sincnetStrides
+
+ // Build conv layers: input channels are [1, 80, 60]
+ let inputChannels = [1] + Array(filters.dropLast())
+ self.conv = zip(zip(inputChannels, filters), zip(kernels, strides)).map { arg in
+ let ((inC, outC), (k, s)) = arg
+ return Conv1d(inputChannels: inC, outputChannels: outC, kernelSize: k, stride: s)
+ }
+
+ self.norm = filters.map { InstanceNorm(dimensions: $0) }
+
+ self._wavNorm.wrappedValue = InstanceNorm(dimensions: 1)
+ }
+
+ /// Forward pass.
+ /// - Parameter x: `[batch, 1, samples]` raw waveform (channels-first, transposed internally)
+ /// - Returns: `[batch, channels, frames]` feature frames (channels-first for LSTM transpose)
+ func callAsFunction(_ x: MLXArray) -> MLXArray {
+ // Convert from [B, C, T] to MLX channels-last [B, T, C]
+ var out = x.transposed(0, 2, 1)
+
+ // Normalize input waveform: [B, T, 1]
+ out = wavNorm(out)
+
+ for i in 0 ..< conv.count {
+ // Conv1d: [B, T, Cin] → [B, T', Cout]
+ out = conv[i](out)
+
+ // First layer uses abs() (sinc filterbank — energy)
+ if i == 0 {
+ out = abs(out)
+ }
+
+ // MaxPool1d(3, stride=3) — pool over time (axis -2)
+ out = maxPool1d(out, kernelSize: 3, stride: 3)
+
+ // InstanceNorm + LeakyReLU
+ out = norm[i](out)
+ out = leakyRelu(out)
+ }
+
+ // Convert back to [B, C, T] for downstream compatibility
+ return out.transposed(0, 2, 1)
+ }
+}
+
+/// Instance normalization for 1D data (channels-last format).
+///
+/// Input shape: `[batch, length, channels]`
+/// Normalizes over the length dimension per-channel, with learnable affine.
+class InstanceNorm: Module {
+ let dimensions: Int
+ var weight: MLXArray
+ var bias: MLXArray
+
+ init(dimensions: Int) {
+ self.dimensions = dimensions
+ self.weight = MLXArray.ones([dimensions])
+ self.bias = MLXArray.zeros([dimensions])
+ }
+
+ func callAsFunction(_ x: MLXArray) -> MLXArray {
+ // x: [B, L, C] — normalize over L dimension (axis 1)
+ let mean = x.mean(axis: 1, keepDims: true)
+ let variance = x.variance(axis: 1, keepDims: true)
+ let eps: Float = 1e-5
+ let normalized = (x - mean) * rsqrt(variance + eps)
+ // Scale and shift: weight and bias are [C], broadcast naturally over [B, L, C]
+ return normalized * weight + bias
+ }
+}
+
+/// 1D max pooling for channels-last data.
+///
+/// Input: `[batch, length, channels]` — pools over the length dimension (axis -2).
+///
+/// - Parameters:
+/// - x: `[batch, length, channels]`
+/// - kernelSize: pooling window
+/// - stride: pooling stride
+/// - Returns: `[batch, floor((length - kernelSize) / stride) + 1, channels]`
+func maxPool1d(_ x: MLXArray, kernelSize: Int, stride: Int) -> MLXArray {
+ let length = x.dim(-2) // time/length axis
+ let outLen = (length - kernelSize) / stride + 1
+
+ // Collect slices for each position in the kernel
+ var slices = [MLXArray]()
+ for k in 0 ..< kernelSize {
+ // Take every stride-th element starting at offset k along time axis
+ let slice = x[0..., .stride(from: k, to: k + outLen * stride, by: stride), 0...]
+ slices.append(slice)
+ }
+
+ // Stack along new axis and take max: [B, outLen, C, kernelSize] → [B, outLen, C]
+ let stacked = MLX.stacked(slices, axis: -1)
+ return stacked.max(axis: -1)
+}
+
+/// LeakyReLU activation.
+func leakyRelu(_ x: MLXArray, negativeSlope: Float = 0.01) -> MLXArray {
+ maximum(x, x * negativeSlope)
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SortformerConfig.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SortformerConfig.swift
new file mode 100644
index 0000000..84c12eb
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SortformerConfig.swift
@@ -0,0 +1,114 @@
+import Foundation
+
+/// Configuration for the Sortformer diarization model.
+///
+/// Sortformer is NVIDIA's end-to-end neural diarization model that directly
+/// predicts speaker activity without requiring separate embedding extraction
+/// or clustering stages.
+public struct SortformerConfig: Sendable {
+
+ // MARK: - Mel Feature Extraction
+
+ /// Number of mel frequency bins
+ public let nMels: Int
+ /// FFT window size in samples
+ public let nFFT: Int
+ /// Hop length in samples
+ public let hopLength: Int
+ /// Expected input sample rate in Hz
+ public let sampleRate: Int
+
+ // MARK: - Streaming Chunking
+
+ /// Chunk length in seconds for streaming inference
+ public let chunkLenSeconds: Float
+ /// Left context in seconds (prepended from previous chunk)
+ public let leftContextSeconds: Float
+ /// Right context in seconds (lookahead)
+ public let rightContextSeconds: Float
+ /// Subsampling factor of the encoder (frames → mel frames)
+ public let subsamplingFactor: Int
+
+ // MARK: - State Dimensions
+
+ /// Speaker cache length (number of frames)
+ public let spkcacheLen: Int
+ /// FIFO buffer length (number of frames)
+ public let fifoLen: Int
+ /// Feature/hidden dimension of the model
+ public let fcDModel: Int
+
+ // MARK: - Model I/O Shapes
+
+ /// Maximum number of speakers the model can predict
+ public let maxSpeakers: Int
+
+ // MARK: - Post-processing
+
+ /// Onset threshold for speaker activity binarization
+ public var onset: Float
+ /// Offset threshold for speaker activity binarization
+ public var offset: Float
+ /// Minimum speech segment duration in seconds
+ public var minSpeechDuration: Float
+ /// Minimum silence gap to split segments, in seconds
+ public var minSilenceDuration: Float
+
+ // MARK: - Presets
+
+ /// Default streaming configuration matching the NeMo checkpoint.
+ public static let `default` = SortformerConfig(
+ nMels: 128,
+ nFFT: 400,
+ hopLength: 160,
+ sampleRate: 16000,
+ chunkLenSeconds: 6.0,
+ leftContextSeconds: 1.0,
+ rightContextSeconds: 7.0,
+ subsamplingFactor: 8,
+ spkcacheLen: 188,
+ fifoLen: 40,
+ fcDModel: 512,
+ maxSpeakers: 4,
+ onset: 0.5,
+ offset: 0.3,
+ minSpeechDuration: 0.3,
+ minSilenceDuration: 0.15
+ )
+
+ public init(
+ nMels: Int = 128,
+ nFFT: Int = 400,
+ hopLength: Int = 160,
+ sampleRate: Int = 16000,
+ chunkLenSeconds: Float = 6.0,
+ leftContextSeconds: Float = 1.0,
+ rightContextSeconds: Float = 7.0,
+ subsamplingFactor: Int = 8,
+ spkcacheLen: Int = 188,
+ fifoLen: Int = 40,
+ fcDModel: Int = 512,
+ maxSpeakers: Int = 4,
+ onset: Float = 0.5,
+ offset: Float = 0.3,
+ minSpeechDuration: Float = 0.3,
+ minSilenceDuration: Float = 0.15
+ ) {
+ self.nMels = nMels
+ self.nFFT = nFFT
+ self.hopLength = hopLength
+ self.sampleRate = sampleRate
+ self.chunkLenSeconds = chunkLenSeconds
+ self.leftContextSeconds = leftContextSeconds
+ self.rightContextSeconds = rightContextSeconds
+ self.subsamplingFactor = subsamplingFactor
+ self.spkcacheLen = spkcacheLen
+ self.fifoLen = fifoLen
+ self.fcDModel = fcDModel
+ self.maxSpeakers = maxSpeakers
+ self.onset = onset
+ self.offset = offset
+ self.minSpeechDuration = minSpeechDuration
+ self.minSilenceDuration = minSilenceDuration
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SortformerDiarizer.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SortformerDiarizer.swift
new file mode 100644
index 0000000..0da6d5a
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SortformerDiarizer.swift
@@ -0,0 +1,432 @@
+#if canImport(CoreML)
+import CoreML
+import Foundation
+import AudioCommon
+
+/// End-to-end neural speaker diarization using NVIDIA Sortformer (CoreML).
+///
+/// Sortformer directly predicts per-frame speaker activity for up to 4 speakers
+/// without requiring separate embedding extraction or clustering. Runs on
+/// Neural Engine at ~120x real-time.
+///
+/// ```swift
+/// let diarizer = try await SortformerDiarizer.fromPretrained()
+/// let result = diarizer.diarize(audio: samples, sampleRate: 16000)
+/// for seg in result.segments {
+/// print("Speaker \(seg.speakerId): [\(seg.startTime)s - \(seg.endTime)s]")
+/// }
+/// ```
+public final class SortformerDiarizer {
+
+ /// Default HuggingFace model ID for the CoreML Sortformer model
+ public static let defaultModelId = "aufklarer/Sortformer-Diarization-CoreML"
+
+ private let model: SortformerCoreMLModel
+ private let melExtractor: SortformerMelExtractor
+ let config: SortformerConfig
+
+ /// Frame duration from model metadata (0.08s = 80ms per diarization frame)
+ private let frameDuration: Float = 0.08
+
+ // MARK: - Streaming State
+
+ /// Speaker cache buffer, flat `[spkcacheLen * fcDModel]`
+ private var spkcache: [Float]
+ /// Number of valid frames in speaker cache
+ private var spkcacheLength: Int = 0
+ /// FIFO buffer, flat `[fifoLen * fcDModel]`
+ private var fifo: [Float]
+ /// Number of valid frames in FIFO
+ private var fifoLength: Int = 0
+ init(model: SortformerCoreMLModel, config: SortformerConfig = .default) {
+ self.model = model
+ self.config = config
+ self.melExtractor = SortformerMelExtractor(config: config)
+ self.spkcache = [Float](repeating: 0, count: config.spkcacheLen * config.fcDModel)
+ self.fifo = [Float](repeating: 0, count: config.fifoLen * config.fcDModel)
+ }
+
+ /// Reset streaming state between different audio files.
+ public func resetState() {
+ spkcache = [Float](repeating: 0, count: config.spkcacheLen * config.fcDModel)
+ spkcacheLength = 0
+ fifo = [Float](repeating: 0, count: config.fifoLen * config.fcDModel)
+ fifoLength = 0
+ }
+
+ // MARK: - Loading
+
+ /// Load a pre-trained Sortformer model from HuggingFace.
+ ///
+ /// - Parameters:
+ /// - modelId: HuggingFace model ID
+ /// - progressHandler: callback for download progress
+ /// - Returns: ready-to-use diarizer
+ public static func fromPretrained(
+ modelId: String = defaultModelId,
+ cacheDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> SortformerDiarizer {
+ progressHandler?(0.0, "Downloading Sortformer model...")
+
+ let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
+
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: modelId,
+ to: cacheDir,
+ additionalFiles: ["Sortformer.mlmodelc/**", "config.json"],
+ offlineMode: offlineMode,
+ progressHandler: { progress in
+ progressHandler?(progress * 0.8, "Downloading Sortformer model...")
+ }
+ )
+
+ progressHandler?(0.8, "Loading CoreML model...")
+
+ let modelURL = cacheDir.appendingPathComponent("Sortformer.mlmodelc", isDirectory: true)
+ guard FileManager.default.fileExists(atPath: modelURL.path) else {
+ throw AudioModelError.modelLoadFailed(
+ modelId: modelId,
+ reason: "CoreML model not found at \(modelURL.path)")
+ }
+
+ let mlConfig = MLModelConfiguration()
+ mlConfig.computeUnits = CoreMLComputeUnitsResolver.resolved(default: .cpuAndNeuralEngine)
+
+ let mlModel: MLModel
+ do {
+ mlModel = try MLModel(contentsOf: modelURL, configuration: mlConfig)
+ } catch {
+ throw AudioModelError.modelLoadFailed(
+ modelId: modelId,
+ reason: "Failed to load CoreML model",
+ underlying: error)
+ }
+
+ let config = SortformerConfig.default
+ let coremlModel = SortformerCoreMLModel(model: mlModel, config: config)
+
+ progressHandler?(1.0, "Ready")
+ return SortformerDiarizer(model: coremlModel, config: config)
+ }
+
+ // MARK: - Diarization
+
+ /// Run speaker diarization on complete audio.
+ ///
+ /// Processes audio in streaming chunks matching NeMo's streaming_feat_loader:
+ /// each chunk is 112 mel frames = (leftCtx + coreChunk + rightCtx) × subsampling.
+ /// Core predictions are extracted per chunk and concatenated.
+ ///
+ /// - Parameters:
+ /// - audio: PCM Float32 audio samples
+ /// - sampleRate: sample rate of the input audio
+ /// - config: optional override for diarization thresholds
+ /// - Returns: diarization result with speaker-labeled segments
+ public func diarize(
+ audio: [Float],
+ sampleRate: Int,
+ config: DiarizationConfig = .default
+ ) -> DiarizationResult {
+ diarize(audio: audio, sampleRate: sampleRate, config: config, progressHandler: nil)
+ }
+
+ /// Run speaker diarization with progress reporting and optional cancellation.
+ ///
+ /// Same as `diarize(audio:sampleRate:config:)` but reports progress per chunk.
+ /// The handler returns a `Bool`: `true` to continue, `false` to cancel.
+ /// When cancelled, an empty `DiarizationResult` is returned immediately.
+ ///
+ /// - Parameters:
+ /// - audio: PCM Float32 audio samples
+ /// - sampleRate: sample rate of the input audio
+ /// - config: optional override for diarization thresholds
+ /// - progressHandler: called with (progress 0.0–1.0, stage description);
+ /// return `true` to continue or `false` to cancel
+ /// - Returns: diarization result with speaker-labeled segments
+ public func diarize(
+ audio: [Float],
+ sampleRate: Int,
+ config: DiarizationConfig = .default,
+ progressHandler: ((Float, String) -> Bool)?
+ ) -> DiarizationResult {
+ let samples = DiarizationHelpers.resample(audio, from: sampleRate, to: self.config.sampleRate)
+
+ guard !samples.isEmpty else {
+ return DiarizationResult(segments: [], numSpeakers: 0, speakerEmbeddings: [])
+ }
+
+ resetState()
+
+ // Extract mel features for the entire audio: [totalMelFrames, 128]
+ let (melSpec, totalMelFrames) = melExtractor.extract(samples)
+
+ guard totalMelFrames > 0 else {
+ return DiarizationResult(segments: [], numSpeakers: 0, speakerEmbeddings: [])
+ }
+
+ // Streaming chunking parameters (matching NeMo)
+ let subFactor = self.config.subsamplingFactor
+ let chunkLen = Int(self.config.chunkLenSeconds) // 6 encoder output frames
+ let leftCtx = Int(self.config.leftContextSeconds) // 1
+ let rightCtx = Int(self.config.rightContextSeconds) // 7
+ let coreMelFrames = chunkLen * subFactor // 48 mel frames per core chunk
+ let coreMLInputFrames = 112 // Fixed CoreML input size
+ let nMels = self.config.nMels
+ let numSpeakers = self.config.maxSpeakers
+
+ // Collect core predictions from each chunk
+ var allChunkProbs = [[Float]]() // Each entry: [coreFrames * numSpeakers]
+ let emptyResult = DiarizationResult(segments: [], numSpeakers: 0, speakerEmbeddings: [])
+
+ // Calculate total chunks for progress reporting
+ let totalChunks = max(1, (totalMelFrames + coreMelFrames - 1) / coreMelFrames)
+ var chunkIndex = 0
+
+ var sttFeat = 0
+ var endFeat = 0
+
+ while endFeat < totalMelFrames {
+ chunkIndex += 1
+ if progressHandler?(Float(chunkIndex) / Float(totalChunks), "Diarizing \(chunkIndex)/\(totalChunks)") == false {
+ return emptyResult
+ }
+ let leftOffset = min(leftCtx * subFactor, sttFeat)
+ endFeat = min(sttFeat + coreMelFrames, totalMelFrames)
+ let rightOffset = min(rightCtx * subFactor, totalMelFrames - endFeat)
+
+ let chunkStart = sttFeat - leftOffset
+ let chunkEnd = endFeat + rightOffset
+ let actualLen = chunkEnd - chunkStart
+
+ // Build padded mel chunk [coreMLInputFrames, nMels]
+ var chunkMel = [Float](repeating: 0, count: coreMLInputFrames * nMels)
+ let framesToCopy = min(actualLen, coreMLInputFrames)
+ for fi in 0.. 0 ? coreLen : 0
+
+ let predOffset = spkcacheLength + fifoLength + lcFrames
+ let totalPredFrames = output.predsFrames
+
+ var chunkProbs = [Float]()
+ for f in 0.. 0 else { return }
+ let dim = config.fcDModel
+ let fifoCapacity = config.fifoLen
+ let cacheCapacity = config.spkcacheLen
+
+ if fifoLength + validFrames <= fifoCapacity {
+ // FIFO has room — just append
+ for f in 0.. 0 {
+ for f in 0.. 0 {
+ for f in 0.. [DiarizedSegment] {
+ // Concatenate all chunk predictions into one flat array
+ var allProbs = [Float]()
+ for chunkProbs in allChunkProbs {
+ allProbs.append(contentsOf: chunkProbs)
+ }
+
+ let totalFrames = allProbs.count / numSpeakers
+ guard totalFrames > 0 else { return [] }
+
+ // Apply sigmoid if predictions are logits
+ for i in 0.. 1.0 || allProbs[i] < 0.0 {
+ allProbs[i] = 1.0 / (1.0 + exp(-allProbs[i]))
+ }
+ }
+
+ // Binarize each speaker track
+ var allSegments = [DiarizedSegment]()
+
+ for spk in 0..= minSpeechDuration else { continue }
+ allSegments.append(DiarizedSegment(
+ startTime: seg.startTime,
+ endTime: min(seg.endTime, audioDuration),
+ speakerId: spk
+ ))
+ }
+ }
+
+ allSegments.sort { $0.startTime < $1.startTime }
+ let merged = DiarizationHelpers.mergeSegments(allSegments, minSilence: minSilenceDuration)
+ return DiarizationHelpers.compactSpeakerIds(merged)
+ }
+}
+
+// MARK: - SpeakerDiarizationModel
+
+extension SortformerDiarizer: SpeakerDiarizationModel {
+ public var inputSampleRate: Int { config.sampleRate }
+
+ public func diarize(audio: [Float], sampleRate: Int) -> [DiarizedSegment] {
+ diarize(audio: audio, sampleRate: sampleRate, config: .default).segments
+ }
+}
+#endif
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SortformerMelExtractor.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SortformerMelExtractor.swift
new file mode 100644
index 0000000..7fef0e8
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SortformerMelExtractor.swift
@@ -0,0 +1,205 @@
+import Foundation
+import Accelerate
+
+/// 128-dim log-mel feature extractor for Sortformer diarization.
+///
+/// Matches NeMo's audio preprocessor: Hann window (no Povey), no pre-emphasis,
+/// nFFT=400, hop=160, 128 mel bins, 16kHz. Uses vDSP for FFT and mel filterbank.
+///
+/// Key differences from `MelFeatureExtractor` (WeSpeaker):
+/// - 128 mel bins (vs 80)
+/// - Hann window (vs Povey window)
+/// - No pre-emphasis (vs 0.97)
+/// - Power spectrum (vs magnitude spectrum)
+class SortformerMelExtractor {
+ let sampleRate: Int
+ let nFFT: Int
+ let hopLength: Int
+ let nMels: Int
+
+ private let paddedFFT: Int = 512
+ private let log2PaddedFFT: vDSP_Length = 9
+ private var fftSetup: FFTSetup
+ private var window: [Float]
+ private var melFilterbank: [Float] // [nMels, nBins]
+
+ init(config: SortformerConfig = .default) {
+ self.sampleRate = config.sampleRate
+ self.nFFT = config.nFFT
+ self.hopLength = config.hopLength
+ self.nMels = config.nMels
+
+ // Hann window (NeMo default, no Povey modification)
+ window = [Float](repeating: 0, count: config.nFFT)
+ for i in 0.. Float {
+ 2595.0 * log10(1.0 + hz / 700.0)
+ }
+
+ func melToHz(_ mel: Float) -> Float {
+ 700.0 * (pow(10.0, mel / 2595.0) - 1.0)
+ }
+
+ let nBins = paddedFFT / 2 + 1 // 257
+
+ var fftFreqs = [Float](repeating: 0, count: nBins)
+ for i in 0.. (melSpec: [Float], nFrames: Int) {
+ let nBins = paddedFFT / 2 + 1
+ let halfPadded = paddedFFT / 2
+
+ // No pre-emphasis for Sortformer (NeMo default)
+
+ guard !audio.isEmpty else { return ([], 0) }
+
+ // Reflect padding (same as torch.stft with center=True)
+ let padLength = nFFT / 2
+ var paddedAudio = [Float](repeating: 0, count: padLength + audio.count + padLength)
+
+ for i in 0.. SortformerOutput {
+ // Create input arrays
+ let chunkArray = try makeMultiArray(
+ shape: [1, NSNumber(value: chunkFrames), NSNumber(value: config.nMels)],
+ from: chunk)
+ let chunkLenArray = try makeScalarInt32Array(value: Int32(chunkLength))
+
+ let spkcacheArray = try makeMultiArray(
+ shape: [1, NSNumber(value: spkcacheFrames), NSNumber(value: featureDim)],
+ from: spkcache)
+ let spkcacheLenArray = try makeScalarInt32Array(value: Int32(spkcacheLength))
+
+ let fifoArray = try makeMultiArray(
+ shape: [1, NSNumber(value: fifoFrames), NSNumber(value: featureDim)],
+ from: fifo)
+ let fifoLenArray = try makeScalarInt32Array(value: Int32(fifoLength))
+
+ let input = try MLDictionaryFeatureProvider(dictionary: [
+ "chunk": MLFeatureValue(multiArray: chunkArray),
+ "chunk_lengths": MLFeatureValue(multiArray: chunkLenArray),
+ "spkcache": MLFeatureValue(multiArray: spkcacheArray),
+ "spkcache_lengths": MLFeatureValue(multiArray: spkcacheLenArray),
+ "fifo": MLFeatureValue(multiArray: fifoArray),
+ "fifo_lengths": MLFeatureValue(multiArray: fifoLenArray),
+ ])
+
+ let result = try model.prediction(from: input)
+
+ // Extract outputs
+ let predsArray = result.featureValue(for: "speaker_preds_out")!.multiArrayValue!
+ let embsArray = result.featureValue(for: "chunk_pre_encoder_embs_out")!.multiArrayValue!
+ let embsLenArray = result.featureValue(for: "chunk_pre_encoder_lengths_out")!.multiArrayValue!
+
+ let predsShape = (0..= 2 ? predsShape[predsShape.count - 2] : totalPreds / config.maxSpeakers,
+ numSpeakers: config.maxSpeakers,
+ encoderEmbs: embs,
+ encoderEmbFrames: embsShape.count >= 2 ? embsShape[embsShape.count - 2] : validEmbFrames,
+ embDim: featureDim,
+ validEmbFrames: validEmbFrames
+ )
+ }
+
+ // MARK: - Helpers
+
+ private func makeMultiArray(shape: [NSNumber], from data: [Float]) throws -> MLMultiArray {
+ let array = try MLMultiArray(shape: shape, dataType: .float32)
+ let ptr = array.dataPointer.assumingMemoryBound(to: Float.self)
+ let count = min(data.count, array.count)
+ for i in 0.. MLMultiArray {
+ let array = try MLMultiArray(shape: [1], dataType: .int32)
+ let ptr = array.dataPointer.assumingMemoryBound(to: Int32.self)
+ ptr[0] = value
+ return array
+ }
+}
+
+/// Output from one Sortformer inference step.
+struct SortformerOutput {
+ /// Speaker predictions, flat `[predsFrames * numSpeakers]`, sigmoid probabilities
+ let speakerPreds: [Float]
+ /// Number of prediction frames
+ let predsFrames: Int
+ /// Number of speaker channels
+ let numSpeakers: Int
+ /// Pre-encoder embeddings for state update, flat `[encoderEmbFrames * embDim]`
+ let encoderEmbs: [Float]
+ /// Total encoder embedding frames
+ let encoderEmbFrames: Int
+ /// Embedding dimension
+ let embDim: Int
+ /// Number of valid (non-padding) embedding frames
+ let validEmbFrames: Int
+
+ /// Get speaker prediction probability at (frame, speaker).
+ func pred(frame: Int, speaker: Int) -> Float {
+ speakerPreds[frame * numSpeakers + speaker]
+ }
+
+ /// Get encoder embedding at (frame, dim).
+ func emb(frame: Int, dim: Int) -> Float {
+ encoderEmbs[frame * embDim + dim]
+ }
+}
+#endif
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SpeechVAD+Protocols.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SpeechVAD+Protocols.swift
new file mode 100644
index 0000000..e82f2bd
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SpeechVAD+Protocols.swift
@@ -0,0 +1,29 @@
+import AudioCommon
+
+// MARK: - VoiceActivityDetectionModel
+
+extension PyannoteVADModel: VoiceActivityDetectionModel {
+ public var inputSampleRate: Int { segConfig.sampleRate }
+}
+
+// MARK: - SpeakerEmbeddingModel
+
+extension WeSpeakerModel: SpeakerEmbeddingModel {}
+
+// MARK: - SpeakerDiarizationModel
+
+extension PyannoteDiarizationPipeline: SpeakerDiarizationModel {
+ public var inputSampleRate: Int { segConfig.sampleRate }
+
+ public func diarize(audio: [Float], sampleRate: Int) -> [DiarizedSegment] {
+ diarize(audio: audio, sampleRate: sampleRate, config: .default).segments
+ }
+}
+
+// MARK: - SpeakerExtractionCapable
+
+extension PyannoteDiarizationPipeline: SpeakerExtractionCapable {
+ public func extractSpeaker(audio: [Float], sampleRate: Int, targetEmbedding: [Float]) -> [SpeechSegment] {
+ extractSpeaker(audio: audio, sampleRate: sampleRate, targetEmbedding: targetEmbedding, config: .default)
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SpeechVAD.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SpeechVAD.swift
new file mode 100644
index 0000000..df6f945
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SpeechVAD.swift
@@ -0,0 +1,142 @@
+import Foundation
+import MLXCommon
+import MLX
+import AudioCommon
+
+/// Voice Activity Detection using pyannote PyanNet segmentation.
+///
+/// Detects speech regions in audio using a sliding-window segmentation model
+/// with hysteresis thresholding and duration filtering.
+///
+/// - Warning: This class is not thread-safe. Create separate instances for concurrent use.
+///
+/// ```swift
+/// let vad = try await PyannoteVADModel.fromPretrained()
+/// let segments = vad.detectSpeech(audio: samples, sampleRate: 16000)
+/// for seg in segments {
+/// print("Speech: \(seg.startTime)s - \(seg.endTime)s")
+/// }
+/// ```
+public final class PyannoteVADModel {
+ /// The segmentation model
+ let model: SegmentationModel
+
+ /// VAD pipeline configuration
+ public let vadConfig: VADConfig
+
+ /// Segmentation model configuration
+ public let segConfig: SegmentationConfig
+
+ /// Default HuggingFace model ID
+ public static let defaultModelId = "aufklarer/Pyannote-Segmentation-MLX"
+
+ /// Whether the model weights are loaded and ready for inference.
+ var _isLoaded = true
+
+ init(model: SegmentationModel, segConfig: SegmentationConfig, vadConfig: VADConfig) {
+ self.model = model
+ self.segConfig = segConfig
+ self.vadConfig = vadConfig
+ }
+
+ /// Load a pre-trained VAD model from HuggingFace.
+ ///
+ /// Downloads model weights on first use, then caches locally.
+ ///
+ /// - Parameters:
+ /// - modelId: HuggingFace model ID
+ /// - vadConfig: VAD pipeline configuration (thresholds, durations)
+ /// - progressHandler: callback for download progress
+ /// - Returns: ready-to-use VAD model
+ public static func fromPretrained(
+ modelId: String = defaultModelId,
+ vadConfig: VADConfig = .default,
+ cacheDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> PyannoteVADModel {
+ progressHandler?(0.0, "Downloading model...")
+
+ let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: modelId)
+
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: modelId,
+ to: cacheDir,
+ offlineMode: offlineMode,
+ progressHandler: { progress in
+ progressHandler?(progress * 0.8, "Downloading weights...")
+ }
+ )
+
+ progressHandler?(0.8, "Loading model...")
+
+ let segConfig = SegmentationConfig.default
+ let model = SegmentationModel(config: segConfig)
+
+ try SegmentationWeightLoader.loadWeights(model: model, from: cacheDir)
+
+ progressHandler?(1.0, "Ready")
+
+ return PyannoteVADModel(model: model, segConfig: segConfig, vadConfig: vadConfig)
+ }
+
+ /// Detect speech segments in audio.
+ ///
+ /// - Parameters:
+ /// - audio: PCM Float32 audio samples
+ /// - sampleRate: sample rate of the input audio (will resample to 16kHz if needed)
+ /// - Returns: array of speech segments with start/end times in seconds
+ public func detectSpeech(audio: [Float], sampleRate: Int) -> [SpeechSegment] {
+ let samples: [Float]
+ if sampleRate != segConfig.sampleRate {
+ samples = AudioFileLoader.resample(audio, from: sampleRate, to: segConfig.sampleRate)
+ } else {
+ samples = audio
+ }
+
+ let pipeline = VADPipeline(
+ config: vadConfig,
+ sampleRate: segConfig.sampleRate,
+ framesPerChunk: 589
+ )
+
+ let positions = pipeline.windowPositions(numSamples: samples.count)
+
+ guard !positions.isEmpty else { return [] }
+
+ let windowSamples = Int(vadConfig.windowDuration * Float(segConfig.sampleRate))
+
+ // Run segmentation on each window
+ var windowProbs = [[Float]]()
+
+ for (start, end) in positions {
+ // Extract window, zero-pad if needed
+ var window = Array(samples[start ..< end])
+ if window.count < windowSamples {
+ window.append(contentsOf: [Float](repeating: 0, count: windowSamples - window.count))
+ }
+
+ // Run model: [1, 1, samples] → [1, frames, 7]
+ let input = MLXArray(window).reshaped(1, 1, windowSamples)
+ let posteriors = model(input)
+
+ // Extract speech probability: [1, frames] → [frames]
+ let speechProb = SegmentationModel.speechProbability(from: posteriors)
+ eval(speechProb)
+
+ let probArray = speechProb[0].asArray(Float.self)
+ windowProbs.append(probArray)
+ }
+
+ // Aggregate overlapping windows
+ let aggregated = pipeline.aggregateFrames(
+ windowProbs: windowProbs,
+ positions: positions,
+ numSamples: samples.count
+ )
+
+ // Binarize with hysteresis
+ return pipeline.binarize(probs: aggregated)
+ }
+
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/StreamingVADProcessor.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/StreamingVADProcessor.swift
new file mode 100644
index 0000000..ccbe523
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/StreamingVADProcessor.swift
@@ -0,0 +1,210 @@
+import Foundation
+import AudioCommon
+
+/// Events emitted by the streaming VAD processor.
+public enum VADEvent: Sendable {
+ /// Speech has been detected and confirmed (duration ≥ minSpeechDuration).
+ case speechStarted(time: Float)
+ /// Speech has ended (silence ≥ minSilenceDuration).
+ case speechEnded(segment: SpeechSegment)
+}
+
+/// Event-driven streaming VAD processor.
+///
+/// Wraps a `SileroVADModel` to provide event-based speech detection.
+/// Accepts audio samples of any length, buffers them into 512-sample chunks,
+/// runs the model, and applies hysteresis with duration filtering via a
+/// four-state machine.
+///
+/// - Warning: This class is not thread-safe. Create separate instances for concurrent use.
+///
+/// ```swift
+/// let model = try await SileroVADModel.fromPretrained()
+/// let processor = StreamingVADProcessor(model: model)
+///
+/// // Feed audio samples (any length)
+/// let events = processor.process(samples: audioBuffer)
+/// for event in events {
+/// switch event {
+/// case .speechStarted(let time):
+/// print("Speech started at \(time)s")
+/// case .speechEnded(let segment):
+/// print("Speech: \(segment.startTime)s - \(segment.endTime)s")
+/// }
+/// }
+///
+/// // At end of stream, flush any pending segment
+/// let finalEvents = processor.flush()
+/// ```
+public final class StreamingVADProcessor {
+
+ private let model: SileroVADModel
+ private let config: VADConfig
+ private let chunkDuration: Float // seconds per chunk (0.032)
+
+ /// Buffer for accumulating samples until we have a full chunk
+ private var buffer: [Float] = []
+ /// Number of chunks processed so far
+ private var chunkCount: Int = 0
+
+ /// State machine for hysteresis + duration filtering
+ private enum State {
+ /// No speech detected
+ case silence
+ /// Onset threshold crossed, waiting for minSpeechDuration
+ case pendingSpeech(startTime: Float)
+ /// Speech confirmed and speechStarted emitted
+ case speech(startTime: Float)
+ /// Offset threshold crossed, waiting for minSilenceDuration
+ case pendingSilence(speechStart: Float, silenceStart: Float)
+ }
+
+ private var state: State = .silence
+
+ /// Create a streaming VAD processor.
+ ///
+ /// - Parameters:
+ /// - model: Silero VAD model instance
+ /// - config: VAD configuration (thresholds, durations)
+ public init(model: SileroVADModel, config: VADConfig = .sileroDefault) {
+ self.model = model
+ self.config = config
+ self.chunkDuration = Float(SileroVADModel.chunkSize) / Float(SileroVADModel.sampleRate)
+ }
+
+ /// Feed audio samples and get VAD events back.
+ ///
+ /// Samples are buffered internally. Events are emitted as soon as the
+ /// state machine confirms speech start/end with the configured thresholds
+ /// and duration constraints.
+ ///
+ /// - Parameter samples: PCM Float32 samples at 16kHz (any length)
+ /// - Returns: zero or more VAD events
+ public func process(samples: [Float]) -> [VADEvent] {
+ buffer.append(contentsOf: samples)
+ var events = [VADEvent]()
+
+ while buffer.count >= SileroVADModel.chunkSize {
+ let chunk = Array(buffer.prefix(SileroVADModel.chunkSize))
+ buffer.removeFirst(SileroVADModel.chunkSize)
+
+ let prob = model.processChunk(chunk)
+ let time = Float(chunkCount) * chunkDuration
+ chunkCount += 1
+
+ events.append(contentsOf: processProb(prob, time: time))
+ }
+
+ return events
+ }
+
+ /// Flush any pending speech segment at end of stream.
+ ///
+ /// Call this when the audio stream ends to close any open speech segment.
+ ///
+ /// - Returns: zero or more final VAD events
+ public func flush() -> [VADEvent] {
+ // Process any remaining buffered samples (zero-padded)
+ var events = [VADEvent]()
+ if !buffer.isEmpty {
+ var lastChunk = buffer
+ lastChunk.append(contentsOf: [Float](repeating: 0, count: SileroVADModel.chunkSize - lastChunk.count))
+ buffer.removeAll()
+
+ let prob = model.processChunk(lastChunk)
+ let time = Float(chunkCount) * chunkDuration
+ chunkCount += 1
+ events.append(contentsOf: processProb(prob, time: time))
+ }
+
+ let endTime = Float(chunkCount) * chunkDuration
+
+ // Close any open state
+ switch state {
+ case .silence:
+ break
+ case .pendingSpeech(let startTime):
+ // Check if pending speech meets minimum duration
+ if endTime - startTime >= config.minSpeechDuration {
+ events.append(.speechStarted(time: startTime))
+ events.append(.speechEnded(segment: SpeechSegment(
+ startTime: startTime, endTime: endTime)))
+ }
+ case .speech(let startTime):
+ events.append(.speechEnded(segment: SpeechSegment(
+ startTime: startTime, endTime: endTime)))
+ case .pendingSilence(let speechStart, let silenceStart):
+ // End at the silence start point
+ events.append(.speechEnded(segment: SpeechSegment(
+ startTime: speechStart, endTime: silenceStart)))
+ }
+
+ state = .silence
+ return events
+ }
+
+ /// Reset all state (model + processor).
+ ///
+ /// Call between processing different audio streams.
+ public func reset() {
+ buffer.removeAll()
+ chunkCount = 0
+ state = .silence
+ model.resetState()
+ }
+
+ /// Current time position in seconds.
+ public var currentTime: Float {
+ Float(chunkCount) * chunkDuration
+ }
+
+ // MARK: - State Machine
+
+ private func processProb(_ prob: Float, time: Float) -> [VADEvent] {
+ var events = [VADEvent]()
+ let nextTime = time + chunkDuration
+
+ switch state {
+ case .silence:
+ if prob >= config.onset {
+ state = .pendingSpeech(startTime: time)
+ }
+
+ case .pendingSpeech(let startTime):
+ if prob < config.offset {
+ // False alarm — speech too brief, return to silence
+ state = .silence
+ } else if nextTime - startTime >= config.minSpeechDuration {
+ // Speech confirmed
+ events.append(.speechStarted(time: startTime))
+ state = .speech(startTime: startTime)
+ }
+ // else: still pending, keep waiting
+
+ case .speech(let startTime):
+ if prob < config.offset {
+ // Speech may be ending
+ state = .pendingSilence(speechStart: startTime, silenceStart: time)
+ }
+
+ case .pendingSilence(let speechStart, let silenceStart):
+ if prob >= config.onset {
+ // Speech resumed — cancel silence
+ state = .speech(startTime: speechStart)
+ } else if nextTime - silenceStart >= config.minSilenceDuration {
+ // Silence confirmed — emit speechEnded
+ events.append(.speechEnded(segment: SpeechSegment(
+ startTime: speechStart, endTime: silenceStart)))
+ // Check if new speech is starting
+ if prob >= config.onset {
+ state = .pendingSpeech(startTime: time)
+ } else {
+ state = .silence
+ }
+ }
+ // else: still waiting for silence confirmation
+ }
+
+ return events
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/VADPipeline.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/VADPipeline.swift
new file mode 100644
index 0000000..7904ad3
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/VADPipeline.swift
@@ -0,0 +1,181 @@
+import Foundation
+import MLX
+import AudioCommon
+
+/// VAD pipeline: sliding window segmentation → aggregation → binarization.
+///
+/// Processes audio in overlapping 10-second windows, runs the segmentation model
+/// on each window, aggregates overlapping frame predictions, then applies
+/// hysteresis thresholding and duration filtering to produce speech segments.
+public struct VADPipeline: Sendable {
+
+ /// Configuration for the pipeline
+ public let config: VADConfig
+
+ /// Sample rate expected by the segmentation model
+ public let sampleRate: Int
+
+ /// Number of output frames per 10s chunk (from the segmentation model)
+ public let framesPerChunk: Int
+
+ public init(config: VADConfig = .default, sampleRate: Int = 16000, framesPerChunk: Int = 589) {
+ self.config = config
+ self.sampleRate = sampleRate
+ self.framesPerChunk = framesPerChunk
+ }
+
+ /// Duration of one frame in seconds.
+ public var frameDuration: Float {
+ config.windowDuration / Float(framesPerChunk)
+ }
+
+ // MARK: - Sliding Window
+
+ /// Generate sliding window positions for the given audio length.
+ /// - Parameter numSamples: total audio samples
+ /// - Returns: array of (start, end) sample indices
+ func windowPositions(numSamples: Int) -> [(start: Int, end: Int)] {
+ let windowSamples = Int(config.windowDuration * Float(sampleRate))
+ let stepSamples = Int(config.windowDuration * config.stepRatio * Float(sampleRate))
+
+ guard numSamples > 0 else { return [] }
+
+ // If audio is shorter than one window, just use one window (zero-padded)
+ if numSamples <= windowSamples {
+ return [(0, numSamples)]
+ }
+
+ var positions = [(start: Int, end: Int)]()
+ var start = 0
+ while start + windowSamples <= numSamples {
+ positions.append((start, start + windowSamples))
+ start += stepSamples
+ }
+ // Handle the last partial window
+ if positions.isEmpty || positions.last!.end < numSamples {
+ positions.append((numSamples - windowSamples, numSamples))
+ }
+
+ return positions
+ }
+
+ // MARK: - Frame Aggregation
+
+ /// Aggregate overlapping frame-level speech probabilities from multiple windows.
+ ///
+ /// Each window produces `framesPerChunk` frames. Overlapping regions are
+ /// averaged across windows.
+ ///
+ /// - Parameters:
+ /// - windowProbs: array of per-window speech probability arrays (each `framesPerChunk` long)
+ /// - positions: corresponding window positions (sample indices)
+ /// - numSamples: total audio length in samples
+ /// - Returns: aggregated speech probability per frame for the entire audio
+ func aggregateFrames(
+ windowProbs: [[Float]],
+ positions: [(start: Int, end: Int)],
+ numSamples: Int
+ ) -> [Float] {
+ let totalDuration = Float(numSamples) / Float(sampleRate)
+ let numFrames = Int(ceil(totalDuration / frameDuration))
+
+ guard numFrames > 0 else { return [] }
+
+ var sumProbs = [Float](repeating: 0, count: numFrames)
+ var counts = [Float](repeating: 0, count: numFrames)
+
+ for (windowIdx, (start, _)) in positions.enumerated() {
+ let probs = windowProbs[windowIdx]
+ let windowStartTime = Float(start) / Float(sampleRate)
+
+ for (frameIdx, prob) in probs.enumerated() {
+ let frameTime = windowStartTime + Float(frameIdx) * frameDuration
+ let globalFrame = Int(frameTime / frameDuration)
+
+ if globalFrame >= 0 && globalFrame < numFrames {
+ sumProbs[globalFrame] += prob
+ counts[globalFrame] += 1
+ }
+ }
+ }
+
+ // Average where we have overlapping windows
+ return zip(sumProbs, counts).map { sum, count in
+ count > 0 ? sum / count : 0
+ }
+ }
+
+ // MARK: - Binarization (Hysteresis Thresholding)
+
+ /// Apply hysteresis thresholding to speech probabilities.
+ ///
+ /// Speech starts when probability exceeds `onset` and ends when it drops
+ /// below `offset`. This two-threshold approach prevents rapid toggling.
+ ///
+ /// - Parameter probs: per-frame speech probabilities
+ /// - Returns: array of `SpeechSegment` with start/end times
+ public func binarize(probs: [Float]) -> [SpeechSegment] {
+ var segments = [SpeechSegment]()
+ var inSpeech = false
+ var speechStart: Float = 0
+
+ for (i, prob) in probs.enumerated() {
+ let time = Float(i) * frameDuration
+
+ if !inSpeech && prob >= config.onset {
+ inSpeech = true
+ speechStart = time
+ } else if inSpeech && prob < config.offset {
+ inSpeech = false
+ let segment = SpeechSegment(startTime: speechStart, endTime: time)
+ segments.append(segment)
+ }
+ }
+
+ // Close any open segment
+ if inSpeech {
+ let endTime = Float(probs.count) * frameDuration
+ segments.append(SpeechSegment(startTime: speechStart, endTime: endTime))
+ }
+
+ return filterDurations(segments)
+ }
+
+ // MARK: - Duration Filtering
+
+ /// Filter segments by minimum speech and silence durations.
+ ///
+ /// 1. Remove speech segments shorter than `minSpeechDuration`
+ /// 2. Merge segments separated by silence shorter than `minSilenceDuration`
+ func filterDurations(_ segments: [SpeechSegment]) -> [SpeechSegment] {
+ guard !segments.isEmpty else { return [] }
+
+ // Filter short speech segments
+ let filtered = segments.filter { $0.duration >= config.minSpeechDuration }
+
+ guard !filtered.isEmpty else { return [] }
+
+ // Merge segments separated by short silence
+ var merged = [SpeechSegment]()
+ var current = filtered[0]
+
+ for i in 1 ..< filtered.count {
+ let next = filtered[i]
+ let gap = next.startTime - current.endTime
+
+ if gap < config.minSilenceDuration {
+ // Merge: extend current segment
+ current = SpeechSegment(
+ startTime: current.startTime,
+ endTime: next.endTime
+ )
+ } else {
+ merged.append(current)
+ current = next
+ }
+ }
+ merged.append(current)
+
+ return merged
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeaker+Memory.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeaker+Memory.swift
new file mode 100644
index 0000000..020048a
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeaker+Memory.swift
@@ -0,0 +1,19 @@
+import AudioCommon
+
+extension WeSpeakerModel: ModelMemoryManageable {
+ public var isLoaded: Bool { _isLoaded }
+
+ public func unload() {
+ guard _isLoaded else { return }
+ network?.clearParameters()
+ #if canImport(CoreML)
+ coremlModel = nil
+ #endif
+ _isLoaded = false
+ }
+
+ public var memoryFootprint: Int {
+ guard _isLoaded else { return 0 }
+ return network?.parameterMemoryBytes() ?? 0
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeaker.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeaker.swift
new file mode 100644
index 0000000..751a107
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeaker.swift
@@ -0,0 +1,231 @@
+import Foundation
+import MLXCommon
+import MLX
+import AudioCommon
+
+#if canImport(CoreML)
+import CoreML
+#endif
+
+/// Inference engine for WeSpeaker speaker embeddings.
+public enum WeSpeakerEngine: String, Sendable {
+ /// MLX backend — runs on GPU via Metal shaders.
+ case mlx
+ /// CoreML backend — runs on Neural Engine + CPU, freeing the GPU.
+ case coreml
+}
+
+/// Speaker embedding model using WeSpeaker ResNet34-LM.
+///
+/// Produces 256-dimensional L2-normalized speaker embeddings from audio.
+/// Uses 80-dim log-mel features at 16kHz.
+///
+/// Supports two backends:
+/// - `.mlx`: GPU-based inference via MLX (default)
+/// - `.coreml`: Neural Engine inference via CoreML (lower power, frees GPU)
+///
+/// This class is thread-safe: all properties are immutable after construction and
+/// inference is pure computation with no mutable state. `MLModel.prediction(from:)` is
+/// documented as thread-safe by Apple.
+///
+/// ```swift
+/// let model = try await WeSpeakerModel.fromPretrained(engine: .coreml)
+/// let embedding = model.embed(audio: samples, sampleRate: 16000)
+/// // embedding: [Float] of length 256
+/// ```
+public final class WeSpeakerModel {
+
+ /// The inference engine in use.
+ public let engine: WeSpeakerEngine
+
+ /// Whether the model weights are loaded and ready for inference.
+ var _isLoaded = true
+
+ /// The ResNet34 network (nil when using CoreML engine)
+ let network: WeSpeakerNetwork?
+
+ /// Mel feature extractor
+ let melExtractor: MelFeatureExtractor
+
+ #if canImport(CoreML)
+ /// CoreML compiled model (nil when using MLX engine)
+ var coremlModel: MLModel?
+ #endif
+
+ /// Default HuggingFace model ID (MLX weights)
+ public static let defaultModelId = "aufklarer/WeSpeaker-ResNet34-LM-MLX"
+
+ /// Default HuggingFace model ID (CoreML weights)
+ public static let defaultCoreMLModelId = "aufklarer/WeSpeaker-ResNet34-LM-CoreML"
+
+ /// Embedding dimension
+ public let embeddingDimension: Int = 256
+
+ /// Expected input sample rate
+ public let inputSampleRate: Int = 16000
+
+ /// Enumerated mel frame lengths supported by the CoreML model.
+ static let enumeratedMelLengths = [20, 50, 100, 200, 300, 500, 750, 1000, 1500, 2000]
+
+ init(network: WeSpeakerNetwork) {
+ self.engine = .mlx
+ self.network = network
+ self.melExtractor = MelFeatureExtractor()
+ #if canImport(CoreML)
+ self.coremlModel = nil
+ #endif
+ }
+
+ #if canImport(CoreML)
+ init(coremlModel: MLModel) {
+ self.engine = .coreml
+ self.network = nil
+ self.coremlModel = coremlModel
+ self.melExtractor = MelFeatureExtractor()
+ }
+ #endif
+
+ /// Load a pre-trained speaker embedding model from HuggingFace.
+ ///
+ /// Downloads model weights on first use, then caches locally.
+ ///
+ /// - Parameters:
+ /// - modelId: HuggingFace model ID (auto-selected by engine if not specified)
+ /// - engine: inference backend (`.mlx` or `.coreml`)
+ /// - progressHandler: callback for download progress
+ /// - Returns: ready-to-use speaker embedding model
+ public static func fromPretrained(
+ modelId: String? = nil,
+ engine: WeSpeakerEngine = .mlx,
+ cacheDir: URL? = nil,
+ offlineMode: Bool = false,
+ progressHandler: ((Double, String) -> Void)? = nil
+ ) async throws -> WeSpeakerModel {
+ let resolvedModelId = modelId ?? (engine == .coreml ? defaultCoreMLModelId : defaultModelId)
+
+ progressHandler?(0.0, "Downloading speaker embedding model...")
+
+ let cacheDir = try cacheDir ?? HuggingFaceDownloader.getCacheDirectory(for: resolvedModelId)
+
+ switch engine {
+ case .mlx:
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: resolvedModelId,
+ to: cacheDir,
+ offlineMode: offlineMode,
+ progressHandler: { progress in
+ progressHandler?(progress * 0.8, "Downloading weights...")
+ }
+ )
+
+ progressHandler?(0.8, "Loading model...")
+
+ let network = WeSpeakerNetwork()
+ try WeSpeakerWeightLoader.loadWeights(model: network, from: cacheDir)
+
+ progressHandler?(1.0, "Ready")
+ return WeSpeakerModel(network: network)
+
+ case .coreml:
+ #if canImport(CoreML)
+ try await HuggingFaceDownloader.downloadWeights(
+ modelId: resolvedModelId,
+ to: cacheDir,
+ additionalFiles: ["wespeaker.mlmodelc/**", "config.json"],
+ offlineMode: offlineMode,
+ progressHandler: { progress in
+ progressHandler?(progress * 0.8, "Downloading CoreML model...")
+ }
+ )
+
+ progressHandler?(0.8, "Loading CoreML model...")
+
+ let modelURL = cacheDir.appendingPathComponent("wespeaker.mlmodelc", isDirectory: true)
+ guard FileManager.default.fileExists(atPath: modelURL.path) else {
+ throw AudioModelError.modelLoadFailed(
+ modelId: resolvedModelId,
+ reason: "CoreML model not found at \(modelURL.path)")
+ }
+
+ let mlConfig = MLModelConfiguration()
+ mlConfig.computeUnits = CoreMLComputeUnitsResolver.resolved(default: .cpuAndNeuralEngine)
+
+ let model: MLModel
+ do {
+ model = try MLModel(contentsOf: modelURL, configuration: mlConfig)
+ } catch {
+ throw AudioModelError.modelLoadFailed(
+ modelId: resolvedModelId,
+ reason: "Failed to load CoreML model",
+ underlying: error)
+ }
+
+ progressHandler?(1.0, "Ready")
+ return WeSpeakerModel(coremlModel: model)
+ #else
+ throw AudioModelError.invalidConfiguration(
+ model: "WeSpeaker", reason: "CoreML not available on this platform")
+ #endif
+ }
+ }
+
+ /// Extract a 256-dimensional speaker embedding from audio.
+ ///
+ /// - Parameters:
+ /// - audio: PCM Float32 audio samples
+ /// - sampleRate: sample rate of the input audio
+ /// - Returns: 256-dim L2-normalized speaker embedding
+ public func embed(audio: [Float], sampleRate: Int) -> [Float] {
+ let samples: [Float]
+ if sampleRate != inputSampleRate {
+ samples = AudioFileLoader.resample(audio, from: sampleRate, to: inputSampleRate)
+ } else {
+ samples = audio
+ }
+
+ switch engine {
+ case .mlx:
+ return embedMLX(samples)
+ case .coreml:
+ #if canImport(CoreML)
+ let (melSpec, nFrames) = melExtractor.extractRaw(samples)
+ return (try? embedCoreML(melSpec: melSpec, nFrames: nFrames)) ?? [Float](repeating: 0, count: embeddingDimension)
+ #else
+ fatalError("CoreML not available on this platform")
+ #endif
+ }
+ }
+
+ /// MLX inference path.
+ private func embedMLX(_ samples: [Float]) -> [Float] {
+ guard let network else { fatalError("MLX network not loaded") }
+
+ // Extract mel features: [T, 80]
+ let mel = melExtractor.extract(samples)
+
+ // Add batch and channel dims: [1, T, 80, 1]
+ let input = mel.reshaped(1, mel.dim(0), mel.dim(1), 1)
+
+ // Forward pass: [1, 256]
+ let emb = network(input)
+ eval(emb)
+
+ return emb[0].asArray(Float.self)
+ }
+
+ /// Compute cosine similarity between two embeddings.
+ public static func cosineSimilarity(_ a: [Float], _ b: [Float]) -> Float {
+ guard a.count == b.count, !a.isEmpty else { return 0 }
+ var dot: Float = 0
+ var normA: Float = 0
+ var normB: Float = 0
+ for i in 0.. 0 ? dot / denom : 0
+ }
+
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeakerModel.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeakerModel.swift
new file mode 100644
index 0000000..39653f0
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeakerModel.swift
@@ -0,0 +1,167 @@
+import MLX
+import MLXNN
+
+/// ResNet BasicBlock with BN fused into Conv2d.
+///
+/// Each block has two 3×3 Conv2d layers with bias (fused BatchNorm).
+/// Shortcut Conv2d(1×1) is added when stride≠1 or channels change.
+class BasicBlock: Module {
+ let conv1: Conv2d
+ let conv2: Conv2d
+ let shortcut: Conv2d?
+ let stride: Int
+
+ init(inChannels: Int, outChannels: Int, stride: Int = 1) {
+ self.stride = stride
+
+ self.conv1 = Conv2d(
+ inputChannels: inChannels, outputChannels: outChannels,
+ kernelSize: 3, stride: IntOrPair(arrayLiteral: stride, stride),
+ padding: 1, bias: true
+ )
+ self.conv2 = Conv2d(
+ inputChannels: outChannels, outputChannels: outChannels,
+ kernelSize: 3, stride: 1, padding: 1, bias: true
+ )
+
+ if stride != 1 || inChannels != outChannels {
+ self.shortcut = Conv2d(
+ inputChannels: inChannels, outputChannels: outChannels,
+ kernelSize: 1, stride: IntOrPair(arrayLiteral: stride, stride),
+ padding: 0, bias: true
+ )
+ } else {
+ self.shortcut = nil
+ }
+ }
+
+ func callAsFunction(_ x: MLXArray) -> MLXArray {
+ var out = relu(conv1(x))
+ out = conv2(out)
+
+ let residual: MLXArray
+ if let shortcut {
+ residual = shortcut(x)
+ } else {
+ residual = x
+ }
+
+ return relu(out + residual)
+ }
+}
+
+/// WeSpeaker ResNet34 speaker embedding network (BN-fused).
+///
+/// Architecture:
+/// ```
+/// Input: [B, T, 80, 1] mel spectrogram
+/// → Conv2d(1→32, k=3, p=1) + ReLU
+/// → Layer1: 3× BasicBlock(32→32)
+/// → Layer2: 4× BasicBlock(32→64, s=2)
+/// → Layer3: 6× BasicBlock(64→128, s=2)
+/// → Layer4: 3× BasicBlock(128→256, s=2)
+/// → Statistics Pooling: mean + std → [B, 5120]
+/// → Linear(5120→256) → L2 normalize
+/// Output: [B, 256] speaker embedding
+/// ```
+class WeSpeakerNetwork: Module {
+ let conv1: Conv2d
+ let layer1: [BasicBlock]
+ let layer2: [BasicBlock]
+ let layer3: [BasicBlock]
+ let layer4: [BasicBlock]
+ let embedding: Linear
+
+ override init() {
+ self.conv1 = Conv2d(
+ inputChannels: 1, outputChannels: 32,
+ kernelSize: 3, stride: 1, padding: 1, bias: true
+ )
+
+ // Layer1: 3 blocks, 32→32
+ var blocks1 = [BasicBlock]()
+ for _ in 0..<3 {
+ blocks1.append(BasicBlock(inChannels: 32, outChannels: 32))
+ }
+ self.layer1 = blocks1
+
+ // Layer2: 4 blocks, 32→64, first stride=2
+ var blocks2 = [BasicBlock]()
+ for i in 0..<4 {
+ blocks2.append(BasicBlock(
+ inChannels: i == 0 ? 32 : 64,
+ outChannels: 64,
+ stride: i == 0 ? 2 : 1
+ ))
+ }
+ self.layer2 = blocks2
+
+ // Layer3: 6 blocks, 64→128, first stride=2
+ var blocks3 = [BasicBlock]()
+ for i in 0..<6 {
+ blocks3.append(BasicBlock(
+ inChannels: i == 0 ? 64 : 128,
+ outChannels: 128,
+ stride: i == 0 ? 2 : 1
+ ))
+ }
+ self.layer3 = blocks3
+
+ // Layer4: 3 blocks, 128→256, first stride=2
+ var blocks4 = [BasicBlock]()
+ for i in 0..<3 {
+ blocks4.append(BasicBlock(
+ inChannels: i == 0 ? 128 : 256,
+ outChannels: 256,
+ stride: i == 0 ? 2 : 1
+ ))
+ }
+ self.layer4 = blocks4
+
+ // Pooling output: T/8 * 10 * 256 → mean+std → 2 * 10 * 256 = 5120
+ self.embedding = Linear(5120, 256)
+ }
+
+ /// Forward pass.
+ /// - Parameter mel: `[B, T, 80, 1]` mel spectrogram (channels-last)
+ /// - Returns: `[B, 256]` L2-normalized speaker embedding
+ func callAsFunction(_ mel: MLXArray) -> MLXArray {
+ // mel: [B, T, 80, 1]
+ // Python WeSpeaker permutes input: (B,T,F) -> (B,F,T) -> (B,1,F,T)
+ // In MLX NHWC: (B,1,F,T) maps to [B, F, T, 1]
+ var x = mel.transposed(0, 2, 1, 3) // [B, 80, T, 1] = [B, F, T, C]
+
+ x = relu(conv1(x))
+
+ // ResNet layers
+ for block in layer1 { x = block(x) }
+ for block in layer2 { x = block(x) }
+ for block in layer3 { x = block(x) }
+ for block in layer4 { x = block(x) }
+ // x: [B, F'=10, T'=T/8, 256] (NHWC)
+ // Corresponds to Python's [B, 256, F'=10, T'=T/8] (NCHW)
+
+ // Flatten freq and channels: [B, 10, T/8, 256] → [B, T/8, 10*256]
+ // Match Python: [B, 256, 10, T'] → [B, 256*10, T'] via reshape (C*F order)
+ // MLX: transpose to [B, T/8, 256, 10] then reshape
+ let B = x.dim(0)
+ let Tp = x.dim(2) // T/8 (time is dim 2 now)
+ x = x.transposed(0, 2, 3, 1) // [B, T/8, 256, 10]
+ x = x.reshaped(B, Tp, -1) // [B, T/8, 2560] in C*F order
+
+ // Statistics pooling: mean + std over time (dim=1) → [B, 5120]
+ let mean = x.mean(axis: 1) // [B, 2560]
+ let variance = x.variance(axis: 1) // [B, 2560]
+ let std = sqrt(variance + 1e-10)
+ let pooled = concatenated([mean, std], axis: -1) // [B, 5120]
+
+ // Embedding projection
+ var emb = embedding(pooled) // [B, 256]
+
+ // L2 normalize
+ let norm = sqrt((emb * emb).sum(axis: -1, keepDims: true) + 1e-10)
+ emb = emb / norm
+
+ return emb
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeakerWeightLoading.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeakerWeightLoading.swift
new file mode 100644
index 0000000..e313ae9
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeakerWeightLoading.swift
@@ -0,0 +1,33 @@
+import Foundation
+import MLXCommon
+import MLX
+import MLXNN
+import AudioCommon
+
+/// Weight loading for the WeSpeaker ResNet34-LM speaker embedding model.
+///
+/// Loads from safetensors files produced by `scripts/convert_wespeaker.py`.
+/// The conversion script fuses BatchNorm into Conv2d and transposes weights,
+/// so loading is a straightforward parameter tree update.
+enum WeSpeakerWeightLoader {
+
+ /// Load weights from a directory containing model.safetensors.
+ static func loadWeights(
+ model: WeSpeakerNetwork,
+ from directory: URL
+ ) throws {
+ let weightsURL = directory.appendingPathComponent("model.safetensors")
+
+ guard FileManager.default.fileExists(atPath: weightsURL.path) else {
+ throw WeightLoadingError.noWeightsFound(directory)
+ }
+
+ let weights = try MLX.loadArrays(url: weightsURL)
+
+ let parameters = ModuleParameters.unflattened(weights)
+
+ try model.update(parameters: parameters, verify: .noUnusedKeys)
+
+ MLX.eval(model.parameters())
+ }
+}
diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeightLoading.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeightLoading.swift
new file mode 100644
index 0000000..6fb0eba
--- /dev/null
+++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeightLoading.swift
@@ -0,0 +1,37 @@
+import Foundation
+import MLXCommon
+import MLX
+import MLXNN
+import AudioCommon
+
+/// Weight loading for the PyanNet segmentation model.
+///
+/// Loads from safetensors files produced by `scripts/convert_pyannote.py`.
+/// The conversion script pre-computes sinc filters and transposes Conv1d weights,
+/// so loading is a straightforward parameter tree update.
+enum SegmentationWeightLoader {
+
+ /// Load weights from a directory containing model.safetensors.
+ static func loadWeights(
+ model: SegmentationModel,
+ from directory: URL
+ ) throws {
+ let weightsURL = directory.appendingPathComponent("model.safetensors")
+
+ guard FileManager.default.fileExists(atPath: weightsURL.path) else {
+ throw WeightLoadingError.noWeightsFound(directory)
+ }
+
+ let weights = try MLX.loadArrays(url: weightsURL)
+
+ // The conversion script produces keys matching our module structure directly.
+ // Use ModuleParameters.unflattened to build the nested parameter tree.
+ let parameters = ModuleParameters.unflattened(weights)
+
+ // Apply to model
+ try model.update(parameters: parameters, verify: .noUnusedKeys)
+
+ // Evaluate all parameters to ensure they're materialized
+ MLX.eval(model.parameters())
+ }
+}
diff --git a/OSGKeyboard/Utilities/ASRLocaleLabels.swift b/OSGKeyboard/Utilities/ASRLocaleLabels.swift
new file mode 100644
index 0000000..4116445
--- /dev/null
+++ b/OSGKeyboard/Utilities/ASRLocaleLabels.swift
@@ -0,0 +1,32 @@
+// ASRLocaleLabels.swift
+// OSGKeyboard · Main App
+//
+// Human-readable labels for ASR locale picker rows. Honors the in-app
+// UI language override instead of caching strings at load time.
+
+import Foundation
+import OSGKeyboardShared
+
+enum ASRLocaleLabels {
+ private static let bundledKeys: [String: String] = [
+ "auto": "locale.auto",
+ "zh-Hans": "locale.zh-Hans",
+ "zh-Hant": "locale.zh-Hant",
+ "en-US": "locale.en-US",
+ "ja-JP": "locale.ja-JP",
+ "ko-KR": "locale.ko-KR",
+ ]
+
+ static func displayName(for localeId: String, language: AppUILanguage) -> String {
+ if let key = bundledKeys[localeId] {
+ return AppUILanguage.localizedString(
+ key,
+ tableName: nil,
+ bundle: .main,
+ language: language
+ )
+ }
+ let uiLocale = Locale(identifier: language.resolvedLanguageCode())
+ return uiLocale.localizedString(forIdentifier: localeId) ?? localeId
+ }
+}
diff --git a/OSGKeyboard/Utilities/AppL10n.swift b/OSGKeyboard/Utilities/AppL10n.swift
new file mode 100644
index 0000000..4d5690e
--- /dev/null
+++ b/OSGKeyboard/Utilities/AppL10n.swift
@@ -0,0 +1,40 @@
+// AppL10n.swift
+// OSGKeyboard · Main App
+//
+// Loads Localizable.strings from the host app bundle while honoring
+// the in-app UI language override (not only the system language).
+
+import Foundation
+import SwiftUI
+import OSGKeyboardShared
+
+enum AppL10n {
+ static func string(
+ _ key: String,
+ language: AppUILanguage? = nil
+ ) -> String {
+ let lang = language ?? ProviderConfig.shared.uiLanguage
+ return AppUILanguage.localizedString(
+ key,
+ tableName: nil,
+ bundle: .main,
+ language: lang
+ )
+ }
+
+ static func format(
+ _ key: String,
+ language: AppUILanguage? = nil,
+ _ args: CVarArg...
+ ) -> String {
+ String(
+ format: string(key, language: language),
+ locale: Locale.current,
+ arguments: args
+ )
+ }
+
+ static func text(_ key: String) -> Text {
+ Text(LocalizedStringKey(key))
+ }
+}
diff --git a/OSGKeyboard/Views/APISettingsCard.swift b/OSGKeyboard/Views/APISettingsCard.swift
index 9716894..cd00e16 100644
--- a/OSGKeyboard/Views/APISettingsCard.swift
+++ b/OSGKeyboard/Views/APISettingsCard.swift
@@ -24,7 +24,7 @@ struct APISettingsCard: View {
var body: some View {
VStack(spacing: 0) {
field(
- title: NSLocalizedString("api.baseUrl", comment: ""),
+ title: AppL10n.string("api.baseUrl"),
placeholder: "https://api.openai.com/v1",
text: $config.baseURL,
autocap: false
@@ -33,7 +33,7 @@ struct APISettingsCard: View {
keyField
Divider().background(palette.divider)
field(
- title: NSLocalizedString("api.model", comment: ""),
+ title: AppL10n.string("api.model"),
placeholder: "gpt-4o-mini",
text: $config.model,
autocap: false
@@ -86,8 +86,8 @@ struct APISettingsCard: View {
}
.buttonStyle(.plain)
.accessibilityLabel(Text(showKey
- ? NSLocalizedString("api.key.hide", comment: "")
- : NSLocalizedString("api.key.show", comment: "")))
+ ? AppL10n.string("api.key.hide")
+ : AppL10n.string("api.key.show")))
}
Group {
if showKey {
@@ -170,10 +170,10 @@ struct APISettingsCard: View {
private var testButtonLabel: String {
switch testStatus {
- case .idle: return NSLocalizedString("api.test.idle", comment: "")
- case .running: return NSLocalizedString("api.test.running", comment: "")
- case .success: return NSLocalizedString("api.test.success", comment: "")
- case .failure: return NSLocalizedString("api.test.failure", comment: "")
+ case .idle: return AppL10n.string("api.test.idle")
+ case .running: return AppL10n.string("api.test.running")
+ case .success: return AppL10n.string("api.test.success")
+ case .failure: return AppL10n.string("api.test.failure")
}
}
@@ -205,21 +205,15 @@ struct APISettingsCard: View {
_ = try await client.polish("ping", systemPrompt: "Reply with the single word PONG.")
testStatus = .success
} catch LLMError.noAPIKey {
- testStatus = .failure(NSLocalizedString("api.test.missing", comment: ""))
+ testStatus = .failure(AppL10n.string("api.test.missing"))
} catch let error as LLMError {
switch error {
case .http(let status):
- testStatus = .failure(String.localizedStringWithFormat(
- NSLocalizedString("api.test.http", comment: ""),
- status
- ))
+ testStatus = .failure(AppL10n.format("api.test.http", status))
case .rateLimited:
- testStatus = .failure(NSLocalizedString("api.test.rateLimited", comment: ""))
+ testStatus = .failure(AppL10n.string("api.test.rateLimited"))
case .transport(let msg):
- testStatus = .failure(String.localizedStringWithFormat(
- NSLocalizedString("api.test.transportWith", comment: ""),
- msg
- ))
+ testStatus = .failure(AppL10n.format("api.test.transportWith", msg))
default:
testStatus = .failure(error.errorDescription ?? "\(error)")
}
diff --git a/OSGKeyboard/Views/Components/LegalWebView.swift b/OSGKeyboard/Views/Components/LegalWebView.swift
new file mode 100644
index 0000000..ee6375f
--- /dev/null
+++ b/OSGKeyboard/Views/Components/LegalWebView.swift
@@ -0,0 +1,50 @@
+// LegalWebView.swift
+// OSGKeyboard · Main App
+//
+// In-app HTML viewer for bundled legal documents (privacy policy).
+
+import SwiftUI
+import WebKit
+
+struct LegalWebView: UIViewRepresentable {
+ let resourceName: String
+ var scrollToAnchor: String?
+
+ func makeCoordinator() -> Coordinator {
+ Coordinator(scrollToAnchor: scrollToAnchor)
+ }
+
+ func makeUIView(context: Context) -> WKWebView {
+ let webView = WKWebView(frame: .zero)
+ webView.isOpaque = false
+ webView.backgroundColor = .clear
+ webView.scrollView.backgroundColor = .clear
+ webView.navigationDelegate = context.coordinator
+ context.coordinator.webView = webView
+
+ guard let url = Bundle.main.url(forResource: resourceName, withExtension: "html") else {
+ return webView
+ }
+ webView.loadFileURL(url, allowingReadAccessTo: url.deletingLastPathComponent())
+ return webView
+ }
+
+ func updateUIView(_ uiView: WKWebView, context: Context) {
+ context.coordinator.scrollToAnchor = scrollToAnchor
+ }
+
+ final class Coordinator: NSObject, WKNavigationDelegate {
+ var scrollToAnchor: String?
+ weak var webView: WKWebView?
+
+ init(scrollToAnchor: String?) {
+ self.scrollToAnchor = scrollToAnchor
+ }
+
+ func webView(_ webView: WKWebView, didFinish navigation: WKNavigation!) {
+ guard let anchor = scrollToAnchor, !anchor.isEmpty else { return }
+ let escaped = anchor.replacingOccurrences(of: "'", with: "\\'")
+ webView.evaluateJavaScript("location.hash = '#\(escaped)';") { _, _ in }
+ }
+ }
+}
diff --git a/OSGKeyboard/Views/Components/RemoteWebView.swift b/OSGKeyboard/Views/Components/RemoteWebView.swift
new file mode 100644
index 0000000..c525ab4
--- /dev/null
+++ b/OSGKeyboard/Views/Components/RemoteWebView.swift
@@ -0,0 +1,52 @@
+// RemoteWebView.swift
+// OSGKeyboard · Main App
+//
+// In-app WKWebView for remote HTTPS pages (e.g. GitHub Issues).
+
+import SwiftUI
+import WebKit
+
+struct RemoteWebView: UIViewRepresentable {
+ let url: URL
+ @Binding var isLoading: Bool
+
+ func makeCoordinator() -> Coordinator {
+ Coordinator(isLoading: $isLoading)
+ }
+
+ func makeUIView(context: Context) -> WKWebView {
+ let webView = WKWebView(frame: .zero)
+ webView.isOpaque = false
+ webView.backgroundColor = .clear
+ webView.scrollView.backgroundColor = .clear
+ webView.navigationDelegate = context.coordinator
+ webView.load(URLRequest(url: url))
+ return webView
+ }
+
+ func updateUIView(_ uiView: WKWebView, context: Context) {}
+
+ final class Coordinator: NSObject, WKNavigationDelegate {
+ @Binding var isLoading: Bool
+
+ init(isLoading: Binding) {
+ _isLoading = isLoading
+ }
+
+ func webView(_ webView: WKWebView, didStartProvisionalNavigation navigation: WKNavigation!) {
+ isLoading = true
+ }
+
+ func webView(_ webView: WKWebView, didFinish navigation: WKNavigation!) {
+ isLoading = false
+ }
+
+ func webView(_ webView: WKWebView, didFail navigation: WKNavigation!, withError error: Error) {
+ isLoading = false
+ }
+
+ func webView(_ webView: WKWebView, didFailProvisionalNavigation navigation: WKNavigation!, withError error: Error) {
+ isLoading = false
+ }
+ }
+}
diff --git a/OSGKeyboard/Views/Components/SafariSheet.swift b/OSGKeyboard/Views/Components/SafariSheet.swift
index f287d3b..1aba856 100644
--- a/OSGKeyboard/Views/Components/SafariSheet.swift
+++ b/OSGKeyboard/Views/Components/SafariSheet.swift
@@ -3,7 +3,6 @@
import SafariServices
import SwiftUI
-import OSGKeyboardShared
extension URL: @retroactive Identifiable {
public var id: String { absoluteString }
@@ -13,9 +12,7 @@ struct SafariSheet: UIViewControllerRepresentable {
let url: URL
func makeUIViewController(context: Context) -> SFSafariViewController {
- let controller = SFSafariViewController(url: url)
- controller.preferredControlTintColor = UIColor(Palette.accent)
- return controller
+ SFSafariViewController(url: url)
}
func updateUIViewController(_ uiViewController: SFSafariViewController, context: Context) {}
diff --git a/OSGKeyboard/Views/DictationCaptureView.swift b/OSGKeyboard/Views/DictationCaptureView.swift
index 48bd7f9..8641928 100644
--- a/OSGKeyboard/Views/DictationCaptureView.swift
+++ b/OSGKeyboard/Views/DictationCaptureView.swift
@@ -29,7 +29,7 @@ struct DictationCaptureView: View {
@ObservedObject var coordinator: DictationSessionCoordinator
@StateObject private var dictation = LiveDictationController()
- @State private var statusText: String = "准备录音..."
+ @State private var statusText: String = ""
@State private var isSaving: Bool = false
var body: some View {
@@ -60,7 +60,7 @@ struct DictationCaptureView: View {
Button {
cancelAndClose()
} label: {
- Text("取消")
+ Text("common.cancel")
.secondaryButton()
}
.buttonStyle(.plain)
@@ -69,7 +69,7 @@ struct DictationCaptureView: View {
Button {
stopAndFinalize()
} label: {
- Text("完成")
+ Text("common.done")
.primaryButton()
}
.buttonStyle(.plain)
@@ -80,6 +80,7 @@ struct DictationCaptureView: View {
}
}
.onAppear {
+ statusText = AppL10n.string("dictation.status.ready")
DictationBridge.setStatus(.requested)
startRecording()
}
@@ -90,12 +91,14 @@ struct DictationCaptureView: View {
switch new {
case .recording:
DictationBridge.setStatus(.recording)
- statusText = config.isLocalEngine ? "正在实时识别..." : "正在听..."
+ statusText = config.isLocalEngine
+ ? AppL10n.string("dictation.status.listeningLive")
+ : AppL10n.string("dictation.status.listening")
case .processing:
DictationBridge.setStatus(.transcribing)
- statusText = "处理中..."
+ statusText = AppL10n.string("dictation.status.processing")
case .requestingPermission:
- statusText = "请求权限中..."
+ statusText = AppL10n.string("dictation.status.requestingPermission")
case .denied(let message):
DictationBridge.setStatus(.error, message: message)
statusText = message
@@ -114,10 +117,15 @@ struct DictationCaptureView: View {
guard !new.isEmpty else { return }
saveAndClose(new)
}
+ .onChange(of: config.uiLanguage) { _, _ in
+ refreshStatusForCurrentPhase()
+ }
}
private var titleText: String {
- isSaving ? "保存中..." : "语音输入"
+ isSaving
+ ? AppL10n.string("dictation.status.saving")
+ : AppL10n.string("dictation.title")
}
private func startRecording() {
@@ -126,7 +134,7 @@ struct DictationCaptureView: View {
private func stopAndFinalize() {
dictation.stop()
- statusText = "等待识别结果..."
+ statusText = AppL10n.string("dictation.status.waitingResult")
}
private func cancelAndClose() {
@@ -139,8 +147,29 @@ struct DictationCaptureView: View {
private func saveAndClose(_ transcript: String) {
guard !isSaving else { return }
isSaving = true
- DictationBridge.storePendingTranscript(transcript)
- coordinator.dismiss()
- dismiss()
+ statusText = AppL10n.string("dictation.status.processing")
+ Task {
+ let delivered = transcript
+ DictationBridge.storePendingTranscript(delivered, polishWarning: nil)
+ coordinator.dismiss()
+ dismiss()
+ }
+ }
+
+ private func refreshStatusForCurrentPhase() {
+ switch dictation.phase {
+ case .idle:
+ statusText = AppL10n.string("dictation.status.ready")
+ case .recording:
+ statusText = config.isLocalEngine
+ ? AppL10n.string("dictation.status.listeningLive")
+ : AppL10n.string("dictation.status.listening")
+ case .processing:
+ statusText = AppL10n.string("dictation.status.processing")
+ case .requestingPermission:
+ statusText = AppL10n.string("dictation.status.requestingPermission")
+ case .denied(let message), .error(let message):
+ statusText = message
+ }
}
}
diff --git a/OSGKeyboard/Views/DownloadConfirmSheet.swift b/OSGKeyboard/Views/DownloadConfirmSheet.swift
new file mode 100644
index 0000000..f8ae1d5
--- /dev/null
+++ b/OSGKeyboard/Views/DownloadConfirmSheet.swift
@@ -0,0 +1,85 @@
+// DownloadConfirmSheet.swift
+// OSGKeyboard · Main App
+//
+// One-step confirmation before an on-device model download starts.
+// Progress and cancellation live on `OnDeviceModelsView` — this
+// sheet dismisses as soon as the user confirms.
+
+import SwiftUI
+import OSGKeyboardShared
+
+struct DownloadConfirmSheet: View {
+ @Environment(\.themePalette) private var palette: ThemePalette
+ @Environment(\.dismiss) private var dismiss
+
+ let model: OnDeviceModel
+ let onConfirm: () -> Void
+
+ private let sheetHeight: CGFloat = 340
+
+ /// Horizontal inset for copy — 50% wider than the default `Spacing.md`.
+ private var textHorizontalInset: CGFloat { Spacing.md * 1.5 }
+
+ /// Top inset — 50% more than the previous `Spacing.xl` (24 → 36).
+ private var topContentInset: CGFloat { Spacing.xl * 1.5 }
+
+ var body: some View {
+ VStack(alignment: .leading, spacing: 0) {
+ topContent
+ Spacer(minLength: Spacing.md)
+ actionButtons
+ }
+ .padding(.horizontal, textHorizontalInset)
+ .padding(.top, topContentInset)
+ .padding(.bottom, Spacing.md)
+ .frame(maxWidth: .infinity, minHeight: sheetHeight, maxHeight: sheetHeight, alignment: .topLeading)
+ .presentationDetents([.height(sheetHeight)])
+ .presentationDragIndicator(.visible)
+ .presentationBackground(.clear)
+ }
+
+ // MARK: - Top
+
+ private var topContent: some View {
+ VStack(alignment: .leading, spacing: Spacing.lg) {
+ Text(AppL10n.format("settings.models.confirm.downloadTitle %@", model.displayName))
+ .font(TypeStyle.title3)
+ .foregroundStyle(palette.textPrimary)
+ .fixedSize(horizontal: false, vertical: true)
+
+ Text(AppL10n.format("settings.models.confirm.body %lld", model.approximateSizeMB))
+ .font(TypeStyle.body)
+ .foregroundStyle(palette.textPrimary)
+ .fixedSize(horizontal: false, vertical: true)
+ }
+ .frame(maxWidth: .infinity, alignment: .topLeading)
+ }
+
+ // MARK: - Actions
+
+ private var actionButtons: some View {
+ VStack(spacing: Spacing.sm) {
+ Button {
+ onConfirm()
+ dismiss()
+ } label: {
+ Text(AppL10n.format("settings.models.confirm.download %lld", model.approximateSizeMB))
+ }
+ .primaryButton()
+ .buttonStyle(.plain)
+
+ Button {
+ dismiss()
+ } label: {
+ Text("settings.models.confirm.cancel")
+ }
+ .secondaryButton()
+ .buttonStyle(.plain)
+ }
+ }
+}
+
+#Preview {
+ DownloadConfirmSheet(model: .qwen3ASR) { }
+ .environment(\.themePalette, Palette.light)
+}
diff --git a/OSGKeyboard/Views/EnginePickerSection.swift b/OSGKeyboard/Views/EnginePickerSection.swift
index 0af8be5..5f75525 100644
--- a/OSGKeyboard/Views/EnginePickerSection.swift
+++ b/OSGKeyboard/Views/EnginePickerSection.swift
@@ -1,19 +1,8 @@
// EnginePickerSection.swift
// OSGKeyboard · Main App
//
-// Engine picker — Local (on-device ASR, no LLM) vs Cloud (ASR + LLM
-// polish). Lives in its own file so the onboarding flow (first-run,
-// no API key yet) and the in-app Settings sheet can render the same
-// component: both need to expose the same two options, and the user
-// must reach the same "Cloud needs an API key" conclusion from either
-// entry point.
-//
-// Selecting "local" also forces `modeId = "transcribe"`: the Local
-// engine skips the LLM round-trip, so leaving modeId on `polish`
-// would surface a confusing "I set everything up and nothing
-// happens" state. The settings UI is the source of truth for
-// `engineMode`; the onboarding page mutates the same `ProviderConfig`
-// singleton.
+// Engine picker — Local (on-device ASR) vs
+// Cloud (ASR + optional LLM polish via the user's API).
import SwiftUI
import OSGKeyboardShared
@@ -31,16 +20,16 @@ struct EnginePickerSection: View {
VStack(spacing: 0) {
engineOptionRow(
id: "local",
- assetName: "apple",
- title: NSLocalizedString("settings.engine.local.title", comment: ""),
+ systemIcon: "iphone.badge.checkmark",
+ title: AppL10n.string("settings.engine.local.title"),
subtitle: localSubtitle
)
Divider().background(palette.divider)
engineOptionRow(
id: "cloud",
systemIcon: "wand.and.stars",
- title: NSLocalizedString("settings.engine.cloud.title", comment: ""),
- subtitle: NSLocalizedString("settings.engine.cloud.subtitle", comment: "")
+ title: AppL10n.string("settings.engine.cloud.title"),
+ subtitle: AppL10n.string("settings.engine.cloud.subtitle")
)
}
.background(palette.surface, in: RoundedRectangle(cornerRadius: Radius.large, style: .continuous))
@@ -62,9 +51,7 @@ struct EnginePickerSection: View {
}
private var localSubtitle: String {
- // iOS 26's `SpeechAnalyzer` is always fully on-device, so the
- // local engine's only contract is "no network, no LLM".
- NSLocalizedString("settings.engine.local.ios26", comment: "")
+ AppL10n.string("settings.engine.local.legacy")
}
private func engineOptionRow(
@@ -115,7 +102,10 @@ struct EnginePickerSection: View {
private func selectEngine(_ id: String) {
withAnimation(.easeInOut(duration: 0.2)) {
config.engineMode = id
- if id == "local" { config.modeId = "transcribe" }
+ // Cloud always runs ASR + LLM polish; no off/transcribe toggle.
+ if id == "cloud" {
+ config.modeId = "polish"
+ }
}
}
diff --git a/OSGKeyboard/Views/HelpFeedbackView.swift b/OSGKeyboard/Views/HelpFeedbackView.swift
new file mode 100644
index 0000000..d2106ba
--- /dev/null
+++ b/OSGKeyboard/Views/HelpFeedbackView.swift
@@ -0,0 +1,35 @@
+// HelpFeedbackView.swift
+// OSGKeyboard · Main App
+//
+// In-app GitHub Issues page. Reached from Settings → About.
+
+import SwiftUI
+import OSGKeyboardShared
+
+struct HelpFeedbackView: View {
+ @Environment(\.themePalette) private var palette: ThemePalette
+ @State private var isLoading = true
+
+ var body: some View {
+ Group {
+ if let url = LegalLinks.supportURL {
+ ZStack {
+ RemoteWebView(url: url, isLoading: $isLoading)
+ if isLoading {
+ ProgressView()
+ .tint(palette.accent)
+ }
+ }
+ } else {
+ ContentUnavailableView(
+ "settings.support.unavailable.title",
+ systemImage: "wifi.slash",
+ description: Text("settings.support.unavailable.message")
+ )
+ }
+ }
+ .background(palette.background.ignoresSafeArea())
+ .navigationTitle("settings.link.support")
+ .navigationBarTitleDisplayMode(.inline)
+ }
+}
diff --git a/OSGKeyboard/Views/HomeView.swift b/OSGKeyboard/Views/HomeView.swift
index 14380ab..855b4ce 100644
--- a/OSGKeyboard/Views/HomeView.swift
+++ b/OSGKeyboard/Views/HomeView.swift
@@ -12,6 +12,8 @@ struct HomeView: View {
@Environment(\.scenePhase) private var scenePhase
@ObservedObject private var config = ProviderConfig.shared
+ @ObservedObject private var modelWarmup = OnDeviceModelWarmup.shared
+ @ObservedObject private var modelManager = ModelManager.shared
@EnvironmentObject private var flowManager: FlowSessionManager
@FocusState private var previewFocused: Bool
@State private var previewText = ""
@@ -37,6 +39,33 @@ struct HomeView: View {
&& !needsPermissionSetup
&& flowManager.sessionWarning == nil
&& !needsCloudSetup
+ && !localModelNeedsAttention
+ }
+
+ /// Local engine still needs model download and/or in-memory warm-up.
+ private var localModelNeedsAttention: Bool {
+ guard config.isLocalEngine else { return false }
+ if config.isLocalEngine, config.localASRBackend == .qwen3ASR,
+ !OnDeviceMLRuntime.supportsOnDeviceQwen3 { return true }
+ if isAnyModelDownloading { return true }
+ if !OnDeviceModelStatus.isLocalStackReady(asrBackend: config.localASRBackend) {
+ return true
+ }
+ switch modelWarmup.phase {
+ case .ready, .notNeeded:
+ return false
+ case .warming, .failed, .idle:
+ // `.idle` with a downloaded stack means warm-up has not finished yet.
+ return true
+ }
+ }
+
+ private var isAnyModelDownloading: Bool {
+ if !modelManager.activeDownloads.isEmpty { return true }
+ return OnDeviceModel.allCases.contains { model in
+ if case .downloading = modelManager.states[model]?.download { return true }
+ return OnDeviceModelStatus.downloadProgress(model) != nil
+ }
}
var body: some View {
@@ -53,9 +82,11 @@ struct HomeView: View {
.padding(.top, Spacing.xxxl)
.padding(.bottom, Spacing.xl)
- flowSessionExtras
- .padding(.horizontal, Spacing.lg)
- .padding(.bottom, Spacing.lg)
+ if showsFlowSessionExtras {
+ flowSessionExtras
+ .padding(.horizontal, Spacing.lg)
+ .padding(.bottom, Spacing.lg)
+ }
previewField
.padding(.horizontal, Spacing.lg)
@@ -70,14 +101,39 @@ struct HomeView: View {
}
.background(palette.background)
}
- .onAppear { refreshPermissionStatuses() }
+ .onAppear {
+ refreshPermissionStatuses()
+ scheduleModelWarmup(force: modelWarmup.phase.isFailed)
+ }
.onChange(of: scenePhase) { _, phase in
guard phase == .active else { return }
refreshPermissionStatuses()
+ if config.isLocalEngine {
+ OnDeviceModelWarmup.shared.ensureReadyAfterBackground()
+ }
}
.onReceive(NotificationCenter.default.publisher(for: UIApplication.didBecomeActiveNotification)) { _ in
refreshPermissionStatuses()
}
+ .onAppear { scheduleModelWarmup() }
+ .onChange(of: config.engineMode) { _, _ in scheduleModelWarmup(force: true) }
+ .onChange(of: config.localASRBackend) { _, _ in
+ modelWarmup.invalidate()
+ scheduleModelWarmup(force: true)
+ }
+ .onChange(of: modelWarmup.phase) { _, phase in
+ if phase == .idle, config.isLocalEngine {
+ scheduleModelWarmup()
+ }
+ }
+ }
+
+ private func scheduleModelWarmup(force: Bool = false) {
+ guard config.isLocalEngine else {
+ modelWarmup.invalidate()
+ return
+ }
+ modelWarmup.warmUpIfNeeded(force: force)
}
private func refreshPermissionStatuses() {
@@ -103,23 +159,36 @@ struct HomeView: View {
private func sessionHeaderGradient(height: CGFloat) -> some View {
LinearGradient(
- colors: sessionIsLive
- ? [
- palette.accent.opacity(0.28),
- palette.accent.opacity(0.10),
- palette.background.opacity(0)
- ]
- : [
- palette.textTertiary.opacity(0.14),
- palette.textTertiary.opacity(0.05),
- palette.background.opacity(0)
- ],
+ colors: headerGradientColors,
startPoint: .top,
endPoint: .bottom
)
.frame(maxWidth: .infinity)
.frame(height: height)
.animation(Motion.soft, value: sessionIsLive)
+ .animation(Motion.soft, value: localModelNeedsAttention)
+ }
+
+ private var headerGradientColors: [Color] {
+ if localModelNeedsAttention {
+ return [
+ palette.warning.opacity(0.32),
+ palette.warning.opacity(0.12),
+ palette.background.opacity(0)
+ ]
+ }
+ if sessionIsLive {
+ return [
+ palette.accent.opacity(0.28),
+ palette.accent.opacity(0.10),
+ palette.background.opacity(0)
+ ]
+ }
+ return [
+ palette.textTertiary.opacity(0.14),
+ palette.textTertiary.opacity(0.05),
+ palette.background.opacity(0)
+ ]
}
// MARK: - Header
@@ -173,7 +242,9 @@ struct HomeView: View {
.fill(flowStatusColor)
.frame(width: 6, height: 6)
- if flowManager.isActive, let expires = flowManager.sessionExpiresAt {
+ if flowManager.isActive,
+ let expires = flowManager.sessionExpiresAt,
+ !localModelNeedsAttention {
Text("home.flow.label")
.font(TypeStyle.status)
.foregroundStyle(palette.textPrimary)
@@ -185,11 +256,12 @@ struct HomeView: View {
.foregroundStyle(palette.textSecondary)
.monospacedDigit()
} else {
- Text(flowStatusTitle)
+ Text(flowCapsuleStatusMessage)
.font(TypeStyle.status)
.foregroundStyle(palette.textPrimary)
- .lineLimit(1)
+ .lineLimit(2)
.minimumScaleFactor(0.85)
+ .multilineTextAlignment(.leading)
}
}
}
@@ -202,6 +274,14 @@ struct HomeView: View {
// MARK: - Flow extras (warnings / hints)
+ private var showsFlowSessionExtras: Bool {
+ needsPermissionSetup
+ || flowManager.sessionWarning != nil
+ || needsCloudSetup
+ || shouldShowKeyboardHint
+ || (!flowManager.isActive && !localModelNeedsAttention)
+ }
+
@ViewBuilder
private var flowSessionExtras: some View {
if needsPermissionSetup {
@@ -251,7 +331,7 @@ struct HomeView: View {
}
.buttonStyle(.plain)
}
- } else if !flowManager.isActive {
+ } else {
Text("home.flow.hint")
.font(TypeStyle.caption2)
.foregroundStyle(palette.textTertiary)
@@ -275,6 +355,7 @@ struct HomeView: View {
}
private var flowStatusColor: Color {
+ if localModelNeedsAttention { return palette.warning }
if flowManager.isActive { return palette.accent }
if flowManager.isStarting { return palette.accent }
if needsPermissionSetup { return palette.warning }
@@ -282,10 +363,37 @@ struct HomeView: View {
return palette.textTertiary
}
- private var flowStatusTitle: LocalizedStringKey {
- if flowManager.isActive { return "home.flow.active" }
- if flowManager.isStarting { return "home.flow.starting" }
- return "home.flow.inactive"
+ /// Single source of truth for the logo status capsule (local model + flow state).
+ private var flowCapsuleStatusMessage: String {
+ if config.isLocalEngine, config.localASRBackend == .qwen3ASR,
+ !OnDeviceMLRuntime.supportsOnDeviceQwen3 {
+ return AppL10n.string("home.engine.unsupportedOS")
+ }
+ if config.isLocalEngine {
+ if isAnyModelDownloading {
+ return AppL10n.string("home.engine.downloading")
+ }
+ if !OnDeviceModelStatus.isLocalStackReady(asrBackend: config.localASRBackend) {
+ return AppL10n.string("home.engine.downloadFirst")
+ }
+ switch modelWarmup.phase {
+ case .warming:
+ return AppL10n.string("home.engine.warming")
+ case .failed(let message):
+ return message
+ case .idle:
+ return AppL10n.string("home.engine.warming")
+ case .ready, .notNeeded:
+ break
+ }
+ }
+ if flowManager.isStarting {
+ return AppL10n.string("home.flow.starting")
+ }
+ if flowManager.isActive {
+ return AppL10n.string("home.flow.label")
+ }
+ return AppL10n.string("home.flow.inactive")
}
// MARK: - Preview field
@@ -304,6 +412,9 @@ struct HomeView: View {
RoundedRectangle(cornerRadius: Radius.large, style: .continuous)
.stroke(previewFocused ? palette.dividerStrong : palette.divider, lineWidth: 0.5)
)
+ // TextField only hit-tests the text line(s); expand taps to the full card.
+ .contentShape(RoundedRectangle(cornerRadius: Radius.large, style: .continuous))
+ .onTapGesture { previewFocused = true }
}
private var engineStatusLine: some View {
@@ -311,7 +422,8 @@ struct HomeView: View {
EngineServiceLabel.summary(
engineMode: config.engineMode,
providerId: config.providerId,
- model: config.model
+ model: config.model,
+ localASRBackend: config.localASRBackend
)
)
.font(TypeStyle.caption2)
diff --git a/OSGKeyboard/Views/KeyboardPreviewSheet.swift b/OSGKeyboard/Views/KeyboardPreviewSheet.swift
index 5bebf61..cbc809d 100644
--- a/OSGKeyboard/Views/KeyboardPreviewSheet.swift
+++ b/OSGKeyboard/Views/KeyboardPreviewSheet.swift
@@ -108,16 +108,24 @@ struct KeyboardPreviewSheet: View {
Button {
toggleRecording()
} label: {
- Label(isRecording ? "停止录音" : "开始录音", systemImage: isRecording ? "stop.circle.fill" : "mic.circle.fill")
- .primaryButton()
+ Label {
+ Text(isRecording ? "preview.record.stop" : "preview.record.start")
+ } icon: {
+ Image(systemName: isRecording ? "stop.circle.fill" : "mic.circle.fill")
+ }
+ .primaryButton()
}
.buttonStyle(.plain)
Button {
showSettings = true
} label: {
- Label("设置", systemImage: "gearshape")
- .secondaryButton()
+ Label {
+ Text("preview.openSettings")
+ } icon: {
+ Image(systemName: "gearshape")
+ }
+ .secondaryButton()
}
.buttonStyle(.plain)
}
@@ -135,7 +143,8 @@ struct KeyboardPreviewSheet: View {
EngineServiceLabel.summary(
engineMode: config.engineMode,
providerId: config.providerId,
- model: config.model
+ model: config.model,
+ localASRBackend: config.localASRBackend
)
}
@@ -153,10 +162,16 @@ struct KeyboardPreviewSheet: View {
private var statusTitle: String {
switch dictation.phase {
- case .idle: return "输入框"
- case .requestingPermission: return "请求权限中..."
- case .recording: return config.isLocalEngine ? "正在实时识别..." : "正在录音..."
- case .processing: return "识别中..."
+ case .idle:
+ return AppL10n.string("preview.status.field")
+ case .requestingPermission:
+ return AppL10n.string("dictation.status.requestingPermission")
+ case .recording:
+ return config.isLocalEngine
+ ? AppL10n.string("dictation.status.listeningLive")
+ : AppL10n.string("preview.status.recording")
+ case .processing:
+ return AppL10n.string("preview.status.processing")
case .denied(let message), .error(let message):
return message
}
diff --git a/OSGKeyboard/Views/LocalEngineSettingsRows.swift b/OSGKeyboard/Views/LocalEngineSettingsRows.swift
new file mode 100644
index 0000000..1056d38
--- /dev/null
+++ b/OSGKeyboard/Views/LocalEngineSettingsRows.swift
@@ -0,0 +1,91 @@
+// LocalEngineSettingsRows.swift
+// OSGKeyboard · Main App
+//
+// Grouped "local models" block for the local engine settings card.
+// Speech recognition row with a compact readiness summary in the header.
+
+import SwiftUI
+import OSGKeyboardShared
+
+// MARK: - Local models group (scheme A)
+
+struct LocalModelsGroup: View {
+ @Environment(\.themePalette) private var palette: ThemePalette
+
+ @ObservedObject var config: ProviderConfig
+ @ObservedObject var manager: ModelManager
+ @Binding var pendingDownload: OnDeviceModel?
+
+ var body: some View {
+ speechRow
+ }
+
+ // MARK: Speech row
+
+ private var speechRow: some View {
+ HStack(spacing: Spacing.xs) {
+ Text("settings.localModels.speechRole")
+ .font(TypeStyle.body)
+ .foregroundStyle(palette.textPrimary)
+ speechBackendMenu
+ Spacer(minLength: Spacing.xs)
+ speechTrailing
+ }
+ .padding(.horizontal, Spacing.md)
+ .frame(minHeight: SettingsListMetrics.singleLineMinHeight)
+ }
+
+ private var speechBackendMenu: some View {
+ Menu {
+ ForEach(LocalASRBackend.allCases) { backend in
+ Button {
+ config.localASRBackend = backend
+ } label: {
+ if backend == config.localASRBackend {
+ Label(
+ AppL10n.string(backend.labelKey),
+ systemImage: "checkmark"
+ )
+ } else {
+ Text(AppL10n.string(backend.labelKey))
+ }
+ }
+ }
+ } label: {
+ HStack(spacing: 4) {
+ Text(AppL10n.string(config.localASRBackend.labelKey))
+ .font(TypeStyle.body)
+ .foregroundStyle(palette.textSecondary)
+ .lineLimit(1)
+ Image(systemName: "chevron.up.chevron.down")
+ .font(.system(size: 11, weight: .bold))
+ .foregroundStyle(palette.textTertiary)
+ }
+ }
+ }
+
+ @ViewBuilder
+ private var speechTrailing: some View {
+ if config.localASRBackend == .speechAnalyzer {
+ builtInBadge
+ } else {
+ ModelListActionButton(
+ model: .qwen3ASR,
+ manager: manager,
+ pendingDownload: $pendingDownload
+ )
+ }
+ }
+
+ // MARK: Helpers
+
+ private var builtInBadge: some View {
+ HStack(spacing: 4) {
+ Image(systemName: "checkmark.circle.fill")
+ .font(.system(size: 12, weight: .semibold))
+ Text("settings.localModels.builtIn")
+ .font(TypeStyle.caption)
+ }
+ .foregroundStyle(palette.accent)
+ }
+}
diff --git a/OSGKeyboard/Views/OnDeviceModelsView.swift b/OSGKeyboard/Views/OnDeviceModelsView.swift
new file mode 100644
index 0000000..e683a7d
--- /dev/null
+++ b/OSGKeyboard/Views/OnDeviceModelsView.swift
@@ -0,0 +1,184 @@
+// OnDeviceModelsView.swift
+// OSGKeyboard · Main App
+//
+// Full-page model download manager. The same `OnDeviceModelsContent`
+// is also embedded inline in `SettingsView` so users see manual
+// download controls without drilling into About.
+
+import SwiftUI
+import OSGKeyboardShared
+
+/// Inline + full-page body for on-device model downloads.
+struct OnDeviceModelsContent: View {
+ @Environment(\.themePalette) private var palette: ThemePalette
+ @ObservedObject var manager: ModelManager
+ @Binding var pendingDownload: OnDeviceModel?
+
+ init(
+ manager: ModelManager = .shared,
+ pendingDownload: Binding
+ ) {
+ _manager = ObservedObject(wrappedValue: manager)
+ _pendingDownload = pendingDownload
+ }
+
+ var body: some View {
+ Group {
+ ForEach(Array(OnDeviceModel.allCases.enumerated()), id: \.element.id) { index, model in
+ if index > 0 {
+ Divider().background(palette.divider)
+ }
+ OnDeviceModelListRow(
+ model: model,
+ manager: manager,
+ pendingDownload: $pendingDownload
+ )
+ }
+ }
+ }
+}
+
+// MARK: - Model row
+
+struct OnDeviceModelListRow: View {
+ @Environment(\.themePalette) private var palette: ThemePalette
+
+ let model: OnDeviceModel
+ @ObservedObject var manager: ModelManager
+ @Binding var pendingDownload: OnDeviceModel?
+
+ var body: some View {
+ HStack(spacing: Spacing.sm) {
+ Text(model.listTitle)
+ .font(TypeStyle.body)
+ .foregroundStyle(palette.textPrimary)
+ .lineLimit(1)
+ Spacer(minLength: Spacing.xs)
+ ModelListActionButton(
+ model: model,
+ manager: manager,
+ pendingDownload: $pendingDownload
+ )
+ }
+ .padding(.horizontal, Spacing.md)
+ .frame(minHeight: SettingsListMetrics.singleLineMinHeight)
+ }
+}
+
+// MARK: - Action button
+
+struct ModelListActionButton: View {
+ @Environment(\.themePalette) private var palette: ThemePalette
+
+ let model: OnDeviceModel
+ @ObservedObject var manager: ModelManager
+ @Binding var pendingDownload: OnDeviceModel?
+
+ var body: some View {
+ let state = manager.states[model]?.download ?? .notDownloaded
+ Button(action: { performAction(for: state) }) {
+ Text(buttonTitle(for: state))
+ .font(TypeStyle.caption)
+ .foregroundStyle(buttonColor(for: state))
+ .monospacedDigit()
+ .padding(.horizontal, Spacing.sm)
+ .padding(.vertical, 5)
+ .background(buttonBackground(for: state), in: Capsule())
+ .overlay(
+ Capsule().stroke(buttonBorder(for: state), lineWidth: 0.5)
+ )
+ }
+ .buttonStyle(.plain)
+ }
+
+ private func performAction(for state: ModelDownloadState) {
+ switch state {
+ case .notDownloaded, .failed:
+ pendingDownload = model
+ case .downloading:
+ manager.cancelDownload(model)
+ case .downloaded:
+ manager.deleteModel(model)
+ }
+ }
+
+ private func buttonTitle(for state: ModelDownloadState) -> String {
+ let size = model.compactSizeLabel
+ switch state {
+ case .notDownloaded, .failed:
+ return AppL10n.format("settings.models.action.downloadSize %@", size)
+ case .downloading(let progress):
+ return AppL10n.format(
+ "settings.models.action.downloadingPercent %lld",
+ Int((progress * 100).rounded())
+ )
+ case .downloaded:
+ return AppL10n.format("settings.models.action.deleteSize %@", size)
+ }
+ }
+
+ private func buttonColor(for state: ModelDownloadState) -> Color {
+ switch state {
+ case .notDownloaded, .failed:
+ return palette.textOnAccent
+ case .downloading:
+ return palette.accent
+ case .downloaded:
+ return palette.warning
+ }
+ }
+
+ private func buttonBackground(for state: ModelDownloadState) -> Color {
+ switch state {
+ case .notDownloaded, .failed:
+ return palette.accent
+ case .downloading:
+ return palette.accent.opacity(0.12)
+ case .downloaded:
+ return palette.warning.opacity(0.1)
+ }
+ }
+
+ private func buttonBorder(for state: ModelDownloadState) -> Color {
+ switch state {
+ case .notDownloaded, .failed:
+ return .clear
+ case .downloading:
+ return palette.accent.opacity(0.35)
+ case .downloaded:
+ return palette.warning.opacity(0.35)
+ }
+ }
+}
+
+// MARK: - Full page
+
+struct OnDeviceModelsView: View {
+ @Environment(\.themePalette) private var palette: ThemePalette
+ @StateObject private var manager = ModelManager.shared
+ @State private var pendingDownload: OnDeviceModel?
+
+ var body: some View {
+ ScrollView {
+ VStack(spacing: 0) {
+ OnDeviceModelsContent(manager: manager, pendingDownload: $pendingDownload)
+ }
+ .background(palette.surface, in: RoundedRectangle(cornerRadius: Radius.large, style: .continuous))
+ .overlay(
+ RoundedRectangle(cornerRadius: Radius.large, style: .continuous)
+ .stroke(palette.divider, lineWidth: 0.5)
+ )
+ .padding(.horizontal, Spacing.md)
+ .padding(.vertical, Spacing.md)
+ }
+ .background(palette.background.ignoresSafeArea())
+ .navigationTitle("settings.models.title")
+ .navigationBarTitleDisplayMode(.inline)
+ .sheet(item: $pendingDownload) { model in
+ DownloadConfirmSheet(model: model) {
+ pendingDownload = nil
+ manager.startDownload(model)
+ }
+ }
+ }
+}
diff --git a/OSGKeyboard/Views/OnboardingView.swift b/OSGKeyboard/Views/OnboardingView.swift
index 8f3e446..0be4773 100644
--- a/OSGKeyboard/Views/OnboardingView.swift
+++ b/OSGKeyboard/Views/OnboardingView.swift
@@ -27,9 +27,12 @@ struct OnboardingView: View {
@Environment(\.scenePhase) private var scenePhase
@ObservedObject var config: ProviderConfig
+ @ObservedObject private var modelWarmup = OnDeviceModelWarmup.shared
+ @StateObject private var modelManager = ModelManager.shared
@State private var micStatus = AppPermissions.micStatus
@State private var speechStatus = AppPermissions.speechStatus
@State private var keyboardReady = KeyboardSetupBridge.isReadyForOnboardingSkip
+ @State private var pendingDownload: OnDeviceModel?
private var currentPage: OnboardingPage {
OnboardingPage(rawValue: config.onboardingPage) ?? .welcome
@@ -58,7 +61,11 @@ struct OnboardingView: View {
case .keyboard:
EnableKeyboardPage()
case .api:
- APISetupPage(config: config)
+ APISetupPage(
+ config: config,
+ manager: modelManager,
+ pendingDownload: $pendingDownload
+ )
}
}
.frame(maxWidth: .infinity, maxHeight: .infinity)
@@ -84,6 +91,12 @@ struct OnboardingView: View {
applyOnboardingDefaultsIfNeeded()
refreshPermissionStatuses()
snapToVisiblePageIfNeeded()
+ scheduleModelWarmup()
+ }
+ .onChange(of: config.engineMode) { _, _ in scheduleModelWarmup(force: true) }
+ .onChange(of: config.localASRBackend) { _, _ in
+ modelWarmup.invalidate()
+ scheduleModelWarmup(force: true)
}
.onChange(of: scenePhase) { _, phase in
if phase == .active { refreshPermissionStatuses() }
@@ -105,8 +118,29 @@ struct OnboardingView: View {
// First-time users with no API key: default to local for a faster path.
if config.apiKey.isEmpty, config.engineMode == "cloud" {
config.engineMode = "local"
- config.modeId = "transcribe"
}
+ if config.engineMode == "local" {
+ config.localASRBackend = .qwen3ASR
+ }
+ }
+
+ private var localModelNeedsAttention: Bool {
+ guard config.isLocalEngine else { return false }
+ if !OnDeviceModelStatus.isLocalStackReady(asrBackend: config.localASRBackend) {
+ return true
+ }
+ switch modelWarmup.phase {
+ case .ready, .notNeeded, .idle: return false
+ case .warming, .failed: return true
+ }
+ }
+
+ private func scheduleModelWarmup(force: Bool = false) {
+ guard config.isLocalEngine else {
+ modelWarmup.invalidate()
+ return
+ }
+ modelWarmup.warmUpIfNeeded(force: force)
}
private func shouldShowPage(_ page: OnboardingPage) -> Bool {
@@ -174,16 +208,23 @@ struct OnboardingView: View {
private func onboardingHeaderGradient(height: CGFloat) -> some View {
LinearGradient(
- colors: [
- palette.textTertiary.opacity(0.14),
- palette.textTertiary.opacity(0.05),
- palette.background.opacity(0)
- ],
+ colors: localModelNeedsAttention
+ ? [
+ palette.warning.opacity(0.32),
+ palette.warning.opacity(0.12),
+ palette.background.opacity(0)
+ ]
+ : [
+ palette.textTertiary.opacity(0.14),
+ palette.textTertiary.opacity(0.05),
+ palette.background.opacity(0)
+ ],
startPoint: .top,
endPoint: .bottom
)
.frame(maxWidth: .infinity)
.frame(height: height)
+ .animation(Motion.soft, value: localModelNeedsAttention)
}
private func refreshPermissionStatuses() {
@@ -195,7 +236,7 @@ struct OnboardingView: View {
private var progressHeader: some View {
Text(
String(
- format: NSLocalizedString("onboarding.progress", comment: ""),
+ format: AppL10n.string("onboarding.progress"),
config.onboardingPage + 1,
OnboardingPage.count
)
@@ -220,8 +261,10 @@ struct OnboardingView: View {
private var canAdvance: Bool {
switch currentPage {
- case .welcome, .keyboard, .api:
- return !isLastPage || config.isConfigured
+ case .welcome, .keyboard:
+ return !isLastPage || onboardingCompleteReady
+ case .api:
+ return !isLastPage || onboardingCompleteReady
case .microphone:
return micStatus != .undetermined
case .speech:
@@ -229,16 +272,23 @@ struct OnboardingView: View {
}
}
+ private var onboardingCompleteReady: Bool {
+ if config.isLocalEngine {
+ return OnDeviceModelStatus.isLocalStackReady(asrBackend: config.localASRBackend)
+ }
+ return config.isConfigured
+ }
+
private var primaryActionTitle: String {
if isLastPage {
- return NSLocalizedString("common.done", comment: "")
+ return AppL10n.string("common.done")
}
switch currentPage {
case .microphone where micStatus == .granted,
.speech where speechStatus == .granted:
- return NSLocalizedString("common.continue", comment: "")
+ return AppL10n.string("common.continue")
default:
- return NSLocalizedString("common.next", comment: "")
+ return AppL10n.string("common.next")
}
}
@@ -609,8 +659,8 @@ private struct EnableKeyboardPage: View {
.padding(.top, OnboardingLayoutMetrics.heroTextGap)
VStack(alignment: .leading, spacing: Spacing.lg) {
- step(num: 1, text: NSLocalizedString("onboarding.enable.step1", comment: ""))
- step(num: 2, text: NSLocalizedString("onboarding.enable.step2", comment: ""))
+ step(num: 1, text: AppL10n.string("onboarding.enable.step1"))
+ step(num: 2, text: AppL10n.string("onboarding.enable.step2"))
switchKeyboardStep(num: 3)
}
.frame(maxWidth: .infinity, alignment: .leading)
@@ -686,6 +736,8 @@ private struct APISetupPage: View {
@Environment(\.themePalette) private var palette: ThemePalette
@ObservedObject var config: ProviderConfig
+ @ObservedObject var manager: ModelManager
+ @Binding var pendingDownload: OnDeviceModel?
var body: some View {
ScrollView {
@@ -704,9 +756,45 @@ private struct APISetupPage: View {
.padding(.horizontal, Spacing.lg)
APISettingsCard(config: config)
.padding(.horizontal, Spacing.lg)
+ } else {
+ VStack(alignment: .leading, spacing: Spacing.sm) {
+ Text("onboarding.api.localModels.hint")
+ .font(TypeStyle.caption2)
+ .foregroundStyle(palette.warning)
+ .fixedSize(horizontal: false, vertical: true)
+ LocalModelsGroup(
+ config: config,
+ manager: manager,
+ pendingDownload: $pendingDownload
+ )
+ .background(palette.surface, in: RoundedRectangle(cornerRadius: Radius.large, style: .continuous))
+ .overlay(
+ RoundedRectangle(cornerRadius: Radius.large, style: .continuous)
+ .stroke(palette.divider, lineWidth: 0.5)
+ )
+ }
+ .padding(.horizontal, Spacing.lg)
}
}
.padding(.bottom, Spacing.xxxl)
}
+ .onAppear {
+ manager.refreshAll()
+ if config.engineMode == "local" {
+ config.localASRBackend = .qwen3ASR
+ }
+ }
+ .onChange(of: config.engineMode) { _, mode in
+ guard mode == "local" else { return }
+ Task { @MainActor in
+ config.localASRBackend = .qwen3ASR
+ }
+ }
+ .sheet(item: $pendingDownload) { model in
+ DownloadConfirmSheet(model: model) {
+ pendingDownload = nil
+ manager.startDownload(model)
+ }
+ }
}
}
diff --git a/OSGKeyboard/Views/OpenSourceLicensesView.swift b/OSGKeyboard/Views/OpenSourceLicensesView.swift
new file mode 100644
index 0000000..2b1939c
--- /dev/null
+++ b/OSGKeyboard/Views/OpenSourceLicensesView.swift
@@ -0,0 +1,115 @@
+// OpenSourceLicensesView.swift
+// OSGKeyboard · Main App
+//
+// Standard acknowledgements list: one tappable row per dependency,
+// detail screen with the verbatim license text. Reached from
+// Settings → About → "Third-Party Licenses".
+
+import SwiftUI
+import OSGKeyboardShared
+
+struct OpenSourceLicensesView: View {
+ @Environment(\.themePalette) private var palette: ThemePalette
+
+ var body: some View {
+ ScrollView {
+ VStack(alignment: .leading, spacing: SettingsListMetrics.sectionLabelSpacing) {
+ Text("settings.licenses.footer")
+ .font(TypeStyle.caption2)
+ .foregroundStyle(palette.textTertiary)
+ .fixedSize(horizontal: false, vertical: true)
+ .padding(.horizontal, Spacing.xs)
+
+ VStack(spacing: 0) {
+ ForEach(Array(OpenSourceLicenseCatalog.entries.enumerated()), id: \.element.id) { index, entry in
+ NavigationLink {
+ OpenSourceLicenseDetailView(entry: entry)
+ } label: {
+ licenseRow(entry)
+ }
+ .buttonStyle(.plain)
+
+ if index < OpenSourceLicenseCatalog.entries.count - 1 {
+ Divider().background(palette.divider)
+ }
+ }
+ }
+ .background(palette.surface, in: RoundedRectangle(cornerRadius: Radius.large, style: .continuous))
+ .overlay(
+ RoundedRectangle(cornerRadius: Radius.large, style: .continuous)
+ .stroke(palette.divider, lineWidth: 0.5)
+ )
+ }
+ .padding(.horizontal, Spacing.md)
+ .padding(.vertical, Spacing.md)
+ }
+ .background(palette.background.ignoresSafeArea())
+ .navigationTitle("settings.licenses.title")
+ .navigationBarTitleDisplayMode(.inline)
+ }
+
+ private func licenseRow(_ entry: OpenSourceLicenseCatalog.Entry) -> some View {
+ HStack(spacing: Spacing.xs) {
+ Text(entry.name)
+ .font(TypeStyle.caption2)
+ .foregroundStyle(palette.textPrimary)
+ .lineLimit(1)
+ .truncationMode(.middle)
+ Spacer(minLength: Spacing.xs)
+ Text(entry.licenseName)
+ .font(TypeStyle.caption2)
+ .foregroundStyle(palette.textTertiary)
+ Image(systemName: "chevron.right")
+ .font(.system(size: 11, weight: .semibold))
+ .foregroundStyle(palette.textTertiary)
+ }
+ .padding(.horizontal, Spacing.md)
+ .padding(.vertical, 10)
+ .contentShape(Rectangle())
+ }
+}
+
+// MARK: - Detail
+
+private struct OpenSourceLicenseDetailView: View {
+ @Environment(\.themePalette) private var palette: ThemePalette
+
+ let entry: OpenSourceLicenseCatalog.Entry
+
+ var body: some View {
+ ScrollView {
+ VStack(alignment: .leading, spacing: Spacing.sm) {
+ if let url = entry.url {
+ Link(destination: url) {
+ HStack(spacing: Spacing.xs) {
+ Text(url.absoluteString)
+ .font(TypeStyle.caption2)
+ .foregroundStyle(palette.accent)
+ .lineLimit(2)
+ .multilineTextAlignment(.leading)
+ Spacer(minLength: 0)
+ MaterialIcon(name: .openInNew, size: 14)
+ .foregroundStyle(palette.textTertiary)
+ }
+ }
+ }
+
+ Text(entry.purpose)
+ .font(TypeStyle.caption2)
+ .foregroundStyle(palette.textSecondary)
+ .fixedSize(horizontal: false, vertical: true)
+
+ Text(entry.licenseText)
+ .font(TypeStyle.monoSmall)
+ .foregroundStyle(palette.textPrimary)
+ .textSelection(.enabled)
+ .frame(maxWidth: .infinity, alignment: .leading)
+ }
+ .padding(.horizontal, Spacing.md)
+ .padding(.vertical, Spacing.md)
+ }
+ .background(palette.background.ignoresSafeArea())
+ .navigationTitle(entry.name)
+ .navigationBarTitleDisplayMode(.inline)
+ }
+}
diff --git a/OSGKeyboard/Views/PrivacyPolicyView.swift b/OSGKeyboard/Views/PrivacyPolicyView.swift
new file mode 100644
index 0000000..6ee1d71
--- /dev/null
+++ b/OSGKeyboard/Views/PrivacyPolicyView.swift
@@ -0,0 +1,33 @@
+// PrivacyPolicyView.swift
+// OSGKeyboard · Main App
+//
+// Bundled privacy policy HTML. Reached from Settings → About.
+
+import SwiftUI
+import OSGKeyboardShared
+
+struct PrivacyPolicyView: View {
+ @Environment(\.themePalette) private var palette: ThemePalette
+ @ObservedObject private var config = ProviderConfig.shared
+
+ var body: some View {
+ LegalWebView(
+ resourceName: "PrivacyPolicy",
+ scrollToAnchor: privacyScrollAnchor
+ )
+ .background(palette.background.ignoresSafeArea())
+ .navigationTitle("settings.privacy.policy")
+ .navigationBarTitleDisplayMode(.inline)
+ }
+
+ private var privacyScrollAnchor: String? {
+ switch config.uiLanguage {
+ case .chinese:
+ return "zh"
+ case .english:
+ return "top"
+ case .auto:
+ return config.uiLanguage.resolvedLanguageCode().hasPrefix("zh") ? "zh" : "top"
+ }
+ }
+}
diff --git a/OSGKeyboard/Views/SettingsView.swift b/OSGKeyboard/Views/SettingsView.swift
index c347a1a..ce319bc 100644
--- a/OSGKeyboard/Views/SettingsView.swift
+++ b/OSGKeyboard/Views/SettingsView.swift
@@ -18,7 +18,7 @@ struct SettingsView: View {
@ObservedObject var config = ProviderConfig.shared
@Environment(\.dismiss) private var dismiss
- @State private var safariURL: URL?
+ @Environment(\.openURL) private var openURL
let presentation: SettingsPresentation
@@ -27,7 +27,9 @@ struct SettingsView: View {
}
// Dynamic locale list loaded from SFSpeechRecognizer on first appear.
- @State private var dynamicLocales: [(id: String, label: String, onDevice: Bool)] = []
+ @State private var dynamicLocales: [(id: String, onDevice: Bool)] = []
+ @StateObject private var modelManager = ModelManager.shared
+ @State private var pendingDownload: OnDeviceModel?
var body: some View {
NavigationStack {
@@ -57,14 +59,15 @@ struct SettingsView: View {
palette.background.ignoresSafeArea()
ScrollView {
VStack(spacing: Spacing.md) {
+ appLanguageSection
engineSection
if config.engineMode == "cloud" {
providerSection
apiSection
}
- languageSection
+ languageAndModelsSection
if config.engineMode == "cloud" {
- promptSection
+ systemPromptLinkSection
}
if presentation == .tab {
footerLinks
@@ -79,54 +82,55 @@ struct SettingsView: View {
.background(palette.background)
.toolbar(.hidden, for: .navigationBar)
.task { await loadDynamicLocales() }
- .sheet(item: $safariURL) { url in
- SafariSheet(url: url)
+ .onAppear { modelManager.refreshAll() }
+ .sheet(item: $pendingDownload) { model in
+ DownloadConfirmSheet(model: model) {
+ pendingDownload = nil
+ modelManager.startDownload(model)
+ }
}
}
}
+ // MARK: - App language
+
+ private var appLanguageSection: some View {
+ VStack(alignment: .leading, spacing: SettingsListMetrics.sectionLabelSpacing) {
+ sectionHeader("settings.appLanguage.title")
+ Picker("", selection: $config.uiLanguage) {
+ ForEach(AppUILanguage.allCases) { language in
+ Text(LocalizedStringKey(language.labelKey)).tag(language)
+ }
+ }
+ .pickerStyle(.segmented)
+ .padding(Spacing.md)
+ .background(palette.surface, in: RoundedRectangle(cornerRadius: Radius.large, style: .continuous))
+ .overlay(
+ RoundedRectangle(cornerRadius: Radius.large, style: .continuous)
+ .stroke(palette.divider, lineWidth: 0.5)
+ )
+ }
+ }
+
// MARK: - Engine
private var engineSection: some View {
EnginePickerSection(config: config)
}
- // MARK: - Provider
+ // MARK: - Language & on-device models
- private var providerSection: some View {
- VStack(alignment: .leading, spacing: SettingsListMetrics.sectionLabelSpacing) {
- sectionHeader("settings.provider.title")
- ProviderPickerSection(config: config)
- }
- }
-
- // MARK: - API
-
- private var apiSection: some View {
- VStack(alignment: .leading, spacing: SettingsListMetrics.sectionLabelSpacing) {
- sectionHeader("settings.api.title")
- APISettingsCard(config: config)
- }
- }
-
- // MARK: - Language (ASR + mode)
-
- private var languageSection: some View {
+ private var languageAndModelsSection: some View {
VStack(alignment: .leading, spacing: SettingsListMetrics.sectionLabelSpacing) {
sectionHeader("settings.language.title")
VStack(spacing: 0) {
- if config.engineMode == "cloud" {
- PickerRow(
- title: NSLocalizedString("settings.mode.title", comment: ""),
- options: modeOptions,
- selection: Binding(
- get: { config.modeId },
- set: { config.modeId = $0 }
- )
+ if config.engineMode == "local" {
+ LocalModelsGroup(
+ config: config,
+ manager: modelManager,
+ pendingDownload: $pendingDownload
)
Divider().background(palette.divider)
- asrEngineRow
- Divider().background(palette.divider)
}
LocalePickerRow(
locales: effectiveLocales,
@@ -165,53 +169,37 @@ struct SettingsView: View {
}
}
- /// ASR engine row. With iOS 26 as the deployment target, the
- /// only ASR backend is `SpeechAnalyzer` and it is always fully
- /// on-device — so this row is now a static badge rather than a
- /// toggle. The old cloud-fallback toggle is gone.
- private var asrEngineRow: some View {
- HStack(spacing: Spacing.sm) {
- Text("settings.engineRow.title")
- .font(TypeStyle.body)
- .foregroundStyle(palette.textPrimary)
- Spacer()
- HStack(spacing: 4) {
- Image(systemName: "iphone.badge.checkmark")
- .font(.system(size: 12, weight: .semibold))
- .foregroundStyle(palette.accent)
- Text("settings.engineBadge.ios26")
- .font(TypeStyle.caption)
- .foregroundStyle(palette.accent)
- }
- .padding(.horizontal, Spacing.xs)
- .padding(.vertical, 4)
- .background(palette.accent.opacity(0.12), in: Capsule())
+ private var providerSection: some View {
+ VStack(alignment: .leading, spacing: SettingsListMetrics.sectionLabelSpacing) {
+ sectionHeader("settings.provider.title")
+ ProviderPickerSection(config: config)
}
- .padding(.horizontal, Spacing.md)
- .frame(minHeight: SettingsListMetrics.singleLineMinHeight)
}
+ // MARK: - API
+
+ private var apiSection: some View {
+ VStack(alignment: .leading, spacing: SettingsListMetrics.sectionLabelSpacing) {
+ sectionHeader("settings.api.title")
+ APISettingsCard(config: config)
+ }
+ }
+
+ // MARK: - Language helpers
+
/// Falls back to a static list while dynamic locales are loading.
- private var effectiveLocales: [(id: String, label: String, onDevice: Bool)] {
+ private var effectiveLocales: [(id: String, onDevice: Bool)] {
dynamicLocales.isEmpty ? staticLocales : dynamicLocales
}
- private var staticLocales: [(id: String, label: String, onDevice: Bool)] {
+ private var staticLocales: [(id: String, onDevice: Bool)] {
[
- ("auto", NSLocalizedString("locale.auto", comment: ""), false),
- ("zh-Hans", NSLocalizedString("locale.zh-Hans", comment: ""), false),
- ("zh-Hant", NSLocalizedString("locale.zh-Hant", comment: ""), false),
- ("en-US", NSLocalizedString("locale.en-US", comment: ""), false),
- ("ja-JP", NSLocalizedString("locale.ja-JP", comment: ""), false),
- ("ko-KR", NSLocalizedString("locale.ko-KR", comment: ""), false)
- ]
- }
-
- private var modeOptions: [(id: String, label: String)] {
- [
- ("off", NSLocalizedString("settings.mode.off", comment: "")),
- ("transcribe", NSLocalizedString("settings.mode.transcribe", comment: "")),
- ("polish", NSLocalizedString("settings.mode.polish", comment: ""))
+ ("auto", false),
+ ("zh-Hans", false),
+ ("zh-Hant", false),
+ ("en-US", false),
+ ("ja-JP", false),
+ ("ko-KR", false),
]
}
@@ -222,20 +210,16 @@ struct SettingsView: View {
// can return 100+ locales, and we probe supportsOnDeviceRecognition for each.
// Creating `SFSpeechRecognizer` instances in a @Sendable closure is
// safe here; we only read locale metadata (no transcription session).
- let entries: [(id: String, label: String, onDevice: Bool)] = await Task.detached(
+ let entries: [(id: String, onDevice: Bool)] = await Task.detached(
priority: .userInitiated
) {
- var result: [(id: String, label: String, onDevice: Bool)] = []
- result.append(("auto", NSLocalizedString("locale.auto", comment: ""), false))
+ var result: [(id: String, onDevice: Bool)] = [("auto", false)]
- let currentLocale = Locale.current // snapshot on background thread is fine
for locale in SFSpeechRecognizer.supportedLocales()
.sorted(by: { $0.identifier < $1.identifier }) {
let id = locale.identifier
let onDevice = SFSpeechRecognizer(locale: locale)?.supportsOnDeviceRecognition ?? false
- // localizedString on Locale.current gives the name in the app's UI language.
- let displayName = currentLocale.localizedString(forIdentifier: id) ?? id
- result.append((id: id, label: displayName, onDevice: onDevice))
+ result.append((id: id, onDevice: onDevice))
}
return result
}.value
@@ -244,30 +228,24 @@ struct SettingsView: View {
dynamicLocales = entries
}
- // MARK: - Prompt
+ // MARK: - System prompt (cloud only)
- private var promptSection: some View {
+ private var systemPromptLinkSection: some View {
VStack(alignment: .leading, spacing: SettingsListMetrics.sectionLabelSpacing) {
- HStack {
- sectionHeader("settings.systemPrompt.title")
- Spacer()
- Button("common.reset") { config.systemPrompt = config.defaultSystemPrompt }
- .font(TypeStyle.caption2)
- .foregroundStyle(palette.accent)
+ sectionHeader("settings.systemPrompt.title")
+ VStack(spacing: 0) {
+ NavigationLink {
+ SystemPromptSettingsView(config: config)
+ } label: {
+ footerNavigationRow(title: "settings.systemPrompt.edit")
+ }
+ .buttonStyle(.plain)
}
- VStack(alignment: .leading, spacing: Spacing.xs) {
- TextEditor(text: $config.systemPrompt)
- .font(TypeStyle.mono)
- .scrollContentBackground(.hidden)
- .frame(minHeight: 140)
- .padding(Spacing.xs)
- .background(palette.surface, in: RoundedRectangle(cornerRadius: Radius.medium, style: .continuous))
- .overlay(
- RoundedRectangle(cornerRadius: Radius.medium, style: .continuous)
- .stroke(palette.divider, lineWidth: 0.5)
- )
- }
- .cardSurface()
+ .background(palette.surface, in: RoundedRectangle(cornerRadius: Radius.large, style: .continuous))
+ .overlay(
+ RoundedRectangle(cornerRadius: Radius.large, style: .continuous)
+ .stroke(palette.divider, lineWidth: 0.5)
+ )
}
}
@@ -298,18 +276,34 @@ struct SettingsView: View {
Divider().background(palette.divider)
- if let url = LegalLinks.privacyPolicyURL {
- footerLinkRow(title: "settings.privacy.policy", url: url)
- Divider().background(palette.divider)
+ NavigationLink {
+ PrivacyPolicyView()
+ } label: {
+ footerNavigationRow(title: "settings.privacy.policy")
}
- if let url = LegalLinks.supportURL {
- footerLinkRow(title: "settings.link.support", url: url)
- Divider().background(palette.divider)
+ .buttonStyle(.plain)
+ Divider().background(palette.divider)
+
+ NavigationLink {
+ HelpFeedbackView()
+ } label: {
+ footerNavigationRow(title: "settings.link.support")
}
- footerLinkRow(
+ .buttonStyle(.plain)
+ Divider().background(palette.divider)
+
+ footerExternalLinkRow(
title: "settings.link.github",
- url: URL(string: "https://github.com/hkgood/OSGKeyboard")!
+ url: LegalLinks.repositoryURL
)
+ Divider().background(palette.divider)
+
+ NavigationLink {
+ OpenSourceLicensesView()
+ } label: {
+ footerNavigationRow(title: "settings.link.licenses")
+ }
+ .buttonStyle(.plain)
}
.background(palette.surface, in: RoundedRectangle(cornerRadius: Radius.xl, style: .continuous))
.overlay(
@@ -319,9 +313,9 @@ struct SettingsView: View {
}
}
- private func footerLinkRow(title: LocalizedStringKey, url: URL) -> some View {
+ private func footerExternalLinkRow(title: LocalizedStringKey, url: URL) -> some View {
Button {
- safariURL = url
+ openURL(url)
} label: {
HStack(spacing: Spacing.sm) {
Text(title)
@@ -338,6 +332,25 @@ struct SettingsView: View {
.buttonStyle(.plain)
}
+ /// In-app disclosure row that pushes a child view onto the
+ /// `NavigationStack` rather than opening Safari. Used for the
+ /// Third-Party Licenses entry so the system "back" button
+ /// returns to Settings.
+ private func footerNavigationRow(title: LocalizedStringKey) -> some View {
+ HStack(spacing: Spacing.sm) {
+ Text(title)
+ .font(TypeStyle.body)
+ .foregroundStyle(palette.textPrimary)
+ Spacer()
+ Image(systemName: "chevron.right")
+ .font(.system(size: 14, weight: .semibold))
+ .foregroundStyle(palette.textTertiary)
+ }
+ .padding(.horizontal, Spacing.md)
+ .frame(minHeight: SettingsListMetrics.singleLineMinHeight)
+ .contentShape(Rectangle())
+ }
+
// MARK: - Header
private func sectionHeader(_ title: LocalizedStringKey) -> some View {
@@ -400,8 +413,9 @@ private struct PickerRow: View {
private struct LocalePickerRow: View {
@Environment(\.themePalette) private var palette: ThemePalette
+ @ObservedObject private var config = ProviderConfig.shared
- let locales: [(id: String, label: String, onDevice: Bool)]
+ let locales: [(id: String, onDevice: Bool)]
@Binding var selection: String
var body: some View {
@@ -417,12 +431,13 @@ private struct LocalePickerRow: View {
} label: {
// iOS Menu converts SwiftUI Label to UIAction (title + image).
// Using Label keeps checkmark + on-device icon both visible.
+ let name = label(for: locale.id)
if locale.id == selection {
- Label(locale.label, systemImage: "checkmark")
+ Label(name, systemImage: "checkmark")
} else if locale.onDevice {
- Label(locale.label, systemImage: "iphone")
+ Label(name, systemImage: "iphone")
} else {
- Text(locale.label)
+ Text(name)
}
}
}
@@ -447,7 +462,11 @@ private struct LocalePickerRow: View {
.frame(minHeight: SettingsListMetrics.singleLineMinHeight)
}
+ private func label(for localeId: String) -> String {
+ ASRLocaleLabels.displayName(for: localeId, language: config.uiLanguage)
+ }
+
private var currentLabel: String {
- locales.first(where: { $0.id == selection })?.label ?? "—"
+ label(for: selection)
}
}
diff --git a/OSGKeyboard/Views/SystemPromptSettingsView.swift b/OSGKeyboard/Views/SystemPromptSettingsView.swift
new file mode 100644
index 0000000..8612363
--- /dev/null
+++ b/OSGKeyboard/Views/SystemPromptSettingsView.swift
@@ -0,0 +1,53 @@
+// SystemPromptSettingsView.swift
+// OSGKeyboard · Main App
+//
+// Cloud-engine system prompt editor. Reached from Settings when the
+// user picks the cloud recognition path.
+
+import SwiftUI
+import OSGKeyboardShared
+
+struct SystemPromptSettingsView: View {
+ @Environment(\.themePalette) private var palette: ThemePalette
+
+ @ObservedObject var config: ProviderConfig
+
+ var body: some View {
+ ScrollView {
+ VStack(alignment: .leading, spacing: Spacing.sm) {
+ Text("settings.systemPrompt.hint")
+ .font(TypeStyle.caption2)
+ .foregroundStyle(palette.textTertiary)
+ .fixedSize(horizontal: false, vertical: true)
+
+ TextEditor(text: $config.systemPrompt)
+ .font(TypeStyle.mono)
+ .scrollContentBackground(.hidden)
+ .frame(minHeight: 320)
+ .padding(Spacing.sm)
+ .background(palette.surface, in: RoundedRectangle(cornerRadius: Radius.large, style: .continuous))
+ .overlay(
+ RoundedRectangle(cornerRadius: Radius.large, style: .continuous)
+ .stroke(palette.divider, lineWidth: 0.5)
+ )
+ }
+ .padding(.horizontal, Spacing.md)
+ .padding(.vertical, Spacing.md)
+ }
+ .background(palette.background.ignoresSafeArea())
+ .navigationTitle("settings.systemPrompt.title")
+ .navigationBarTitleDisplayMode(.inline)
+ .toolbar(.visible, for: .navigationBar)
+ .toolbar {
+ ToolbarItem(placement: .topBarTrailing) {
+ Button("common.reset") {
+ config.systemPrompt = config.defaultSystemPrompt
+ }
+ .font(TypeStyle.body)
+ .foregroundStyle(palette.accent)
+ }
+ }
+ .toolbarBackground(palette.background, for: .navigationBar)
+ .toolbarBackground(.visible, for: .navigationBar)
+ }
+}
diff --git a/OSGKeyboard/en.lproj/Localizable.strings b/OSGKeyboard/en.lproj/Localizable.strings
index 1784fe1..0c6d9d5 100644
--- a/OSGKeyboard/en.lproj/Localizable.strings
+++ b/OSGKeyboard/en.lproj/Localizable.strings
@@ -29,6 +29,7 @@
"onboarding.enable.step3.suffix" = "and select OSGKeyboard";
"onboarding.enable.openSettings" = "Open Settings";
"onboarding.api.title" = "Choose Engine";
+"onboarding.api.localModels.hint" = "Download the on-device CoreML ASR model below before you finish setup (only needed when using Qwen3-ASR).";
"settings.onboarding.replay" = "Restart permission setup";
/* Common navigation */
@@ -83,15 +84,19 @@
/* Settings */
"settings.title" = "Settings";
+"settings.appLanguage.title" = "App language";
+"settings.appLanguage.auto" = "Auto";
+"settings.appLanguage.english" = "English";
+"settings.appLanguage.chinese" = "Chinese";
"settings.reset.title" = "Reset all settings?";
"settings.reset.message" = "API key, model, base URL, voice history, and cloud consent will be cleared.";
"settings.reset.confirm" = "Reset all settings";
-"settings.engine.title" = "Engine";
-"settings.engine.subtitle" = "Pick the recognition engine. Local engine does transcription only, no API key needed.";
-"settings.engine.local.title" = "On-device";
+"settings.engine.title" = "Recognition method";
+"settings.engine.subtitle" = "Local: on-device ASR only. Cloud: ASR + polish via your API.";
+"settings.engine.local.title" = "On-device recognition";
"settings.engine.local.ios26" = "Always on-device, no network.";
"settings.engine.local.legacy" = "On-device ASR. Transcription only, no polish.";
-"settings.engine.cloud.title" = "Cloud polish";
+"settings.engine.cloud.title" = "Cloud recognition & polish";
"settings.engine.cloud.subtitle" = "On-device ASR + polish via your API. Text is sent to the endpoint you configure.";
"settings.provider.title" = "Provider";
"settings.provider.subtitle" = "Pick the LLM that polishes your dictation.";
@@ -103,6 +108,12 @@
"provider.custom" = "Custom";
"settings.api.title" = "API";
"settings.language.title" = "Language";
+"settings.languageModels.title" = "Language & models";
+"settings.localModels.title" = "On-device models";
+"settings.localModels.speechRole" = "Speech";
+"settings.localModels.builtIn" = "Built-in";
+"settings.localModels.allReady" = "Ready";
+"settings.localModels.readiness %lld %lld" = "%lld/%lld ready";
"settings.language.subtitle.cloud" = "Recognition language and text processing mode.";
"settings.language.subtitle.local" = "Recognition language.";
"settings.mode.title" = "Mode";
@@ -110,6 +121,8 @@
"settings.mode.transcribe" = "Transcribe";
"settings.mode.polish" = "Polish";
"settings.systemPrompt.title" = "System Prompt";
+"settings.systemPrompt.edit" = "Edit system prompt";
+"settings.systemPrompt.hint" = "Used by the cloud LLM when polishing your dictation. Local engine skips this step.";
"settings.about.title" = "About";
"settings.systemPrompt.reset" = "Reset";
"settings.asrLocale" = "ASR locale";
@@ -128,6 +141,47 @@
"settings.privacy.cloud.alert.message" = "Cloud polish sends transcribed text to the third-party API you configure. OSGKeyboard never stores data on our servers. Continue?";
"settings.link.support" = "Help & Feedback";
"settings.link.github" = "GitHub";
+"settings.support.footer" = "OSGKeyboard is open source. Report bugs and share ideas on GitHub.";
+"settings.support.intro" = "The fastest way to get help or send feedback is to open a GitHub issue. We read every report.";
+"settings.support.issuesHint" = "Search existing issues first. For bugs, include your iOS version, engine mode (local or cloud), and steps to reproduce.";
+"settings.support.openIssues" = "Open GitHub Issues";
+"settings.support.unavailable.title" = "Cannot load page";
+"settings.support.unavailable.message" = "GitHub Issues could not be opened. Check your network connection and try again.";
+"settings.link.licenses" = "Third-Party Licenses";
+"settings.link.models" = "On-device models";
+"settings.models.link.subtitle" = "Download source, models, and storage";
+"settings.licenses.title" = "Third-Party Licenses";
+"settings.licenses.footer" = "Open-source components used by OSGKeyboard. License texts are reproduced from upstream repositories; model weights are downloaded at runtime and cached on device. OSGKeyboard itself is source-available — commercial licensing: rocky.hk@gmail.com.";
+
+/* On-device models manager (settings → "On-device models") */
+"settings.models.title" = "On-device models";
+"settings.models.intro.title" = "Run speech recognition entirely on your iPhone";
+"settings.models.intro.body" = "Download the Qwen3-ASR model before using local recognition. Files are cached in this app's storage and can be deleted anytime.";
+"settings.models.mirror.title" = "Download source";
+"settings.models.source.modelScope" = "ModelScope";
+"settings.models.source.huggingface" = "Hugging Face";
+"settings.models.action.downloadSize %@" = "Download - %@";
+"settings.models.action.deleteSize %@" = "Delete - %@";
+"settings.models.action.downloadingPercent %lld" = "%lld%%";
+"settings.models.status.notDownloaded" = "Not downloaded";
+"settings.models.status.downloading" = "Downloading…";
+"settings.models.status.downloadingPercent %lld" = "Downloading… %lld%%";
+"settings.models.status.downloaded" = "Downloaded";
+"settings.models.status.failed" = "Download failed";
+"settings.models.action.download" = "Download";
+"settings.models.action.cancel" = "Cancel";
+"settings.models.action.cancelDownload" = "Cancel download";
+"settings.models.action.delete" = "Delete";
+
+/* Local ASR backend picker (local engine mode) */
+"asr.backend.speechAnalyzer.label" = "iOS SpeechAnalyzer";
+"asr.backend.speechAnalyzer.blurb" = "Built into iOS 26, always on-device, lowest latency.";
+"asr.backend.speechAnalyzer.hint" = "No download required.";
+"asr.backend.qwen3.label" = "Qwen3-ASR 0.6B (CoreML)";
+"asr.backend.qwen3.blurb" = "Stronger on Chinese dialects and noisy audio. ~1.6 GB CoreML model, works in background.";
+"asr.backend.qwen3.downloadHint" = "Download the CoreML model in the section above before dictating.";
+"asr.error.modelNotDownloaded" = "Qwen3-ASR is not downloaded yet. Download it from On-device models in Settings.";
+"asr.error.unsupportedOS" = "Qwen3-ASR CoreML requires iOS 18 or later. Use iOS SpeechAnalyzer or switch to a cloud engine.";
/* App group error */
"appGroup.error.title" = "App Group not configured";
@@ -143,6 +197,22 @@
"preview.placeholder" = "Type here or tap the record button";
"preview.clear" = "Clear text";
"preview.openSettingsA11y" = "Open OSGKeyboard settings";
+"preview.openSettings" = "Settings";
+"preview.record.start" = "Start Recording";
+"preview.record.stop" = "Stop Recording";
+"preview.status.field" = "Text field";
+"preview.status.recording" = "Recording…";
+"preview.status.processing" = "Transcribing…";
+
+/* Dictation handoff */
+"dictation.title" = "Voice Input";
+"dictation.status.ready" = "Ready to record…";
+"dictation.status.listening" = "Listening…";
+"dictation.status.listeningLive" = "Live transcription…";
+"dictation.status.processing" = "Processing…";
+"dictation.status.requestingPermission" = "Requesting permission…";
+"dictation.status.waitingResult" = "Waiting for transcription…";
+"dictation.status.saving" = "Saving…";
"preview.modeChip.cycle" = "Cycle input mode";
"preview.localeChip.cycle" = "Cycle recognition language";
"preview.error.audioSession" = "Audio session error: %@";
@@ -157,6 +227,7 @@
"keyboard.placeholder.processing" = "Processing";
"keyboard.placeholder.error" = "Polishing failed";
"keyboard.placeholder.localBadge" = "On-device";
+"keyboard.placeholder.cloudBadge" = "Cloud";
"keyboard.rec" = "REC";
"keyboard.space" = "Space";
"keyboard.denied.mic" = "Mic denied";
@@ -206,6 +277,14 @@
"home.flow.endShort" = "End";
"home.preview.label" = "Try typing";
"home.preview.placeholder" = "Tap to type and test…";
+"home.engine.unsupportedOS" = "Qwen3-ASR CoreML requires iOS 18+";
+"home.engine.warming" = "Loading ASR model into memory…";
+"home.engine.downloading" = "Downloading ASR model…";
+"home.engine.loadFailed" = "Model downloaded but failed to load. Restart the app and try again.";
+"home.engine.modelsReady" = "ASR model ready";
+"home.engine.setupRequired" = "Local ASR model not ready";
+"home.engine.downloadFirst" = "Download the local ASR model in Settings first";
+"home.engine.attentionBanner" = "Local ASR model is not ready yet. Download it and wait for warm-up before dictating.";
/* Tabs */
"tab.keyboard" = "Keyboard";
@@ -224,3 +303,11 @@
"flow.error.micRequired" = "Microphone access is required for background voice sessions.";
"flow.error.micUnavailable" = "Microphone is unavailable on this device.";
"flow.error.audioUnavailable" = "Could not start background audio.";
+
+"settings.models.error.metadata" = "Download cache was corrupted or incomplete. Try switching the download source and download again.";
+"settings.models.error.offline" = "Could not reach the model server. Check your network or try a different download source.";
+"settings.models.confirm.downloadTitle %@" = "Download %@";
+"settings.models.confirm.downloadingTitleFormat %@" = "Downloading %@";
+"settings.models.confirm.body %lld" = "This download uses about %lld MB on disk. We recommend Wi‑Fi to avoid cellular data. You can delete the file anytime from On-device models to free space.";
+"settings.models.confirm.cancel" = "Cancel";
+"settings.models.confirm.download %lld" = "Download %lld MB";
diff --git a/OSGKeyboard/zh-Hans.lproj/Localizable.strings b/OSGKeyboard/zh-Hans.lproj/Localizable.strings
index 6976453..58b9a2c 100644
--- a/OSGKeyboard/zh-Hans.lproj/Localizable.strings
+++ b/OSGKeyboard/zh-Hans.lproj/Localizable.strings
@@ -29,6 +29,7 @@
"onboarding.enable.step3.suffix" = ",选中 OSGKeyboard";
"onboarding.enable.openSettings" = "去设置";
"onboarding.api.title" = "选择语音转文字 AI 引擎";
+"onboarding.api.localModels.hint" = "使用 Qwen3-ASR 时需先下载下方 CoreML 语音识别模型,完成后才能开始使用。";
"settings.onboarding.replay" = "重新开始权限引导";
/* Common navigation */
@@ -83,15 +84,19 @@
/* Settings */
"settings.title" = "设置";
+"settings.appLanguage.title" = "界面语言";
+"settings.appLanguage.auto" = "自动";
+"settings.appLanguage.english" = "英文";
+"settings.appLanguage.chinese" = "中文";
"settings.reset.title" = "重置所有设置?";
"settings.reset.message" = "将清空 API Key、模型、接口地址、语音历史及云端确认状态。";
"settings.reset.confirm" = "重置所有设置";
-"settings.engine.title" = "引擎";
-"settings.engine.subtitle" = "本地不用 Key,只转文字;云端会润色,要配 Key。";
+"settings.engine.title" = "识别方式";
+"settings.engine.subtitle" = "本地:端侧识别,仅转录;云端:通过 API 润色。";
"settings.engine.local.title" = "本地识别";
"settings.engine.local.ios26" = "全程在手机本地,不用联网";
-"settings.engine.local.legacy" = "端侧 ASR,仅转录,无润色。";
-"settings.engine.cloud.title" = "云端润色";
+"settings.engine.local.legacy" = "端侧 ASR,仅转录,无润色。";
+"settings.engine.cloud.title" = "云端识别与润色";
"settings.engine.cloud.subtitle" = "本地转写 + 你配置的 API 润色,文字发往该第三方服务";
"settings.provider.title" = "提供商";
"settings.provider.subtitle" = "选择 LLM 提供商。";
@@ -103,6 +108,12 @@
"provider.custom" = "自定义";
"settings.api.title" = "接口";
"settings.language.title" = "语言";
+"settings.languageModels.title" = "语言与模型";
+"settings.localModels.title" = "本地模型";
+"settings.localModels.speechRole" = "语音识别";
+"settings.localModels.builtIn" = "内置";
+"settings.localModels.allReady" = "已就绪";
+"settings.localModels.readiness %lld %lld" = "%lld/%lld 已就绪";
"settings.language.subtitle.cloud" = "选择识别语言和文字处理模式。";
"settings.language.subtitle.local" = "选择识别语言。";
"settings.mode.title" = "模式";
@@ -110,6 +121,8 @@
"settings.mode.transcribe" = "转写";
"settings.mode.polish" = "润色";
"settings.systemPrompt.title" = "系统提示";
+"settings.systemPrompt.edit" = "编辑系统提示";
+"settings.systemPrompt.hint" = "云端润色时使用。本地识别模式不会用到此提示词。";
"settings.about.title" = "关于";
"settings.systemPrompt.reset" = "重置";
"settings.asrLocale" = "识别语言";
@@ -128,6 +141,47 @@
"settings.privacy.cloud.alert.message" = "云端润色会把转写文字发到你配置的第三方 API。OSGKeyboard 不会在自有服务器上存储数据。是否继续?";
"settings.link.support" = "帮助与反馈";
"settings.link.github" = "GitHub";
+"settings.support.footer" = "OSGKeyboard 为开源项目,欢迎在 GitHub 提交问题与建议。";
+"settings.support.intro" = "获取帮助或反馈问题的最快方式是提交 GitHub Issue,我们会阅读每一条反馈。";
+"settings.support.issuesHint" = "提交前请先搜索是否已有相同问题。反馈 Bug 时请附上 iOS 版本、引擎模式(本地/云端)以及复现步骤。";
+"settings.support.openIssues" = "打开 GitHub Issues";
+"settings.support.unavailable.title" = "无法加载页面";
+"settings.support.unavailable.message" = "无法打开 GitHub Issues,请检查网络连接后重试。";
+"settings.link.licenses" = "第三方许可";
+"settings.link.models" = "本地模型管理";
+"settings.models.link.subtitle" = "下载源、模型文件与存储管理";
+"settings.licenses.title" = "第三方许可";
+"settings.licenses.footer" = "以下为 OSGKeyboard 使用的开源组件。许可正文摘自上游仓库;模型权重在运行时下载并缓存在本机。OSGKeyboard 本身采用源码可见许可,商业授权请联系 rocky.hk@gmail.com。";
+
+/* 本地模型管理(设置 → "本地模型管理") */
+"settings.models.title" = "本地模型管理";
+"settings.models.intro.title" = "完全在 iPhone 上跑语音识别";
+"settings.models.intro.body" = "使用本地 Qwen3-ASR 前请先点「下载」。文件缓存在本 App 存储,可随时删除以释放空间。";
+"settings.models.mirror.title" = "下载源";
+"settings.models.source.modelScope" = "ModelScope";
+"settings.models.source.huggingface" = "Hugging Face";
+"settings.models.action.downloadSize %@" = "下载 - %@";
+"settings.models.action.deleteSize %@" = "删除 - %@";
+"settings.models.action.downloadingPercent %lld" = "%lld%%";
+"settings.models.status.notDownloaded" = "未下载";
+"settings.models.status.downloading" = "下载中…";
+"settings.models.status.downloadingPercent %lld" = "下载中… %lld%%";
+"settings.models.status.downloaded" = "已下载";
+"settings.models.status.failed" = "下载失败";
+"settings.models.action.download" = "下载";
+"settings.models.action.cancel" = "取消";
+"settings.models.action.cancelDownload" = "取消下载";
+"settings.models.action.delete" = "删除";
+
+/* 本地识别引擎选择(本地识别模式) */
+"asr.backend.speechAnalyzer.label" = "iOS SpeechAnalyzer";
+"asr.backend.speechAnalyzer.blurb" = "iOS 26 内置,全程在本地,延迟最低。";
+"asr.backend.speechAnalyzer.hint" = "无需下载,直接可用。";
+"asr.backend.qwen3.label" = "Qwen3-ASR 0.6B (CoreML)";
+"asr.backend.qwen3.blurb" = "对方言和嘈杂环境更友好,需下载约 1.6 GB CoreML 模型,支持后台转写。";
+"asr.backend.qwen3.downloadHint" = "请先在上方「本地模型管理」中下载 CoreML 模型,再开始录音。";
+"asr.error.modelNotDownloaded" = "Qwen3-ASR 尚未下载,请先在设置中的「本地模型管理」下载。";
+"asr.error.unsupportedOS" = "Qwen3-ASR CoreML 需要 iOS 18 或更高版本。请改用 iOS SpeechAnalyzer 或云端引擎。";
/* App group error */
"appGroup.error.title" = "App Group 未配置";
@@ -142,6 +196,22 @@
"preview.placeholder" = "试着输入或点击按钮录音";
"preview.clear" = "清空";
"preview.openSettingsA11y" = "打开 OSGKeyboard 设置";
+"preview.openSettings" = "设置";
+"preview.record.start" = "开始录音";
+"preview.record.stop" = "停止录音";
+"preview.status.field" = "输入框";
+"preview.status.recording" = "正在录音…";
+"preview.status.processing" = "识别中…";
+
+/* 语音输入(键盘跳转主 App) */
+"dictation.title" = "语音输入";
+"dictation.status.ready" = "准备录音…";
+"dictation.status.listening" = "正在听…";
+"dictation.status.listeningLive" = "正在实时识别…";
+"dictation.status.processing" = "处理中…";
+"dictation.status.requestingPermission" = "请求权限中…";
+"dictation.status.waitingResult" = "等待识别结果…";
+"dictation.status.saving" = "保存中…";
"preview.modeChip.cycle" = "切换输入模式";
"preview.localeChip.cycle" = "切换识别语言";
"preview.error.audioSession" = "音频会话错误: %@";
@@ -156,6 +226,7 @@
"keyboard.placeholder.processing" = "处理中…";
"keyboard.placeholder.error" = "润色失败";
"keyboard.placeholder.localBadge" = "本地";
+"keyboard.placeholder.cloudBadge" = "云端";
"keyboard.rec" = "REC";
"keyboard.space" = "空格";
"keyboard.denied.mic" = "麦克风被拒绝";
@@ -205,6 +276,14 @@
"home.flow.endShort" = "结束";
"home.preview.label" = "输入测试";
"home.preview.placeholder" = "点这里试试键盘";
+"home.engine.unsupportedOS" = "Qwen3-ASR CoreML 需要 iOS 18 或更高版本";
+"home.engine.warming" = "正在加载语音识别模型…";
+"home.engine.downloading" = "正在下载语音识别模型…";
+"home.engine.loadFailed" = "模型已下载,但加载失败。请重启应用后再试。";
+"home.engine.modelsReady" = "语音识别模型已就绪";
+"home.engine.setupRequired" = "本地语音识别模型未就绪";
+"home.engine.downloadFirst" = "请先在设置中下载本地语音识别模型";
+"home.engine.attentionBanner" = "本地语音识别模型尚未就绪,请先下载并等待加载完成后再使用语音输入。";
/* Tabs */
"tab.keyboard" = "键盘";
@@ -223,3 +302,11 @@
"flow.error.micRequired" = "需要麦克风权限才能维持后台语音会话。";
"flow.error.micUnavailable" = "当前设备无法使用麦克风。";
"flow.error.audioUnavailable" = "无法启动后台音频。";
+
+"settings.models.error.metadata" = "下载缓存损坏或不完整。请更换下载源后重试。";
+"settings.models.error.offline" = "无法连接模型服务器,请检查网络或更换下载源。";
+"settings.models.confirm.downloadTitle %@" = "下载 %@";
+"settings.models.confirm.downloadingTitleFormat %@" = "正在下载 %@";
+"settings.models.confirm.body %lld" = "下载后约占用 %lld MB 空间。建议在 Wi-Fi 下下载,避免消耗蜂窝流量。随时可以在本设置页删除模型文件以释放空间。";
+"settings.models.confirm.cancel" = "取消";
+"settings.models.confirm.download %lld" = "下载 %lld MB";
diff --git a/OSGKeyboardExt/KeyboardViewController.swift b/OSGKeyboardExt/KeyboardViewController.swift
index ccabe9b..9c50316 100644
--- a/OSGKeyboardExt/KeyboardViewController.swift
+++ b/OSGKeyboardExt/KeyboardViewController.swift
@@ -28,7 +28,16 @@ public final class KeyboardViewController: UIInputViewController {
static let pollIntervalNs: UInt64 = 200_000_000
/// Give the user time to manually open the host app when auto-jump fails.
static let startTimeout: TimeInterval = 30
- static let resultTimeout: TimeInterval = 45
+
+ static func resultTimeout(
+ engineMode: String,
+ localASRBackend: LocalASRBackend
+ ) -> TimeInterval {
+ FlowSessionKeys.keyboardResultTimeout(
+ engineMode: engineMode,
+ localASRBackend: localASRBackend
+ )
+ }
}
private enum DictationWatchdog {
@@ -128,6 +137,7 @@ public final class KeyboardViewController: UIInputViewController {
state.setMode = { [weak self] m in self?.persistMode(m) }
state.setLocale = { [weak self] l in self?.persistLocale(l) }
state.setEngineMode = { [weak self] m in self?.persistEngineMode(m) }
+ state.setLocalASRBackend = { [weak self] b in self?.persistLocalASRBackend(b) }
state.insertNewline = { [weak self] in self?.textDocumentProxy.insertText("\n") }
state.insertSpace = { [weak self] in self?.textDocumentProxy.insertText(" ") }
state.deleteBackward = { [weak self] in self?.textDocumentProxy.deleteBackward() }
@@ -191,6 +201,9 @@ public final class KeyboardViewController: UIInputViewController {
}
private func refreshFlowSessionState() {
+ persistor.refreshRuntimeFlags(into: state)
+ consumePendingFlowDeliveryIfNeeded()
+
let active = FlowSessionBridge.isSessionActive()
state.flowSessionActive = active
@@ -209,11 +222,33 @@ public final class KeyboardViewController: UIInputViewController {
}
}
+ /// Pick up transcripts/errors the host wrote while the extension was paused.
+ private func consumePendingFlowDeliveryIfNeeded() {
+ if isAwaitingFlowResult {
+ if let delivery = FlowSessionBridge.consumeTranscriptionDelivery() {
+ isAwaitingFlowResult = false
+ stopFlowWatchdog()
+ handleFlowTranscript(delivery)
+ return
+ }
+ if let error = FlowSessionBridge.consumeTranscriptionError() {
+ isAwaitingFlowResult = false
+ stopFlowWatchdog()
+ state.phase = .error(.unknown(error), message: error)
+ scheduleAutoClearError()
+ return
+ }
+ }
+
+ if isPendingFlowStart, FlowSessionBridge.isSessionActive() {
+ completeFlowStartHandoff()
+ }
+ }
+
/// When the host session is down, proactively jump to the app to start it.
private func maybeAutoStartFlowSession() {
guard !FlowSessionBridge.isSessionActive() else { return }
guard !isPendingFlowStart, !isFlowRecording, !isAwaitingFlowResult else { return }
- guard state.mode != .off else { return }
guard hasFullAccess, AppGroup.isAvailable else { return }
guard case .idle = state.phase else { return }
@@ -249,7 +284,6 @@ public final class KeyboardViewController: UIInputViewController {
default:
return
}
- guard state.mode != .off else { return }
guard hasFullAccess else {
let msg = ExtL10n.string("keyboard.error.fullAccessRequired")
state.phase = .error(.unknown(msg), message: msg)
@@ -392,12 +426,16 @@ public final class KeyboardViewController: UIInputViewController {
stopFlowWatchdog()
isAwaitingFlowResult = true
let startedAt = Date().timeIntervalSince1970
+ let resultTimeout = FlowWatchdog.resultTimeout(
+ engineMode: state.engineMode,
+ localASRBackend: state.localASRBackend
+ )
flowWatchdogTask = Task { @MainActor [weak self] in
while let self, !Task.isCancelled {
- if let result = FlowSessionBridge.consumeTranscriptionResult() {
+ if let delivery = FlowSessionBridge.consumeTranscriptionDelivery() {
self.isAwaitingFlowResult = false
self.stopFlowWatchdog()
- self.handleFlowTranscript(result)
+ self.handleFlowTranscript(delivery)
return
}
if let error = FlowSessionBridge.consumeTranscriptionError() {
@@ -408,7 +446,7 @@ public final class KeyboardViewController: UIInputViewController {
return
}
let now = Date().timeIntervalSince1970
- if now - startedAt > FlowWatchdog.resultTimeout {
+ if now - startedAt > resultTimeout {
self.isAwaitingFlowResult = false
self.stopFlowWatchdog()
let msg = ExtL10n.string("keyboard.flow.resultTimeout")
@@ -421,8 +459,8 @@ public final class KeyboardViewController: UIInputViewController {
}
}
- private func handleFlowTranscript(_ transcript: String) {
- let trimmed = transcript.trimmingCharacters(in: .whitespacesAndNewlines)
+ private func handleFlowTranscript(_ delivery: TranscriptionDelivery) {
+ let trimmed = delivery.text.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty else {
state.phase = .idle
state.level = 0
@@ -432,7 +470,12 @@ public final class KeyboardViewController: UIInputViewController {
textDocumentProxy.insertText(trimmed)
state.lastTranscript = ""
state.level = 0
- state.phase = .idle
+ if let warning = delivery.polishWarning {
+ state.phase = .error(.unknown(warning), message: warning)
+ scheduleAutoClearError()
+ } else {
+ state.phase = .idle
+ }
debug("flow insert length=\(trimmed.count)")
}
@@ -457,8 +500,8 @@ public final class KeyboardViewController: UIInputViewController {
}
}
- private func handleFinalTranscript(_ transcript: String) {
- let trimmed = transcript.trimmingCharacters(in: .whitespacesAndNewlines)
+ private func handleFinalTranscript(_ delivery: TranscriptionDelivery) {
+ let trimmed = delivery.text.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty else {
debug("received empty transcript")
awaitingDictationResult = false
@@ -469,14 +512,19 @@ public final class KeyboardViewController: UIInputViewController {
debug("received transcript length=\(trimmed.count)")
awaitingDictationResult = false
stopDictationWatchdog()
- // Local engine or transcribe mode: insert directly, no LLM call.
- if state.isLocalEngine || state.mode == .transcribe {
+ // Local engine: host app delivers raw ASR transcript; insert as-is.
+ if state.isLocalEngine {
textDocumentProxy.insertText(trimmed)
state.lastTranscript = ""
- state.phase = .idle
+ if let warning = delivery.polishWarning {
+ state.phase = .error(.unknown(warning), message: warning)
+ scheduleAutoClearError()
+ } else {
+ state.phase = .idle
+ }
return
}
- // `.polish` (default): call the LLM.
+ // Cloud engine: always polish via the configured LLM.
state.phase = .processing
Task { @MainActor [weak self] in
guard let self else { return }
@@ -555,6 +603,11 @@ public final class KeyboardViewController: UIInputViewController {
persistor.persist(engineMode: mode)
}
+ private func persistLocalASRBackend(_ backend: LocalASRBackend) {
+ state.localASRBackend = backend
+ persistor.persist(localASRBackend: backend)
+ }
+
// MARK: - Open host app
private func openHostApp(path: String = "settings") {
@@ -592,9 +645,9 @@ public final class KeyboardViewController: UIInputViewController {
}
private func consumePendingDictationResultIfNeeded() {
- guard let transcript = DictationBridge.consumePendingTranscript() else { return }
+ guard let delivery = DictationBridge.consumePendingDelivery() else { return }
debug("consumePendingDictationResultIfNeeded success")
- handleFinalTranscript(transcript)
+ handleFinalTranscript(delivery)
}
private func refreshDictationProgressStateIfNeeded() {
diff --git a/OSGKeyboardExt/Services/AppGroupPersistor.swift b/OSGKeyboardExt/Services/AppGroupPersistor.swift
index 59ef98f..c47df1a 100644
--- a/OSGKeyboardExt/Services/AppGroupPersistor.swift
+++ b/OSGKeyboardExt/Services/AppGroupPersistor.swift
@@ -31,8 +31,14 @@ public struct AppGroupPersistor {
}
let store = AppGroupStore()
state.localeId = store.localeId
- state.mode = KeyboardViewController.State.InputMode(rawValue: store.modeId) ?? .polish
+ // Both engines always polish; ignore legacy off/transcribe modeId.
+ state.mode = .polish
state.engineMode = store.engineMode
+ state.localASRBackend = store.localASRBackend
+ state.localModelsReady = OnDeviceModelStatus.isLocalStackReady(
+ asrBackend: store.localASRBackend
+ )
+ state.localModelsLoaded = OnDeviceModelStatus.modelsLoadedInMemory()
#if DEBUG
// Print a masked view of the live App Group config so we can see
@@ -49,17 +55,31 @@ public struct AppGroupPersistor {
}
print("""
🔍 [AppGroupPersistor.load]
- providerId = \(store.providerId)
- baseURL = \(store.baseURL)
- apiKey = \(masked)
- model = \(store.model)
- modeId = \(store.modeId)
- localeId = \(store.localeId)
+ providerId = \(store.providerId)
+ baseURL = \(store.baseURL)
+ apiKey = \(masked)
+ model = \(store.model)
+ modeId = \(store.modeId)
+ localeId = \(store.localeId)
+ localASRBackend = \(store.localASRBackend.rawValue)
""")
#endif
return .loaded
}
+ /// Lightweight refresh for flags the host app may update while the
+ /// keyboard stays open (model downloads, engine switches).
+ public func refreshRuntimeFlags(into state: KeyboardViewController.State) {
+ guard AppGroup.isAvailable else { return }
+ let store = AppGroupStore()
+ state.engineMode = store.engineMode
+ state.localASRBackend = store.localASRBackend
+ state.localModelsReady = OnDeviceModelStatus.isLocalStackReady(
+ asrBackend: store.localASRBackend
+ )
+ state.localModelsLoaded = OnDeviceModelStatus.modelsLoadedInMemory()
+ }
+
/// Persist `mode` to the App Group store.
public func persist(mode: KeyboardViewController.State.InputMode) {
guard AppGroup.isAvailable else { return }
@@ -77,4 +97,10 @@ public struct AppGroupPersistor {
guard AppGroup.isAvailable else { return }
AppGroupStore().setEngineMode(engineMode)
}
+
+ /// Persist `localASRBackend` to the App Group store.
+ public func persist(localASRBackend: LocalASRBackend) {
+ guard AppGroup.isAvailable else { return }
+ AppGroupStore().setLocalASRBackend(localASRBackend)
+ }
}
\ No newline at end of file
diff --git a/OSGKeyboardExt/Utilities/ExtL10n.swift b/OSGKeyboardExt/Utilities/ExtL10n.swift
index 61cb664..ad84518 100644
--- a/OSGKeyboardExt/Utilities/ExtL10n.swift
+++ b/OSGKeyboardExt/Utilities/ExtL10n.swift
@@ -7,20 +7,27 @@
import Foundation
import SwiftUI
+import OSGKeyboardShared
enum ExtL10n {
private static let table = "Keyboard"
- private static let bundle = Bundle(for: KeyboardViewController.self)
+ private static let container = Bundle(for: KeyboardViewController.self)
+
+ private static var bundle: Bundle {
+ AppUILanguage.localizedBundle(
+ in: container,
+ language: AppGroupStore().uiLanguage
+ )
+ }
static func string(_ key: String) -> String {
- let value = NSLocalizedString(
+ NSLocalizedString(
key,
tableName: table,
bundle: bundle,
value: key,
comment: ""
)
- return value
}
static func text(_ key: String) -> Text {
diff --git a/OSGKeyboardExt/Views/KeyboardRootView.swift b/OSGKeyboardExt/Views/KeyboardRootView.swift
index 6aebaa8..83858f5 100644
--- a/OSGKeyboardExt/Views/KeyboardRootView.swift
+++ b/OSGKeyboardExt/Views/KeyboardRootView.swift
@@ -69,12 +69,9 @@ public struct KeyboardRootView: View {
private var topBar: some View {
HStack(spacing: Spacing.xs) {
if state.isLocalEngine {
- // Local engine: always transcribe, no mode menu needed.
LocalEngineChip()
} else {
- ModeChip(mode: state.mode) { newMode in
- state.setMode(newMode)
- }
+ CloudEngineChip()
}
LocaleChip(localeId: state.localeId) { newId in
state.setLocale(newId)
@@ -103,6 +100,9 @@ public struct KeyboardRootView: View {
phase: state.phase,
transcript: state.lastTranscript,
flowSessionActive: state.flowSessionActive,
+ isLocalEngine: state.isLocalEngine,
+ localModelsReady: state.localModelsReady,
+ localModelsLoaded: state.localModelsLoaded,
openSettings: state.openSettings,
startFlowSession: state.startFlowSession
)
@@ -200,6 +200,9 @@ private struct TranscriptLine: View {
let phase: KeyboardViewController.State.Phase
let transcript: String
let flowSessionActive: Bool
+ let isLocalEngine: Bool
+ let localModelsReady: Bool
+ let localModelsLoaded: Bool
let openSettings: () -> Void
let startFlowSession: () -> Void
@@ -207,7 +210,30 @@ private struct TranscriptLine: View {
ZStack {
switch phase {
case .idle:
- if flowSessionActive {
+ if isLocalEngine, !localModelsReady {
+ Button(action: openSettings) {
+ HStack(spacing: 4) {
+ Text(ExtL10n.string("keyboard.models.notDownloaded"))
+ Image(systemName: "chevron.right")
+ .font(.system(size: 10, weight: .semibold))
+ }
+ .font(TypeStyle.caption)
+ .foregroundStyle(palette.warning)
+ .lineLimit(1)
+ .truncationMode(.tail)
+ .frame(maxWidth: .infinity)
+ .contentShape(Rectangle())
+ }
+ .buttonStyle(.plain)
+ .accessibilityHint(ExtL10n.text("keyboard.models.downloadHint"))
+ } else if isLocalEngine, localModelsReady, !localModelsLoaded {
+ HStack(spacing: 6) {
+ ProgressView().controlSize(.mini).tint(palette.textSecondary)
+ ExtL10n.text("keyboard.models.warming")
+ .font(TypeStyle.caption)
+ .foregroundStyle(palette.textSecondary)
+ }
+ } else if flowSessionActive {
ExtL10n.text("keyboard.placeholder.idle")
.font(TypeStyle.caption)
.foregroundStyle(palette.textTertiary)
@@ -393,6 +419,26 @@ private struct StatusBadge: View {
}
}
+// MARK: - Cloud engine chip (cloud always ASR + LLM polish)
+
+private struct CloudEngineChip: View {
+ @Environment(\.themePalette) private var palette: ThemePalette
+
+ var body: some View {
+ HStack(spacing: 4) {
+ Image(systemName: "wand.and.stars")
+ ExtL10n.text("keyboard.placeholder.cloudBadge")
+ }
+ .font(TypeStyle.caption2)
+ .foregroundStyle(palette.accent)
+ .padding(.horizontal, Spacing.xs + 2)
+ .padding(.vertical, 5)
+ .frame(minHeight: 26)
+ .background(palette.accent.opacity(0.15), in: Capsule())
+ .overlay(Capsule().stroke(palette.accent.opacity(0.35), lineWidth: 0.5))
+ }
+}
+
// MARK: - Local engine chip (shown instead of ModeChip when engineMode == "local")
private struct LocalEngineChip: View {
@@ -413,58 +459,6 @@ private struct LocalEngineChip: View {
}
}
-// MARK: - Mode chip
-
-private struct ModeChip: View {
- @Environment(\.themePalette) private var palette: ThemePalette
-
- let mode: KeyboardViewController.State.InputMode
- let onChange: (KeyboardViewController.State.InputMode) -> Void
-
- var body: some View {
- Menu {
- ForEach(KeyboardViewController.State.InputMode.allCases) { m in
- Button {
- onChange(m)
- } label: {
- if m == mode {
- Label(label(for: m), systemImage: "checkmark")
- } else {
- Text(label(for: m))
- }
- }
- }
- } label: {
- HStack(spacing: 4) {
- Image(systemName: icon(for: mode))
- Text(label(for: mode))
- Image(systemName: "chevron.down")
- .font(.system(size: 8, weight: .bold))
- }
- .font(TypeStyle.caption2)
- .foregroundStyle(mode == .off ? palette.textTertiary : palette.textPrimary)
- .padding(.horizontal, Spacing.xs + 2)
- .padding(.vertical, 5)
- .frame(minHeight: 26)
- .background(palette.surfaceElevated, in: Capsule())
- .overlay(Capsule().stroke(palette.divider, lineWidth: 0.5))
- }
- .menuStyle(.button)
- }
-
- private func label(for m: KeyboardViewController.State.InputMode) -> String {
- ExtL10n.string(m.labelKey)
- }
-
- private func icon(for m: KeyboardViewController.State.InputMode) -> String {
- switch m {
- case .off: return "mic.slash.fill"
- case .transcribe: return "text.bubble.fill"
- case .polish: return "wand.and.stars"
- }
- }
-}
-
// MARK: - Locale chip
private struct LocaleChip: View {
diff --git a/OSGKeyboardExt/en.lproj/Keyboard.strings b/OSGKeyboardExt/en.lproj/Keyboard.strings
index 0c4811c..1f34ba2 100644
--- a/OSGKeyboardExt/en.lproj/Keyboard.strings
+++ b/OSGKeyboardExt/en.lproj/Keyboard.strings
@@ -69,12 +69,12 @@
"settings.reset.title" = "Reset all settings?";
"settings.reset.message" = "API key, model, and base URL will be cleared.";
"settings.reset.confirm" = "Reset all settings";
-"settings.engine.title" = "Engine";
+"settings.engine.title" = "Recognition method";
"settings.engine.subtitle" = "Pick the recognition engine. Local engine does transcription only, no API key needed.";
-"settings.engine.local.title" = "On-device";
+"settings.engine.local.title" = "On-device recognition";
"settings.engine.local.ios26" = "Always on-device, no network.";
"settings.engine.local.legacy" = "On-device ASR. Transcription only, no polish.";
-"settings.engine.cloud.title" = "Cloud polish";
+"settings.engine.cloud.title" = "Cloud recognition & polish";
"settings.engine.cloud.subtitle" = "ASR + LLM polish. API key required.";
"settings.provider.title" = "Provider";
"settings.provider.subtitle" = "Pick the LLM that polishes your dictation.";
@@ -118,6 +118,10 @@
"keyboard.placeholder.processing" = "Processing";
"keyboard.placeholder.error" = "Polishing failed";
"keyboard.placeholder.localBadge" = "On-device";
+"keyboard.placeholder.cloudBadge" = "Cloud";
+"keyboard.models.notDownloaded" = "On-device models not downloaded";
+"keyboard.models.downloadHint" = "Open OSGKeyboard to download models";
+"keyboard.models.warming" = "Loading models…";
"keyboard.rec" = "REC";
"keyboard.space" = "Space";
"keyboard.denied.mic" = "Mic denied";
diff --git a/OSGKeyboardExt/zh-Hans.lproj/Keyboard.strings b/OSGKeyboardExt/zh-Hans.lproj/Keyboard.strings
index 746fbe2..c0cd6e9 100644
--- a/OSGKeyboardExt/zh-Hans.lproj/Keyboard.strings
+++ b/OSGKeyboardExt/zh-Hans.lproj/Keyboard.strings
@@ -69,12 +69,12 @@
"settings.reset.title" = "重置所有设置?";
"settings.reset.message" = "API key、model 和 base URL 都会被清空。";
"settings.reset.confirm" = "重置所有设置";
-"settings.engine.title" = "引擎";
+"settings.engine.title" = "识别方式";
"settings.engine.subtitle" = "选择识别方式。本地引擎无需 API Key,仅做语音转录。";
"settings.engine.local.title" = "本地识别";
"settings.engine.local.ios26" = "始终端侧,无需联网。";
"settings.engine.local.legacy" = "端侧 ASR,仅转录,无润色。";
-"settings.engine.cloud.title" = "云端润色";
+"settings.engine.cloud.title" = "云端识别与润色";
"settings.engine.cloud.subtitle" = "ASR 转录 + LLM 润色,需要 API Key。";
"settings.provider.title" = "提供商";
"settings.provider.subtitle" = "选择 LLM 提供商。";
@@ -118,6 +118,10 @@
"keyboard.placeholder.processing" = "处理中…";
"keyboard.placeholder.error" = "润色失败";
"keyboard.placeholder.localBadge" = "本地";
+"keyboard.placeholder.cloudBadge" = "云端";
+"keyboard.models.notDownloaded" = "本地模型尚未下载";
+"keyboard.models.downloadHint" = "打开 OSGKeyboard 下载模型";
+"keyboard.models.warming" = "正在加载模型…";
"keyboard.rec" = "REC";
"keyboard.space" = "空格";
"keyboard.denied.mic" = "麦克风被拒绝";
diff --git a/OSGKeyboardShared/Localization/SharedL10n.swift b/OSGKeyboardShared/Localization/SharedL10n.swift
new file mode 100644
index 0000000..2cf23e9
--- /dev/null
+++ b/OSGKeyboardShared/Localization/SharedL10n.swift
@@ -0,0 +1,41 @@
+// SharedL10n.swift
+// OSGKeyboard · Shared
+//
+// Localized strings shipped inside the shared framework (Shared.strings).
+// Respects the in-app UI language override from App Group settings.
+
+import Foundation
+
+public enum SharedL10n {
+ private static let table = "Shared"
+ private static let container = Bundle(for: SharedBundleToken.self)
+
+ public static func string(
+ _ key: String,
+ language: AppUILanguage? = nil
+ ) -> String {
+ let lang = language ?? AppGroupStore().uiLanguage
+ let bundle = AppUILanguage.localizedBundle(in: container, language: lang)
+ return NSLocalizedString(
+ key,
+ tableName: table,
+ bundle: bundle,
+ value: key,
+ comment: ""
+ )
+ }
+
+ public static func format(
+ _ key: String,
+ language: AppUILanguage? = nil,
+ _ args: CVarArg...
+ ) -> String {
+ String(
+ format: string(key, language: language),
+ locale: Locale.current,
+ arguments: args
+ )
+ }
+}
+
+private final class SharedBundleToken {}
diff --git a/OSGKeyboardShared/Models/AppUILanguage.swift b/OSGKeyboardShared/Models/AppUILanguage.swift
new file mode 100644
index 0000000..1dd456c
--- /dev/null
+++ b/OSGKeyboardShared/Models/AppUILanguage.swift
@@ -0,0 +1,87 @@
+// AppUILanguage.swift
+// OSGKeyboard · Shared
+//
+// In-app UI language override (main app + keyboard extension strings).
+// Distinct from `localeId`, which controls speech recognition language.
+
+import Foundation
+
+public enum AppUILanguage: String, CaseIterable, Identifiable, Sendable, Codable {
+ case auto
+ case english = "en"
+ case chinese = "zh-Hans"
+
+ public var id: String { rawValue }
+
+ public var labelKey: String {
+ switch self {
+ case .auto: return "settings.appLanguage.auto"
+ case .english: return "settings.appLanguage.english"
+ case .chinese: return "settings.appLanguage.chinese"
+ }
+ }
+
+ /// Locale for SwiftUI `.environment(\.locale, …)` in the host app.
+ public var swiftUILocale: Locale {
+ switch self {
+ case .auto:
+ return Locale.autoupdatingCurrent
+ case .english:
+ return Locale(identifier: "en")
+ case .chinese:
+ return Locale(identifier: "zh-Hans")
+ }
+ }
+
+ /// `.lproj` folder name used for manual bundle lookups (extension).
+ public func resolvedLanguageCode(
+ preferredLanguages: [String] = Locale.preferredLanguages
+ ) -> String {
+ switch self {
+ case .english:
+ return "en"
+ case .chinese:
+ return "zh-Hans"
+ case .auto:
+ if preferredLanguages.contains(where: { $0.hasPrefix("zh") }) {
+ return "zh-Hans"
+ }
+ return "en"
+ }
+ }
+
+ public static func fromStored(_ raw: String?) -> AppUILanguage {
+ guard let raw, let value = AppUILanguage(rawValue: raw) else { return .auto }
+ return value
+ }
+
+ /// Picks the best-matching `.lproj` inside `container` for this preference.
+ public static func localizedBundle(
+ in container: Bundle,
+ language: AppUILanguage,
+ preferredLanguages: [String] = Locale.preferredLanguages
+ ) -> Bundle {
+ let code = language.resolvedLanguageCode(preferredLanguages: preferredLanguages)
+ guard let path = container.path(forResource: code, ofType: "lproj"),
+ let bundle = Bundle(path: path) else {
+ return container
+ }
+ return bundle
+ }
+
+ public static func localizedString(
+ _ key: String,
+ tableName: String?,
+ bundle container: Bundle,
+ language: AppUILanguage = AppGroupStore().uiLanguage
+ ) -> String {
+ let bundle = localizedBundle(in: container, language: language)
+ return NSLocalizedString(
+ key,
+ tableName: tableName,
+ bundle: bundle,
+ value: key,
+ comment: ""
+ )
+ }
+}
diff --git a/OSGKeyboardShared/Models/EngineServiceLabel.swift b/OSGKeyboardShared/Models/EngineServiceLabel.swift
index 5840ec2..540a54a 100644
--- a/OSGKeyboardShared/Models/EngineServiceLabel.swift
+++ b/OSGKeyboardShared/Models/EngineServiceLabel.swift
@@ -9,18 +9,37 @@ public enum EngineServiceLabel {
public static func summary(
engineMode: String,
providerId: String,
- model: String
+ model: String,
+ localASRBackend: LocalASRBackend = .speechAnalyzer,
+ language: AppUILanguage? = nil
) -> String {
- let isChinese = Locale.preferredLanguages.first?.hasPrefix("zh") == true
- let prefix = isChinese ? "当前:" : "Active: "
+ let lang = language ?? AppGroupStore().uiLanguage
if engineMode == "local" {
- return isChinese
- ? "\(prefix)本地引擎 · Apple SpeechAnalyzer"
- : "\(prefix)On-device · Apple SpeechAnalyzer"
+ let asrName = asrDisplayName(for: localASRBackend, language: lang)
+ return SharedL10n.format("engine.summary.local", language: lang, asrName)
}
- let providerName = ProviderDisplayName.name(for: providerId)
+ let providerName = ProviderDisplayName.name(for: providerId, language: lang)
let trimmedModel = model.trimmingCharacters(in: .whitespacesAndNewlines)
- if trimmedModel.isEmpty { return "\(prefix)\(providerName)" }
- return "\(prefix)\(providerName) · \(trimmedModel)"
+ if trimmedModel.isEmpty {
+ return SharedL10n.format("engine.summary.cloud", language: lang, providerName)
+ }
+ return SharedL10n.format(
+ "engine.summary.cloudWithModel",
+ language: lang,
+ providerName,
+ trimmedModel
+ )
+ }
+
+ private static func asrDisplayName(
+ for backend: LocalASRBackend,
+ language: AppUILanguage
+ ) -> String {
+ switch backend {
+ case .speechAnalyzer:
+ return SharedL10n.string("engine.asr.appleSpeech", language: language)
+ case .qwen3ASR:
+ return SharedL10n.string("model.qwen3asr.name", language: language)
+ }
}
}
diff --git a/OSGKeyboardShared/Models/FlowUtteranceChunkConfig.swift b/OSGKeyboardShared/Models/FlowUtteranceChunkConfig.swift
new file mode 100644
index 0000000..10c8505
--- /dev/null
+++ b/OSGKeyboardShared/Models/FlowUtteranceChunkConfig.swift
@@ -0,0 +1,69 @@
+// FlowUtteranceChunkConfig.swift
+// OSGKeyboard · Shared
+//
+// Chunking policy for pipelined Flow utterance ASR (up to 3 minutes).
+
+import Foundation
+
+public struct FlowUtteranceChunkConfig: Sendable, Equatable {
+ /// Target maximum duration per ASR chunk.
+ public let maxChunkDurationSeconds: TimeInterval
+ /// Tail overlap fed into the next chunk for boundary dedup when stitching.
+ public let overlapDurationSeconds: TimeInterval
+ /// After hitting the max window, wait up to this long for a pause before hard-splitting.
+ public let pauseExtensionMaxSeconds: TimeInterval
+ /// RMS below this is treated as a pause candidate (Float32 mono @ 16 kHz).
+ public let pauseRMSThreshold: Float
+ public let sampleRate: Int
+
+ public init(
+ maxChunkDurationSeconds: TimeInterval,
+ overlapDurationSeconds: TimeInterval,
+ pauseExtensionMaxSeconds: TimeInterval,
+ pauseRMSThreshold: Float,
+ sampleRate: Int
+ ) {
+ self.maxChunkDurationSeconds = maxChunkDurationSeconds
+ self.overlapDurationSeconds = overlapDurationSeconds
+ self.pauseExtensionMaxSeconds = pauseExtensionMaxSeconds
+ self.pauseRMSThreshold = pauseRMSThreshold
+ self.sampleRate = sampleRate
+ }
+
+ public var maxChunkSamples: Int {
+ Int(maxChunkDurationSeconds * Double(sampleRate))
+ }
+
+ public var overlapSamples: Int {
+ Int(overlapDurationSeconds * Double(sampleRate))
+ }
+
+ public var pauseExtensionSamples: Int {
+ Int(pauseExtensionMaxSeconds * Double(sampleRate))
+ }
+
+ /// Default for keyboard Flow utterances (≤ 3 min, pipelined ASR).
+ public static let flowDefault = FlowUtteranceChunkConfig(
+ maxChunkDurationSeconds: 30,
+ overlapDurationSeconds: 0.5,
+ pauseExtensionMaxSeconds: 2,
+ pauseRMSThreshold: 0.015,
+ sampleRate: 16_000
+ )
+}
+
+public struct UtteranceAudioChunk: Sendable, Equatable {
+ public let index: Int
+ public let samples: [Float]
+ public let isLast: Bool
+
+ public init(index: Int, samples: [Float], isLast: Bool) {
+ self.index = index
+ self.samples = samples
+ self.isLast = isLast
+ }
+
+ public var durationSeconds: Double {
+ Double(samples.count) / 16_000.0
+ }
+}
diff --git a/OSGKeyboardShared/Models/LocalASRBackend.swift b/OSGKeyboardShared/Models/LocalASRBackend.swift
new file mode 100644
index 0000000..5889e27
--- /dev/null
+++ b/OSGKeyboardShared/Models/LocalASRBackend.swift
@@ -0,0 +1,56 @@
+// LocalASRBackend.swift
+// OSGKeyboard · Shared
+//
+// Identifies which on-device speech recognition engine to use when the
+// user picks the "local" engine (no cloud LLM polish). The shared
+// factory `ASRServiceFactory` dispatches on this enum; the settings UI
+// renders it as a picker.
+//
+// Why an enum in `Shared` rather than living next to the concrete
+// `ASRService` implementations: the value must be serialisable into
+// the App Group store (so the keyboard extension can observe the
+// selection), exposed via `ProviderConfig` (UI binding) and consumed
+// by every layer that asks for an ASR backend.
+
+import Foundation
+
+public enum LocalASRBackend: String, CaseIterable, Identifiable, Sendable, Codable {
+ /// iOS 26 `SpeechAnalyzer` + `DictationTranscriber`. Always
+ /// on-device, no asset download, ships with iOS. Default for every
+ /// fresh install — anything else is opt-in.
+ case speechAnalyzer
+
+ /// Qwen3-ASR-0.6B via CoreML (Neural Engine + CPU). Stronger on Chinese
+ /// dialects and noisy audio than `SpeechAnalyzer`, works in Flow while
+ /// the host app is backgrounded, but requires a ~1.6 GB download on first
+ /// use and iOS 18+.
+ case qwen3ASR
+
+ public var id: String { rawValue }
+
+ /// Localisation key for the human label in the settings picker.
+ public var labelKey: String {
+ switch self {
+ case .speechAnalyzer: return "asr.backend.speechAnalyzer.label"
+ case .qwen3ASR: return "asr.backend.qwen3.label"
+ }
+ }
+
+ /// Localisation key for the one-line subtitle shown under the label.
+ public var blurbKey: String {
+ switch self {
+ case .speechAnalyzer: return "asr.backend.speechAnalyzer.blurb"
+ case .qwen3ASR: return "asr.backend.qwen3.blurb"
+ }
+ }
+
+ /// Whether this backend needs the user to download a model file
+ /// before it can run. Used to gate the "Downloading Qwen3-ASR" UI
+ /// in a follow-up; for now we just expose the flag.
+ public var requiresModelDownload: Bool {
+ switch self {
+ case .speechAnalyzer: return false
+ case .qwen3ASR: return true
+ }
+ }
+}
diff --git a/OSGKeyboardShared/Models/OnDeviceModel.swift b/OSGKeyboardShared/Models/OnDeviceModel.swift
new file mode 100644
index 0000000..1734456
--- /dev/null
+++ b/OSGKeyboardShared/Models/OnDeviceModel.swift
@@ -0,0 +1,54 @@
+// OnDeviceModel.swift
+// OSGKeyboard · Shared
+//
+// Identity of an on-device model the host app downloads and the
+// keyboard extension observes via App Group flags (the extension
+// cannot read the main app's Caches directory).
+
+import Foundation
+
+public enum OnDeviceModel: String, CaseIterable, Identifiable, Sendable {
+ case qwen3ASR
+
+ public var id: String { rawValue }
+
+ /// CoreML inference bundle (`aufklarer/Qwen3-ASR-CoreML`), derived from
+ /// official `Qwen/Qwen3-ASR-0.6B`.
+ public var repoId: String {
+ switch self {
+ case .qwen3ASR: return "aufklarer/Qwen3-ASR-CoreML"
+ }
+ }
+
+ /// Tokenizer files (vocab / merges) pulled from the upstream Qwen repo.
+ public var tokenizerRepoId: String {
+ switch self {
+ case .qwen3ASR: return "Qwen/Qwen3-ASR-0.6B"
+ }
+ }
+
+ public var displayName: String {
+ switch self {
+ case .qwen3ASR: return "Qwen3-ASR 0.6B (CoreML)"
+ }
+ }
+
+ public var approximateSizeMB: Int {
+ switch self {
+ case .qwen3ASR: return 1_600
+ }
+ }
+
+ public var compactSizeLabel: String {
+ "\(approximateSizeMB)M"
+ }
+
+ /// Settings list title: model name plus compact size.
+ public var listTitle: String {
+ "\(displayName) · \(compactSizeLabel)"
+ }
+
+ public var repoAndSizeLabel: String {
+ "\(repoId) · \(approximateSizeMB) MB"
+ }
+}
diff --git a/OSGKeyboardShared/Models/ProviderConfig.swift b/OSGKeyboardShared/Models/ProviderConfig.swift
index 60886d0..9530f43 100644
--- a/OSGKeyboardShared/Models/ProviderConfig.swift
+++ b/OSGKeyboardShared/Models/ProviderConfig.swift
@@ -30,6 +30,11 @@ public final class ProviderConfig: ObservableObject, @unchecked Sendable {
static let hasCompletedOnboarding = "config.hasCompletedOnboarding"
static let onboardingPage = "config.onboardingPage"
static let hasAcknowledgedCloudSharing = "config.hasAcknowledgedCloudSharing"
+ // Which on-device ASR engine to use when `engineMode == "local"`.
+ // Persisted in the App Group so the keyboard can read the
+ // selection even though it never instantiates the backend itself.
+ static let localASRBackend = "config.localASRBackend"
+ static let uiLanguage = "config.uiLanguage"
}
@Published public var providerId: String {
@@ -64,8 +69,8 @@ public final class ProviderConfig: ObservableObject, @unchecked Sendable {
@Published public var localeId: String {
didSet { defaults.set(localeId, forKey: Key.localeId) }
}
- /// "local" → on-device ASR only, no LLM polishing.
- /// "cloud" → ASR + LLM polish (default).
+ /// "local" → on-device ASR only (raw transcript delivery).
+ /// "cloud" → ASR + LLM polish (always on; modeId kept for compatibility).
@Published public var engineMode: String {
didSet { defaults.set(engineMode, forKey: Key.engineMode) }
}
@@ -85,6 +90,16 @@ public final class ProviderConfig: ObservableObject, @unchecked Sendable {
@Published public var hasAcknowledgedCloudSharing: Bool {
didSet { defaults.set(hasAcknowledgedCloudSharing, forKey: Key.hasAcknowledgedCloudSharing) }
}
+ /// Which on-device ASR engine backs the "local" engine mode. Only
+ /// consulted when `isLocalEngine == true`; the cloud engine always
+ /// uses `SpeechAnalyzer`.
+ @Published public var localASRBackend: LocalASRBackend {
+ didSet { defaults.set(localASRBackend.rawValue, forKey: Key.localASRBackend) }
+ }
+ /// Host-app UI language. Also mirrored to the App Group for the keyboard extension.
+ @Published public var uiLanguage: AppUILanguage {
+ didSet { defaults.set(uiLanguage.rawValue, forKey: Key.uiLanguage) }
+ }
public var isConfigured: Bool {
// Local engine (on-device ASR only) doesn't need an API key,
@@ -96,7 +111,7 @@ public final class ProviderConfig: ObservableObject, @unchecked Sendable {
return !baseURL.isEmpty && !apiKey.isEmpty && !model.isEmpty
}
- /// On-device ASR only — no cloud LLM polish.
+ /// On-device ASR only; no cloud API required.
public var isLocalEngine: Bool { engineMode == "local" }
/// The system prompt the user *sees* in the editor — fall back to the
@@ -131,6 +146,19 @@ public final class ProviderConfig: ObservableObject, @unchecked Sendable {
let savedPage = resolvedDefaults.integer(forKey: Key.onboardingPage)
self.onboardingPage = savedPage > 0 ? savedPage : 0
self.hasAcknowledgedCloudSharing = resolvedDefaults.bool(forKey: Key.hasAcknowledgedCloudSharing)
+ // Tolerate missing / unknown raw values (e.g. an enum case that
+ // was renamed in a later build) by falling back to the default
+ // rather than crashing inside `RawRepresentable.init`.
+ let rawBackend = resolvedDefaults.string(forKey: Key.localASRBackend) ?? LocalASRBackend.speechAnalyzer.rawValue
+ self.localASRBackend = LocalASRBackend(rawValue: rawBackend) ?? .speechAnalyzer
+ self.uiLanguage = AppUILanguage.fromStored(
+ resolvedDefaults.string(forKey: Key.uiLanguage)
+ )
+
+ // Cloud no longer exposes off/transcribe; migrate legacy values.
+ if self.engineMode == "cloud", self.modeId != "polish" {
+ self.modeId = "polish"
+ }
}
/// Read the API key from the Keychain, falling back to a one-time
diff --git a/OSGKeyboardShared/Models/TranscriptionDelivery.swift b/OSGKeyboardShared/Models/TranscriptionDelivery.swift
new file mode 100644
index 0000000..a421435
--- /dev/null
+++ b/OSGKeyboardShared/Models/TranscriptionDelivery.swift
@@ -0,0 +1,17 @@
+// TranscriptionDelivery.swift
+// OSGKeyboard · Shared
+//
+// Host-app → keyboard handoff payload: final text plus an optional soft
+// warning when cloud polish failed but the raw transcript is still delivered.
+
+import Foundation
+
+public struct TranscriptionDelivery: Sendable, Equatable {
+ public let text: String
+ public let polishWarning: String?
+
+ public init(text: String, polishWarning: String? = nil) {
+ self.text = text
+ self.polishWarning = polishWarning
+ }
+}
diff --git a/OSGKeyboardShared/Services/ASRService.swift b/OSGKeyboardShared/Services/ASRService.swift
index b2c6cb4..f8b53bd 100644
--- a/OSGKeyboardShared/Services/ASRService.swift
+++ b/OSGKeyboardShared/Services/ASRService.swift
@@ -16,6 +16,7 @@
import Foundation
import AVFoundation
+import CoreMedia
import Speech
import os
@@ -42,6 +43,63 @@ public protocol ASRService: Sendable {
/// Cancel any in-flight recognition and tear down its tasks.
func cancel()
+
+ /// Clears cancellation / cached session state before a new utterance.
+ func resetForNewUtterance()
+
+ /// Transcribe one PCM chunk (Flow pipelined path). Default wraps `transcribe(stream:)`.
+ func transcribeChunk(samples: [Float], locale: Locale) async -> ASRChunkResult
+}
+
+public enum ASRChunkResult: Sendable, Equatable {
+ case success(String)
+ case failure(String)
+ case cancelled
+}
+
+extension ASRService {
+ public func resetForNewUtterance() {}
+
+ public func transcribeChunk(samples: [Float], locale: Locale) async -> ASRChunkResult {
+ guard !samples.isEmpty else { return .success("") }
+ if Task.isCancelled { return .cancelled }
+
+ let snapshot = AudioBufferSnapshot(samples: samples, sampleRate: 16_000)
+ let (stream, continuation) = AsyncStream.makeStream()
+ continuation.yield(snapshot)
+ continuation.finish()
+
+ var lastPartial = ""
+ var finalText = ""
+ var failure: String?
+
+ for await event in transcribe(stream: stream, locale: locale) {
+ if Task.isCancelled { return .cancelled }
+ switch event {
+ case .capability:
+ break
+ case .partial(let text):
+ lastPartial = text
+ case .final(let text):
+ finalText = text
+ case .error(let message):
+ failure = message
+ }
+ }
+
+ if let failure {
+ return .failure(failure)
+ }
+ let trimmed = finalText.trimmingCharacters(in: .whitespacesAndNewlines)
+ if !trimmed.isEmpty {
+ return .success(trimmed)
+ }
+ let partial = lastPartial.trimmingCharacters(in: .whitespacesAndNewlines)
+ if !partial.isEmpty {
+ return .success(partial)
+ }
+ return .success("")
+ }
}
public enum ASREvent: Sendable, Equatable {
@@ -58,13 +116,68 @@ public enum ASREvent: Sendable, Equatable {
// MARK: - Factory
public enum ASRServiceFactory {
- /// Returns the ASR backend. With iOS 26 as the deployment target,
- /// there is exactly one backend (`SpeechAnalyzer`).
+ /// Registry of backend-specific providers. The host app installs
+ /// a provider for the Qwen3-ASR backend at launch time (the
+ /// `Qwen3ASRProvider` lives in the app target because linking
+ /// `Qwen3ASR` pulls in mlx-swift, which the shared framework
+ /// deliberately stays off to keep `APPLICATION_EXTENSION_API_ONLY`
+ /// clean). The shared framework always provides a built-in
+ /// `SpeechAnalyzer` provider; custom providers override it.
+ ///
+ /// `nonisolated(unsafe)` because the only writer is
+ /// `OSGKeyboardApp.init` (single-threaded, runs once at launch).
+ /// After launch, all callers read the dictionary from any
+ /// actor.
+ public nonisolated(unsafe) static var providers: [LocalASRBackend: any ASRServiceProvider] = [
+ .speechAnalyzer: SpeechAnalyzerProvider()
+ ]
+
+ /// Returns the ASR backend chosen by the user. The cloud engine
+ /// always uses the iOS `SpeechAnalyzer` path — it has the lowest
+ /// latency and never hits the network, which matches the user's
+ /// expectation that "ASR" is the local half of the pipeline
+ /// regardless of where the LLM polish happens.
+ ///
+ /// For the local engine, we honour `LocalASRBackend`:
+ /// - `.speechAnalyzer` (default) → on-device iOS pipeline.
+ /// - `.qwen3ASR` → CoreML-backed Qwen3-ASR via `soniqo/speech-swift` (host app only)
+ /// (registered by the host app at launch).
+ public static func make(
+ engineMode: String,
+ localBackend: LocalASRBackend = .speechAnalyzer
+ ) -> ASRService {
+ if engineMode == "local" {
+ if let provider = providers[localBackend] {
+ return provider.make()
+ }
+ }
+ return SpeechAnalyzerASR()
+ }
+
+ /// Back-compat overload for callers that only ever want the
+ /// SpeechAnalyzer path. The previous single-backend build used
+ /// this signature; new code should pass the engine mode explicitly
+ /// so the user's selection is honoured.
public static func make() -> ASRService {
SpeechAnalyzerASR()
}
}
+/// Backend-specific ASR factory. The shared framework ships a default
+/// `SpeechAnalyzerProvider`; the host app installs a `Qwen3ASRProvider`
+/// at launch time so the Qwen3 backend is wired in only where its
+/// large MLX dependency is also linked.
+public protocol ASRServiceProvider: Sendable {
+ var backend: LocalASRBackend { get }
+ func make() -> ASRService
+}
+
+/// Built-in provider for the iOS SpeechAnalyzer path. Always present.
+struct SpeechAnalyzerProvider: ASRServiceProvider {
+ let backend: LocalASRBackend = .speechAnalyzer
+ func make() -> ASRService { SpeechAnalyzerASR() }
+}
+
// MARK: - PCM format conversion (testable helpers)
//
// Extracted from the audio-thread hot path so the scaling + clipping
@@ -115,6 +228,114 @@ final class SpeechAnalyzerASR: ASRService, @unchecked Sendable {
private var analyzer: SpeechAnalyzer?
private var analyzerTask: Task?
private var analyzerFinished = false
+ /// Reused across pipelined chunks within one utterance (assets + format).
+ private var chunkPreparedLocaleID: String?
+ private var chunkAnalyzerFormat: AVAudioFormat?
+
+ func resetForNewUtterance() {
+ lock.withLock {
+ chunkPreparedLocaleID = nil
+ chunkAnalyzerFormat = nil
+ }
+ }
+
+ func transcribeChunk(samples: [Float], locale: Locale) async -> ASRChunkResult {
+ guard !samples.isEmpty else { return .success("") }
+ if Task.isCancelled { return .cancelled }
+
+ do {
+ let text = try await transcribeSamples(samples, locale: locale, reuseChunkPrep: true)
+ let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
+ return trimmed.isEmpty ? .success("") : .success(trimmed)
+ } catch is CancellationError {
+ return .cancelled
+ } catch {
+ return .failure(error.localizedDescription)
+ }
+ }
+
+ /// Analyze a single PCM buffer without the streaming `transcribe` wrapper.
+ private func transcribeSamples(
+ _ samples: [Float],
+ locale: Locale,
+ reuseChunkPrep: Bool
+ ) async throws -> String {
+ guard let resolvedLocale = await DictationTranscriber.supportedLocale(equivalentTo: locale) else {
+ throw ASRChunkError.localeUnsupported
+ }
+ let localeID = resolvedLocale.identifier(.bcp47)
+ let transcriber = DictationTranscriber(
+ locale: resolvedLocale,
+ preset: .progressiveLongDictation
+ )
+
+ let analyzerFormat: AVAudioFormat
+ let cachedPrep = lock.withLock { (chunkPreparedLocaleID, chunkAnalyzerFormat) }
+ if reuseChunkPrep,
+ cachedPrep.0 == localeID,
+ let cached = cachedPrep.1 {
+ analyzerFormat = cached
+ } else {
+ try await Self.prepareAssetsIfNeeded(for: transcriber, locale: resolvedLocale)
+ guard let format = await SpeechAnalyzer.bestAvailableAudioFormat(
+ compatibleWith: [transcriber],
+ considering: Self.captureFormat
+ ) else {
+ throw ASRChunkError.formatUnsupported
+ }
+ analyzerFormat = format
+ lock.withLock {
+ chunkPreparedLocaleID = localeID
+ chunkAnalyzerFormat = format
+ }
+ }
+
+ let snapshot = AudioBufferSnapshot(samples: samples, sampleRate: 16_000)
+ guard let pcm = Self.makeAnalyzerPCMBuffer(from: snapshot, format: analyzerFormat) else {
+ throw ASRChunkError.formatUnsupported
+ }
+
+ let analyzer = SpeechAnalyzer(modules: [transcriber])
+ try await analyzer.prepareToAnalyze(in: analyzerFormat)
+
+ let resultsTask = Task {
+ var accumulator = ProgressiveDictationTranscriptAccumulator()
+ for try await result in transcriber.results {
+ if Task.isCancelled { break }
+ let text = String(result.text.characters)
+ _ = accumulator.ingest(range: result.range, text: text)
+ }
+ return accumulator.finalize()
+ }
+
+ let inputStream = AsyncStream { continuation in
+ continuation.yield(AnalyzerInput(buffer: pcm))
+ continuation.finish()
+ }
+
+ let lastSampleTime = try await analyzer.analyzeSequence(inputStream)
+ if let lastSampleTime {
+ try await analyzer.finalizeAndFinish(through: lastSampleTime)
+ } else {
+ await analyzer.cancelAndFinishNow()
+ }
+
+ return try await resultsTask.value
+ }
+
+ private enum ASRChunkError: LocalizedError {
+ case localeUnsupported
+ case formatUnsupported
+
+ var errorDescription: String? {
+ switch self {
+ case .localeUnsupported:
+ return SharedL10n.string("error.asr.localeUnsupported")
+ case .formatUnsupported:
+ return SharedL10n.string("error.asr.formatUnsupported")
+ }
+ }
+ }
/// Canonical capture format: 16 kHz mono Float32 from `AudioCaptureService`
/// / `PreviewASRController` before it reaches SpeechAnalyzer.
@@ -146,16 +367,21 @@ final class SpeechAnalyzerASR: ASRService, @unchecked Sendable {
do {
guard let resolvedLocale = await DictationTranscriber.supportedLocale(equivalentTo: locale) else {
Self.debug("locale unsupported: \(locale.identifier(.bcp47))")
- continuation.yield(.error("当前系统未分配可用语音语言模型,请稍后重试或切换语言"))
+ continuation.yield(.error(SharedL10n.string("error.asr.localeUnsupported")))
continuation.finish()
return
}
- let transcriber = DictationTranscriber(locale: resolvedLocale, preset: .progressiveShortDictation)
+ // Each pipelined chunk is ≤ 30 s; long dictation preset keeps a
+ // single chunk coherent (Flow utterances run up to 3 min).
+ let transcriber = DictationTranscriber(
+ locale: resolvedLocale,
+ preset: .progressiveLongDictation
+ )
do {
try await Self.prepareAssetsIfNeeded(for: transcriber, locale: resolvedLocale)
} catch {
Self.debug("asset prepare failed: \(error.localizedDescription)")
- continuation.yield(.error("语音语言资源未就绪,请稍后重试"))
+ continuation.yield(.error(SharedL10n.string("error.asr.assetsNotReady")))
continuation.finish()
return
}
@@ -167,7 +393,7 @@ final class SpeechAnalyzerASR: ASRService, @unchecked Sendable {
compatibleWith: [transcriber],
considering: Self.captureFormat
) else {
- continuation.yield(.error("当前设备不支持该语音输入格式"))
+ continuation.yield(.error(SharedL10n.string("error.asr.formatUnsupported")))
continuation.finish()
return
}
@@ -179,15 +405,16 @@ final class SpeechAnalyzerASR: ASRService, @unchecked Sendable {
// Apple recommends consuming `transcriber.results` concurrently
// while `analyzeSequence` drains the input stream.
let resultsTask = Task {
- var lastText = ""
+ var accumulator = ProgressiveDictationTranscriptAccumulator()
for try await result in transcriber.results {
if Task.isCancelled { break }
let text = String(result.text.characters)
- guard !text.isEmpty, text != lastText else { continue }
- lastText = text
- continuation.yield(.partial(text))
+ guard let full = accumulator.ingest(range: result.range, text: text) else {
+ continue
+ }
+ continuation.yield(.partial(full))
}
- return lastText
+ return accumulator.finalize()
}
let lastSampleTime = try await newAnalyzer.analyzeSequence(inputStream)
@@ -195,7 +422,7 @@ final class SpeechAnalyzerASR: ASRService, @unchecked Sendable {
if let lastSampleTime {
try await newAnalyzer.finalizeAndFinish(through: lastSampleTime)
} else {
- try await newAnalyzer.cancelAndFinishNow()
+ await newAnalyzer.cancelAndFinishNow()
}
let lastText: String
@@ -210,7 +437,7 @@ final class SpeechAnalyzerASR: ASRService, @unchecked Sendable {
let trimmed = lastText.trimmingCharacters(in: .whitespacesAndNewlines)
if trimmed.isEmpty {
- continuation.yield(.error("未识别到语音内容,请重试"))
+ continuation.yield(.error(SharedL10n.string("error.asr.noSpeech")))
} else {
continuation.yield(.final(trimmed))
}
diff --git a/OSGKeyboardShared/Services/AppGroupStore.swift b/OSGKeyboardShared/Services/AppGroupStore.swift
index e475de2..9cbc4b9 100644
--- a/OSGKeyboardShared/Services/AppGroupStore.swift
+++ b/OSGKeyboardShared/Services/AppGroupStore.swift
@@ -35,6 +35,8 @@ public struct AppGroupStore: @unchecked Sendable {
static let modeId = "config.modeId"
static let localeId = "config.localeId"
static let engineMode = "config.engineMode"
+ static let localASRBackend = "config.localASRBackend"
+ static let uiLanguage = "config.uiLanguage"
}
// MARK: - Reads
@@ -70,12 +72,25 @@ public struct AppGroupStore: @unchecked Sendable {
defaults.string(forKey: Key.localeId) ?? "auto"
}
- /// "local" → on-device ASR only, no LLM polishing.
+ /// "local" → on-device ASR only (raw transcript delivery).
/// "cloud" → ASR + LLM polish (default behaviour).
public var engineMode: String {
defaults.string(forKey: Key.engineMode) ?? "cloud"
}
+ /// Which on-device ASR engine backs the "local" engine mode. Falls
+ /// back to the iOS SpeechAnalyzer path so legacy installs (which
+ /// never wrote this key) keep working.
+ public var localASRBackend: LocalASRBackend {
+ let raw = defaults.string(forKey: Key.localASRBackend) ?? LocalASRBackend.speechAnalyzer.rawValue
+ return LocalASRBackend(rawValue: raw) ?? .speechAnalyzer
+ }
+
+ /// Host-app UI language override (`auto` / `en` / `zh-Hans`).
+ public var uiLanguage: AppUILanguage {
+ AppUILanguage.fromStored(defaults.string(forKey: Key.uiLanguage))
+ }
+
// MARK: - Writes
public func setModeId(_ id: String) {
@@ -90,6 +105,14 @@ public struct AppGroupStore: @unchecked Sendable {
defaults.set(mode, forKey: Key.engineMode)
}
+ public func setLocalASRBackend(_ backend: LocalASRBackend) {
+ defaults.set(backend.rawValue, forKey: Key.localASRBackend)
+ }
+
+ public func setUILanguage(_ language: AppUILanguage) {
+ defaults.set(language.rawValue, forKey: Key.uiLanguage)
+ }
+
// MARK: - Client
public func makeClient() -> LLMClient {
diff --git a/OSGKeyboardShared/Services/ChunkedUtterancePipeline.swift b/OSGKeyboardShared/Services/ChunkedUtterancePipeline.swift
new file mode 100644
index 0000000..eb0d5f9
--- /dev/null
+++ b/OSGKeyboardShared/Services/ChunkedUtterancePipeline.swift
@@ -0,0 +1,160 @@
+// ChunkedUtterancePipeline.swift
+// OSGKeyboard · Shared
+//
+// Pipelined Flow utterance ASR: split PCM while recording, transcribe chunks
+// serially on a background queue, stitch partials for display and delivery.
+
+import Foundation
+
+public struct ChunkedUtteranceSuccess: Sendable, Equatable {
+ public let text: String
+ /// Non-fatal per-chunk ASR issues (delivered as soft warning when non-empty).
+ public let chunkWarnings: [String]
+
+ public init(text: String, chunkWarnings: [String] = []) {
+ self.text = text
+ self.chunkWarnings = chunkWarnings
+ }
+}
+
+public enum ChunkedUtterancePipelineOutcome: Sendable, Equatable {
+ case success(ChunkedUtteranceSuccess)
+ case failure(String)
+ case cancelled
+}
+
+/// Thread-safe queue between the chunk feeder and ASR worker.
+private actor ChunkWorkQueue {
+ private var items: [UtteranceAudioChunk] = []
+ private var finished = false
+ private var waiters: [CheckedContinuation] = []
+
+ func enqueue(_ chunk: UtteranceAudioChunk) {
+ items.append(chunk)
+ resumeWaiters()
+ }
+
+ func markFinished() {
+ finished = true
+ resumeWaiters()
+ }
+
+ func dequeue() async -> UtteranceAudioChunk? {
+ if !items.isEmpty {
+ return items.removeFirst()
+ }
+ if finished {
+ return nil
+ }
+ return await withCheckedContinuation { continuation in
+ waiters.append(continuation)
+ }
+ }
+
+ private func resumeWaiters() {
+ while !waiters.isEmpty {
+ if !items.isEmpty {
+ let waiter = waiters.removeFirst()
+ waiter.resume(returning: items.removeFirst())
+ } else if finished {
+ let waiter = waiters.removeFirst()
+ waiter.resume(returning: nil)
+ } else {
+ break
+ }
+ }
+ }
+}
+
+public actor ChunkedUtterancePipeline {
+ private let asr: ASRService
+ private let locale: Locale
+ private let config: FlowUtteranceChunkConfig
+ private var cancelled = false
+
+ public init(
+ asr: ASRService,
+ locale: Locale,
+ config: FlowUtteranceChunkConfig = .flowDefault
+ ) {
+ self.asr = asr
+ self.locale = locale
+ self.config = config
+ }
+
+ public func cancel() {
+ cancelled = true
+ asr.cancel()
+ }
+
+ /// Consume `stream` until finished; ASR runs off the caller's actor while recording continues.
+ public func transcribe(
+ stream: AsyncStream,
+ onPartial: @Sendable @escaping (String) -> Void
+ ) async -> ChunkedUtterancePipelineOutcome {
+ asr.resetForNewUtterance()
+
+ let queue = ChunkWorkQueue()
+ var stitcher = UtteranceTranscriptStitcher()
+ var chunkWarnings: [String] = []
+ var failedChunks = 0
+ var processedChunks = 0
+
+ let feeder = Task {
+ for await chunk in UtteranceStreamChunker.chunks(from: stream, config: config) {
+ if Task.isCancelled { break }
+ await queue.enqueue(chunk)
+ }
+ await queue.markFinished()
+ }
+
+ while true {
+ if cancelled || Task.isCancelled {
+ feeder.cancel()
+ return .cancelled
+ }
+
+ guard let chunk = await queue.dequeue() else { break }
+
+ processedChunks += 1
+ let asr = self.asr
+ let locale = self.locale
+ let result = await Task.detached(priority: .userInitiated) {
+ await asr.transcribeChunk(samples: chunk.samples, locale: locale)
+ }.value
+
+ switch result {
+ case .success(let text):
+ stitcher.append(index: chunk.index, text: text)
+ let partial = stitcher.composed()
+ if !partial.isEmpty {
+ onPartial(partial)
+ }
+ case .failure(let message):
+ failedChunks += 1
+ chunkWarnings.append(
+ SharedL10n.format(
+ "error.asr.chunkFailed",
+ chunk.index + 1,
+ message
+ )
+ )
+ case .cancelled:
+ feeder.cancel()
+ return .cancelled
+ }
+ }
+
+ _ = await feeder.value
+
+ let finalText = stitcher.composed().trimmingCharacters(in: .whitespacesAndNewlines)
+ if finalText.isEmpty {
+ if failedChunks > 0, processedChunks == failedChunks {
+ return .failure(SharedL10n.string("error.asr.noSpeech"))
+ }
+ return .failure(SharedL10n.string("error.asr.noSpeech"))
+ }
+
+ return .success(ChunkedUtteranceSuccess(text: finalText, chunkWarnings: chunkWarnings))
+ }
+}
diff --git a/OSGKeyboardShared/Services/DictationBridge.swift b/OSGKeyboardShared/Services/DictationBridge.swift
index a0d3b8f..668b7fe 100644
--- a/OSGKeyboardShared/Services/DictationBridge.swift
+++ b/OSGKeyboardShared/Services/DictationBridge.swift
@@ -21,6 +21,7 @@ public enum DictationBridge {
private enum Key {
static let pendingText = "dictation.pendingText"
+ static let polishWarning = "dictation.polishWarning"
static let updatedAt = "dictation.updatedAt"
static let status = "dictation.status"
static let statusUpdatedAt = "dictation.statusUpdatedAt"
@@ -67,12 +68,21 @@ public enum DictationBridge {
}
/// Store a transcript for the keyboard extension to consume.
- public static func storePendingTranscript(_ text: String, defaults: UserDefaults? = nil) {
+ public static func storePendingTranscript(
+ _ text: String,
+ polishWarning: String? = nil,
+ defaults: UserDefaults? = nil
+ ) {
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty else { return }
let store = resolvedDefaults(defaults)
store.set(trimmed, forKey: Key.pendingText)
store.set(Date().timeIntervalSince1970, forKey: Key.updatedAt)
+ if let polishWarning, !polishWarning.isEmpty {
+ store.set(polishWarning, forKey: Key.polishWarning)
+ } else {
+ store.removeObject(forKey: Key.polishWarning)
+ }
setStatus(.done, defaults: store)
}
@@ -81,6 +91,15 @@ public enum DictationBridge {
maxAge: TimeInterval = 180,
defaults: UserDefaults? = nil
) -> String? {
+ consumePendingDelivery(maxAge: maxAge, defaults: defaults)?.text
+ }
+
+ /// Returns and clears the pending delivery (text + optional polish
+ /// warning) if present.
+ public static func consumePendingDelivery(
+ maxAge: TimeInterval = 180,
+ defaults: UserDefaults? = nil
+ ) -> TranscriptionDelivery? {
let store = resolvedDefaults(defaults)
guard let text = store.string(forKey: Key.pendingText) else {
return nil
@@ -92,14 +111,18 @@ public enum DictationBridge {
return nil
}
}
+ let warning = store.string(forKey: Key.polishWarning)
store.removeObject(forKey: Key.pendingText)
+ store.removeObject(forKey: Key.polishWarning)
+ store.removeObject(forKey: Key.updatedAt)
setStatus(.idle, defaults: store)
- return text
+ return TranscriptionDelivery(text: text, polishWarning: warning)
}
public static func clear(defaults: UserDefaults? = nil) {
let store = resolvedDefaults(defaults)
store.removeObject(forKey: Key.pendingText)
+ store.removeObject(forKey: Key.polishWarning)
store.removeObject(forKey: Key.updatedAt)
setStatus(.idle, defaults: store)
}
diff --git a/OSGKeyboardShared/Services/FlowAppLifecycle.swift b/OSGKeyboardShared/Services/FlowAppLifecycle.swift
new file mode 100644
index 0000000..11fad03
--- /dev/null
+++ b/OSGKeyboardShared/Services/FlowAppLifecycle.swift
@@ -0,0 +1,39 @@
+// FlowAppLifecycle.swift
+// OSGKeyboard · Shared
+//
+// Tracks whether the host app process is in the foreground.
+// Retained for any future GPU-backed paths; CoreML ASR does not require it.
+
+import Foundation
+
+public final class FlowAppLifecycle: @unchecked Sendable {
+
+ public static let shared = FlowAppLifecycle()
+
+ private let lock = NSLock()
+ private var isForeground = true
+
+ private init() {}
+
+ /// `true` when the host app scene is active (`.active`).
+ public var allowsGPUInference: Bool {
+ lock.lock()
+ defer { lock.unlock() }
+ return isForeground
+ }
+
+ public func setForeground(_ foreground: Bool) {
+ lock.lock()
+ isForeground = foreground
+ lock.unlock()
+ }
+
+ /// Blocks until foreground or cancellation.
+ public func waitUntilForeground() async -> Bool {
+ while !allowsGPUInference {
+ if Task.isCancelled { return false }
+ try? await Task.sleep(nanoseconds: 200_000_000)
+ }
+ return true
+ }
+}
diff --git a/OSGKeyboardShared/Services/FlowContinuousCapture.swift b/OSGKeyboardShared/Services/FlowContinuousCapture.swift
index a42394d..80c48bb 100644
--- a/OSGKeyboardShared/Services/FlowContinuousCapture.swift
+++ b/OSGKeyboardShared/Services/FlowContinuousCapture.swift
@@ -34,7 +34,7 @@ private final class FlowCaptureStreamRelay: @unchecked Sendable {
}
func yield(_ snapshot: AudioBufferSnapshot) {
- lock.withLock { continuation?.yield(snapshot) }
+ _ = lock.withLock { continuation?.yield(snapshot) }
}
func finish() {
@@ -250,6 +250,22 @@ public final class FlowContinuousCapture {
)
}
+ /// Re-activate capture after returning from background without
+ /// reinstalling the tap (iOS may deactivate the audio session).
+ public func reassertIfRunning() {
+ guard isRunning else { return }
+ let session = AVAudioSession.sharedInstance()
+ try? session.setCategory(
+ .playAndRecord,
+ mode: .measurement,
+ options: [.defaultToSpeaker, .allowBluetoothHFP, .mixWithOthers]
+ )
+ try? session.setActive(true, options: .notifyOthersOnDeactivation)
+ if !audioEngine.isRunning {
+ try? audioEngine.start()
+ }
+ }
+
/// Begin forwarding downsampled buffers to ASR for one utterance.
public func beginUtterance() -> AsyncStream {
let (stream, continuation) = AsyncStream.makeStream()
diff --git a/OSGKeyboardShared/Services/FlowSessionBridge.swift b/OSGKeyboardShared/Services/FlowSessionBridge.swift
index c455d64..fd8f16c 100644
--- a/OSGKeyboardShared/Services/FlowSessionBridge.swift
+++ b/OSGKeyboardShared/Services/FlowSessionBridge.swift
@@ -65,12 +65,24 @@ public enum FlowSessionBridge {
// MARK: - Session validity (keyboard)
- /// True when expires is in the future and heartbeat is fresh.
+ /// True when the session contract is still valid (not expired).
+ /// Does not require a fresh heartbeat — the host may be suspended in
+ /// background while the continuous audio session is frozen.
public static func isSessionActive(defaults: UserDefaults? = nil) -> Bool {
let store = resolvedDefaults(defaults)
flush(store)
+ guard store.bool(forKey: FlowSessionKeys.flowSessionActive) else { return false }
+
let expires = store.double(forKey: FlowSessionKeys.flowSessionExpires)
- guard expires > Date().timeIntervalSince1970 else { return false }
+ return expires > Date().timeIntervalSince1970
+ }
+
+ /// True when the host app recently wrote a heartbeat (foreground or
+ /// actively processing). Used for auto-start heuristics, not gating record.
+ public static func isHostReachable(defaults: UserDefaults? = nil) -> Bool {
+ let store = resolvedDefaults(defaults)
+ flush(store)
+ guard isSessionActive(defaults: store) else { return false }
let heartbeat = store.double(forKey: FlowSessionKeys.flowHeartbeat)
guard heartbeat > 0 else { return false }
@@ -124,6 +136,7 @@ public enum FlowSessionBridge {
public static func storeTranscriptionResult(
_ text: String,
+ polishWarning: String? = nil,
defaults: UserDefaults? = nil
) {
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
@@ -131,6 +144,11 @@ public enum FlowSessionBridge {
let store = resolvedDefaults(defaults)
store.set(trimmed, forKey: FlowSessionKeys.transcriptionResult)
store.removeObject(forKey: FlowSessionKeys.transcriptionError)
+ if let polishWarning, !polishWarning.isEmpty {
+ store.set(polishWarning, forKey: FlowSessionKeys.transcriptionPolishWarning)
+ } else {
+ store.removeObject(forKey: FlowSessionKeys.transcriptionPolishWarning)
+ }
setRecordingState(.idle, defaults: store)
flush(store)
}
@@ -147,14 +165,24 @@ public enum FlowSessionBridge {
/// Returns and clears a pending transcription result, if any.
public static func consumeTranscriptionResult(defaults: UserDefaults? = nil) -> String? {
+ consumeTranscriptionDelivery(defaults: defaults)?.text
+ }
+
+ /// Returns and clears a pending transcription delivery (text + optional
+ /// polish warning), if any.
+ public static func consumeTranscriptionDelivery(
+ defaults: UserDefaults? = nil
+ ) -> TranscriptionDelivery? {
let store = resolvedDefaults(defaults)
flush(store)
guard let text = store.string(forKey: FlowSessionKeys.transcriptionResult), !text.isEmpty else {
return nil
}
+ let warning = store.string(forKey: FlowSessionKeys.transcriptionPolishWarning)
store.removeObject(forKey: FlowSessionKeys.transcriptionResult)
+ store.removeObject(forKey: FlowSessionKeys.transcriptionPolishWarning)
flush(store)
- return text
+ return TranscriptionDelivery(text: text, polishWarning: warning)
}
/// Returns and clears a pending transcription error, if any.
@@ -212,6 +240,7 @@ public enum FlowSessionBridge {
private static func clearTranscription(defaults: UserDefaults) {
defaults.removeObject(forKey: FlowSessionKeys.transcriptionResult)
+ defaults.removeObject(forKey: FlowSessionKeys.transcriptionPolishWarning)
defaults.removeObject(forKey: FlowSessionKeys.transcriptionError)
}
}
diff --git a/OSGKeyboardShared/Services/FlowSessionKeys.swift b/OSGKeyboardShared/Services/FlowSessionKeys.swift
index 8f947cb..5565940 100644
--- a/OSGKeyboardShared/Services/FlowSessionKeys.swift
+++ b/OSGKeyboardShared/Services/FlowSessionKeys.swift
@@ -13,17 +13,39 @@ public enum FlowSessionKeys {
public static let keyboardRecordingState = "flow.keyboardRecordingState"
public static let transcriptionLanguage = "flow.transcriptionLanguage"
public static let transcriptionResult = "flow.transcriptionResult"
+ /// Soft warning when polish failed but raw transcript was delivered.
+ public static let transcriptionPolishWarning = "flow.transcriptionPolishWarning"
public static let transcriptionError = "flow.transcriptionError"
public static let audioLevels = "flow.audioLevels"
- /// Heartbeat older than this implies the host app was killed.
+ /// Heartbeat older than this while the host is foreground → likely killed.
public static let heartbeatStaleInterval: TimeInterval = 3
/// Default Flow session length when started from the keyboard.
public static let defaultSessionDuration: TimeInterval = 480
- /// Maximum duration for a single keyboard utterance.
- public static let maxUtteranceDuration: TimeInterval = 60
+ /// Maximum duration for a single keyboard utterance (3 minutes).
+ public static let maxUtteranceDuration: TimeInterval = 180
+
+ /// Host polls for pipelined ASR drain after mic stop. Pipelining usually
+ /// finishes most chunks during recording; this is a soft deadline before
+ /// blocking on `asrTask.value` (which waits until the pipeline exits).
+ public static let localASRWaitTimeout: TimeInterval = 120
+ public static let localQwen3ASRWaitTimeout: TimeInterval = 180
+ public static let cloudASRWaitTimeout: TimeInterval = 120
+
+ /// Keyboard watchdog after the user stops recording (not utterance max length).
+ /// Must cover worst-case post-stop backlog: remaining MLX/SpeechAnalyzer chunks
+ /// plus cloud LLM polish (see `PolishingService.effectiveTimeout` cap).
+ public static func keyboardResultTimeout(
+ engineMode: String,
+ localASRBackend: LocalASRBackend
+ ) -> TimeInterval {
+ if engineMode == "local" {
+ return localASRBackend == .qwen3ASR ? 240 : 180
+ }
+ return 240
+ }
public enum RecordingState: String, Sendable, Equatable {
case idle
diff --git a/OSGKeyboardShared/Services/KeyboardState.swift b/OSGKeyboardShared/Services/KeyboardState.swift
index fe24146..60ef7c3 100644
--- a/OSGKeyboardShared/Services/KeyboardState.swift
+++ b/OSGKeyboardShared/Services/KeyboardState.swift
@@ -71,8 +71,17 @@ public final class KeyboardState: ObservableObject {
@Published public var utteranceRemainingSeconds: Int = Int(FlowSessionKeys.maxUtteranceDuration)
/// Whether the host app's Flow voice session is currently valid.
@Published public var flowSessionActive: Bool = false
- /// "local" → ASR only, no LLM. "cloud" → ASR + optional LLM polish.
+ /// "local" → on-device ASR only. "cloud" → ASR + LLM polish.
@Published public var engineMode: String = "cloud"
+ /// Which on-device ASR engine to use when `engineMode == "local"`.
+ /// Mirrored from `ProviderConfig.localASRBackend` for UI display
+ /// and for `state` consumers that want a single source of truth.
+ @Published public var localASRBackend: LocalASRBackend = .speechAnalyzer
+ /// `false` when the local engine needs on-device models that are
+ /// not yet downloaded (mirrored from App Group by the extension).
+ @Published public var localModelsReady: Bool = true
+ /// `true` when host app has preloaded Qwen weights into memory.
+ @Published public var localModelsLoaded: Bool = false
/// Convenience shorthand used by the pipeline and views.
public var isLocalEngine: Bool { engineMode == "local" }
@@ -86,6 +95,7 @@ public final class KeyboardState: ObservableObject {
public var setMode: (InputMode) -> Void = { _ in }
public var setLocale: (String) -> Void = { _ in }
public var setEngineMode: (String) -> Void = { _ in }
+ public var setLocalASRBackend: (LocalASRBackend) -> Void = { _ in }
public var insertNewline: () -> Void = {}
public var insertSpace: () -> Void = {}
public var deleteBackward: () -> Void = {}
diff --git a/OSGKeyboardShared/Services/LLMClient.swift b/OSGKeyboardShared/Services/LLMClient.swift
index d54eac5..f023b63 100644
--- a/OSGKeyboardShared/Services/LLMClient.swift
+++ b/OSGKeyboardShared/Services/LLMClient.swift
@@ -17,13 +17,20 @@ public enum LLMError: Error, LocalizedError, Sendable, Equatable {
public var errorDescription: String? {
switch self {
- case .invalidURL: return "API 地址无效。请在设置中检查 Base URL。"
- case .noAPIKey: return "未填写 API Key。"
- case .http(let s): return "API 返回 HTTP \(s)。请稍后重试或联系服务方。"
- case .decoding: return "解析 API 响应失败。"
- case .transport: return "网络错误,请检查连接后重试。"
- case .rateLimited: return "API 调用过于频繁,请稍候再试。"
- case .cancelled: return "请求已取消。"
+ case .invalidURL:
+ return SharedL10n.string("error.llm.invalidURL")
+ case .noAPIKey:
+ return SharedL10n.string("error.llm.noAPIKey")
+ case .http(let status):
+ return SharedL10n.format("error.llm.http", status)
+ case .decoding:
+ return SharedL10n.string("error.llm.decoding")
+ case .transport:
+ return SharedL10n.string("error.llm.transport")
+ case .rateLimited:
+ return SharedL10n.string("error.llm.rateLimited")
+ case .cancelled:
+ return SharedL10n.string("error.llm.cancelled")
}
}
}
diff --git a/OSGKeyboardShared/Services/LiveDictationController.swift b/OSGKeyboardShared/Services/LiveDictationController.swift
index bee1d52..2f06ad7 100644
--- a/OSGKeyboardShared/Services/LiveDictationController.swift
+++ b/OSGKeyboardShared/Services/LiveDictationController.swift
@@ -36,7 +36,7 @@ private final class CaptureStreamRelay: @unchecked Sendable {
}
func yield(_ snapshot: AudioBufferSnapshot) {
- lock.withLock { continuation?.yield(snapshot) }
+ _ = lock.withLock { continuation?.yield(snapshot) }
}
func finish() {
@@ -73,7 +73,7 @@ public final class LiveDictationController: ObservableObject {
/// next recording starts from zero.
@Published public var lastFinal: String = ""
- private let asr: ASRService = ASRServiceFactory.make()
+ private let asr: ASRService
private let audioEngine = AVAudioEngine()
/// `internal` (not `private`) so the regression test in
/// `OSGKeyboardTests/PreviewASRControllerStateTests.swift` can
@@ -83,10 +83,20 @@ public final class LiveDictationController: ObservableObject {
/// code outside the class from racing on it.
public var asrTask: Task?
private let streamRelay = CaptureStreamRelay()
+ private var chunkedPipeline: ChunkedUtterancePipeline?
private var didConfigureAudioSession = false
private var didInstallTap = false
- public init() {}
+ public init(asr: ASRService? = nil) {
+ // Resolve through the factory so the user's `LocalASRBackend`
+ // selection is honoured. Tests can pass a stub `asr` directly
+ // to bypass the factory and exercise the controller in
+ // isolation.
+ self.asr = asr ?? ASRServiceFactory.make(
+ engineMode: ProviderConfig.shared.engineMode,
+ localBackend: ProviderConfig.shared.localASRBackend
+ )
+ }
/// Start dictation using a persisted settings locale id (`auto`, `zh-Hans`, …).
public func start(localeId: String) async {
@@ -112,6 +122,10 @@ public final class LiveDictationController: ObservableObject {
// same `events` stream.
asrTask?.cancel()
asrTask = nil
+ if let pipeline = chunkedPipeline {
+ Task { await pipeline.cancel() }
+ }
+ chunkedPipeline = nil
teardownCapturePipeline()
phase = .requestingPermission
currentPartial = ""
@@ -335,29 +349,36 @@ public final class LiveDictationController: ObservableObject {
return
}
- // 5. Wire up ASR.
- let events = asr.transcribe(
- stream: stream,
- locale: locale
- )
- asrTask = Task { @MainActor [weak self] in
- guard let self else { return }
- for await event in events {
- switch event {
- case .capability:
- break
- case .partial(let s):
- self.currentPartial = s
- case .final(let s):
- let trimmed = s.trimmingCharacters(in: .whitespacesAndNewlines)
- self.lastFinal = trimmed
- self.currentPartial = ""
- self.phase = .idle
- case .error(let m):
- self.debug("asr error: \(m)")
- self.teardownCapturePipeline()
- self.errorMessage = m
- self.phase = .error(m)
+ // 5. Pipelined ASR (same chunk path as Flow host).
+ let pipeline = ChunkedUtterancePipeline(asr: asr, locale: locale)
+ chunkedPipeline = pipeline
+ asrTask = Task.detached(priority: .userInitiated) { [weak controller = self] in
+ let outcome = await pipeline.transcribe(stream: stream) { partial in
+ Task { @MainActor in
+ controller?.currentPartial = partial
+ }
+ }
+ await MainActor.run {
+ guard let controller else { return }
+ switch outcome {
+ case .success(let success):
+ let trimmed = success.text.trimmingCharacters(in: .whitespacesAndNewlines)
+ if !trimmed.isEmpty {
+ controller.lastFinal = trimmed
+ controller.currentPartial = ""
+ }
+ if controller.phase == .processing || controller.phase == .recording {
+ controller.phase = .idle
+ }
+ case .failure(let message):
+ controller.debug("asr error: \(message)")
+ controller.teardownCapturePipeline()
+ controller.errorMessage = message
+ controller.phase = .error(message)
+ case .cancelled:
+ if controller.phase == .processing {
+ controller.phase = .idle
+ }
}
}
}
diff --git a/OSGKeyboardShared/Services/OnDeviceModelStatus.swift b/OSGKeyboardShared/Services/OnDeviceModelStatus.swift
new file mode 100644
index 0000000..9b817f2
--- /dev/null
+++ b/OSGKeyboardShared/Services/OnDeviceModelStatus.swift
@@ -0,0 +1,104 @@
+// OnDeviceModelStatus.swift
+// OSGKeyboard · Shared
+//
+// Mirrors on-device model download state into the App Group so the
+// keyboard extension can show readiness hints without reading the
+// host app's Caches directory.
+
+import Foundation
+
+public enum OnDeviceModelStatus {
+
+ private enum Key {
+ static func downloaded(_ model: OnDeviceModel) -> String {
+ "models.\(model.rawValue).downloaded"
+ }
+ static func progress(_ model: OnDeviceModel) -> String {
+ "models.\(model.rawValue).downloadProgress"
+ }
+ static let modelsLoadedInMemory = "models.loadedInMemory"
+ }
+
+ // MARK: - Writes (host app)
+
+ public static func setDownloaded(_ downloaded: Bool, for model: OnDeviceModel) {
+ guard AppGroup.isAvailable else { return }
+ AppGroup.defaults.set(downloaded, forKey: Key.downloaded(model))
+ if downloaded {
+ clearProgress(for: model)
+ }
+ }
+
+ public static func setProgress(_ progress: Double?, for model: OnDeviceModel) {
+ guard AppGroup.isAvailable else { return }
+ if let progress {
+ AppGroup.defaults.set(progress, forKey: Key.progress(model))
+ } else {
+ AppGroup.defaults.removeObject(forKey: Key.progress(model))
+ }
+ }
+
+ public static func clearProgress(for model: OnDeviceModel) {
+ guard AppGroup.isAvailable else { return }
+ AppGroup.defaults.removeObject(forKey: Key.progress(model))
+ }
+
+ public static func setModelsLoadedInMemory(_ loaded: Bool) {
+ guard AppGroup.isAvailable else { return }
+ AppGroup.defaults.set(loaded, forKey: Key.modelsLoadedInMemory)
+ }
+
+ public static func modelsLoadedInMemory(defaults: UserDefaults? = nil) -> Bool {
+ let store = defaults ?? (AppGroup.isAvailable ? AppGroup.defaults : .standard)
+ return store.bool(forKey: Key.modelsLoadedInMemory)
+ }
+
+ // MARK: - Reads (keyboard + host app)
+
+ public static func isDownloaded(_ model: OnDeviceModel, defaults: UserDefaults? = nil) -> Bool {
+ let store = defaults ?? (AppGroup.isAvailable ? AppGroup.defaults : .standard)
+ return store.bool(forKey: Key.downloaded(model))
+ }
+
+ public static func downloadProgress(_ model: OnDeviceModel, defaults: UserDefaults? = nil) -> Double? {
+ let store = defaults ?? (AppGroup.isAvailable ? AppGroup.defaults : .standard)
+ guard store.object(forKey: Key.progress(model)) != nil else { return nil }
+ return store.double(forKey: Key.progress(model))
+ }
+
+ /// Whether the currently selected local-engine stack has every
+ /// required on-device model downloaded.
+ public static func isLocalStackReady(
+ asrBackend: LocalASRBackend,
+ defaults: UserDefaults? = nil
+ ) -> Bool {
+ if asrBackend == .qwen3ASR {
+ return isDownloaded(.qwen3ASR, defaults: defaults)
+ }
+ return true
+ }
+
+ /// First missing model for the active local stack, if any.
+ public static func firstMissingModel(
+ asrBackend: LocalASRBackend,
+ defaults: UserDefaults? = nil
+ ) -> OnDeviceModel? {
+ if asrBackend == .qwen3ASR, !isDownloaded(.qwen3ASR, defaults: defaults) {
+ return .qwen3ASR
+ }
+ return nil
+ }
+}
+
+// MARK: - On-device Qwen3 runtime
+
+/// CoreML ASR requires iOS 18+ / macOS 15+ (MLState KV cache).
+public enum OnDeviceMLRuntime {
+ /// Whether Qwen3-ASR CoreML can run in this process.
+ public static var supportsOnDeviceQwen3: Bool {
+ if #available(iOS 18.0, *) {
+ return true
+ }
+ return false
+ }
+}
diff --git a/OSGKeyboardShared/Services/PolishingService.swift b/OSGKeyboardShared/Services/PolishingService.swift
index bc70666..9083e88 100644
--- a/OSGKeyboardShared/Services/PolishingService.swift
+++ b/OSGKeyboardShared/Services/PolishingService.swift
@@ -5,9 +5,9 @@
// to produce polished, well-punctuated text. Falls back to the raw transcript
// if the LLM call fails or times out.
//
-// Mode-aware: when `modeId == "off"` the service short-circuits and returns
-// the trimmed input without touching the network. This is the runtime
-// guarantee behind the keyboard's "Off · 关闭" mode.
+// Cloud engine always runs the LLM polish step (settings no longer expose
+// off / transcribe). Local engine (`engineMode == "local"`) is ASR-only —
+// the raw transcript is returned unchanged and cloud API settings are ignored.
import Foundation
@@ -16,7 +16,6 @@ public actor PolishingService {
public enum PolishError: Error, Equatable {
case noTranscript
case timeout
- case modeOff
}
private let store: AppGroupStore
@@ -44,25 +43,25 @@ public actor PolishingService {
let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty else { throw PolishError.noTranscript }
- // Mode-aware short-circuit. When the user has selected "Off", the
- // keyboard must never hit the network — we return the trimmed
- // input as-is. This is the same value the view controller would
- // produce if it skipped `polish()` entirely, but having the
- // guarantee at the service layer means future call sites (CLI,
- // tests, alternate keyboards) inherit it for free.
- if store.modeId == "off" {
+ // Local engine: ASR-only — no on-device or cloud polish.
+ if store.engineMode == "local" {
return trimmed
}
+ return try await polishRemote(trimmed)
+ }
+
+ private func polishRemote(_ trimmed: String) async throws -> String {
let client = injectedClient ?? store.makeClient()
let prompt = store.systemPrompt
+ let budget = effectiveTimeout(for: trimmed)
return try await withThrowingTaskGroup(of: String.self) { group in
group.addTask {
try await client.polish(trimmed, systemPrompt: prompt)
}
group.addTask {
- try await Task.sleep(nanoseconds: UInt64(self.timeout * 1_000_000_000))
+ try await Task.sleep(nanoseconds: UInt64(budget * 1_000_000_000))
throw PolishError.timeout
}
let result = try await group.next()!
@@ -70,4 +69,10 @@ public actor PolishingService {
return result
}
}
-}
\ No newline at end of file
+
+ /// Scale polish budget with transcript length (3-minute Flow utterances).
+ private func effectiveTimeout(for text: String) -> TimeInterval {
+ let scaled = timeout + (Double(text.count) / 200.0) * 2.0
+ return min(max(scaled, timeout), 120)
+ }
+}
diff --git a/OSGKeyboardShared/Utilities/ProgressiveDictationTranscriptAccumulator.swift b/OSGKeyboardShared/Utilities/ProgressiveDictationTranscriptAccumulator.swift
new file mode 100644
index 0000000..30d61ff
--- /dev/null
+++ b/OSGKeyboardShared/Utilities/ProgressiveDictationTranscriptAccumulator.swift
@@ -0,0 +1,64 @@
+// ProgressiveDictationTranscriptAccumulator.swift
+// OSGKeyboard · Shared
+//
+// Merges progressive `DictationTranscriber` results into one transcript.
+// Short-form presets may emit a new time range after ~30 s; treating the
+// latest partial as the full transcript drops earlier segments.
+
+import Foundation
+import CoreMedia
+
+/// Combines volatile partials and finalized segments from
+/// `DictationTranscriber.results` into a single growing transcript.
+public struct ProgressiveDictationTranscriptAccumulator: Sendable {
+
+ private struct Segment: Sendable {
+ let startSeconds: Double
+ var text: String
+ }
+
+ private var segments: [Segment] = []
+ private var lastEmitted = ""
+
+ public init() {}
+
+ /// Ingest one analyzer result. Returns a non-nil full transcript when the
+ /// composed text changed since the previous emission.
+ public mutating func ingest(range: CMTimeRange, text: String) -> String? {
+ let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
+ guard !trimmed.isEmpty else { return nil }
+
+ let start = range.start.seconds
+
+ if let idx = segments.lastIndex(where: { abs($0.startSeconds - start) < 0.001 }) {
+ // Same audio window — volatile refinement of the current segment.
+ segments[idx].text = trimmed
+ } else if let last = segments.last,
+ trimmed.hasPrefix(last.text) || last.text.hasPrefix(trimmed) {
+ // Cumulative progressive update without a range change.
+ let longer = trimmed.count >= last.text.count ? trimmed : last.text
+ segments[segments.count - 1].text = longer
+ } else {
+ // New time range — append instead of replacing earlier speech.
+ segments.append(Segment(startSeconds: start, text: trimmed))
+ }
+
+ let full = composedText()
+ guard full != lastEmitted else { return nil }
+ lastEmitted = full
+ return full
+ }
+
+ /// Final composed transcript after the results stream finishes.
+ public mutating func finalize() -> String {
+ let full = composedText()
+ lastEmitted = full
+ return full
+ }
+
+ private func composedText() -> String {
+ segments.reduce(into: "") { partial, segment in
+ partial = DictationTextComposer.compose(anchor: partial, live: segment.text)
+ }
+ }
+}
diff --git a/OSGKeyboardShared/Utilities/ProviderDisplayName.swift b/OSGKeyboardShared/Utilities/ProviderDisplayName.swift
index 64d824c..d0786b8 100644
--- a/OSGKeyboardShared/Utilities/ProviderDisplayName.swift
+++ b/OSGKeyboardShared/Utilities/ProviderDisplayName.swift
@@ -6,9 +6,12 @@
import Foundation
public enum ProviderDisplayName {
- public static func name(for providerId: String) -> String {
+ public static func name(
+ for providerId: String,
+ language: AppUILanguage? = nil
+ ) -> String {
let key = "provider.\(providerId)"
- let localized = NSLocalizedString(key, comment: "")
+ let localized = SharedL10n.string(key, language: language)
if localized != key { return localized }
return LLMProvider.provider(id: providerId).name
}
diff --git a/OSGKeyboardShared/Utilities/UtteranceStreamChunker.swift b/OSGKeyboardShared/Utilities/UtteranceStreamChunker.swift
new file mode 100644
index 0000000..4ce2e44
--- /dev/null
+++ b/OSGKeyboardShared/Utilities/UtteranceStreamChunker.swift
@@ -0,0 +1,101 @@
+// UtteranceStreamChunker.swift
+// OSGKeyboard · Shared
+//
+// Splits a Flow utterance PCM stream into ASR-sized chunks. When possible,
+// extends slightly past the max window to the next pause instead of cutting
+// mid-word.
+
+import Foundation
+
+public enum UtteranceStreamChunker {
+
+ /// Yields chunks as audio arrives; the final chunk is marked `isLast`.
+ public static func chunks(
+ from stream: AsyncStream,
+ config: FlowUtteranceChunkConfig = .flowDefault
+ ) -> AsyncStream {
+ AsyncStream { continuation in
+ let task = Task {
+ var buffer: [Float] = []
+ buffer.reserveCapacity(config.maxChunkSamples + config.pauseExtensionSamples)
+ var chunkIndex = 0
+
+ func emit(upTo splitEnd: Int, isLast: Bool) {
+ guard splitEnd > 0, splitEnd <= buffer.count else { return }
+ let chunkSamples = Array(buffer[..= buffer.count {
+ buffer.removeAll(keepingCapacity: true)
+ } else {
+ let overlapStart = max(0, splitEnd - config.overlapSamples)
+ buffer = Array(buffer[overlapStart...])
+ }
+ }
+
+ for await snap in stream {
+ if Task.isCancelled { break }
+ guard !snap.samples.isEmpty else { continue }
+ buffer.append(contentsOf: snap.samples)
+
+ while buffer.count >= config.maxChunkSamples {
+ let split = pauseAwareSplitIndex(in: buffer, config: config)
+ emit(upTo: split, isLast: false)
+ }
+ }
+
+ if !buffer.isEmpty {
+ emit(upTo: buffer.count, isLast: true)
+ } else if chunkIndex == 0 {
+ // Empty utterance — no chunks.
+ } else {
+ // Stream ended exactly on boundary; mark prior path complete.
+ }
+
+ continuation.finish()
+ }
+
+ continuation.onTermination = { _ in
+ task.cancel()
+ }
+ }
+ }
+
+ /// Pick a split index at or after `maxChunkSamples`, preferring a pause.
+ static func pauseAwareSplitIndex(
+ in buffer: [Float],
+ config: FlowUtteranceChunkConfig
+ ) -> Int {
+ let minSplit = config.maxChunkSamples
+ guard buffer.count >= minSplit else { return buffer.count }
+
+ let searchEnd = min(buffer.count, minSplit + config.pauseExtensionSamples)
+ if searchEnd <= minSplit {
+ return minSplit
+ }
+
+ let windowSize = max(config.sampleRate / 50, 160) // ~20 ms
+ var bestPause: Int?
+ var idx = minSplit
+ while idx + windowSize <= searchEnd {
+ if rms(of: buffer, start: idx, count: windowSize) < config.pauseRMSThreshold {
+ bestPause = idx + windowSize
+ }
+ idx += windowSize / 2
+ }
+
+ return bestPause ?? minSplit
+ }
+
+ static func rms(of samples: [Float], start: Int, count: Int) -> Float {
+ guard start >= 0, count > 0, start + count <= samples.count else { return 1 }
+ var sum: Float = 0
+ for i in start..<(start + count) {
+ let v = samples[i]
+ sum += v * v
+ }
+ return sqrtf(sum / Float(count))
+ }
+}
diff --git a/OSGKeyboardShared/Utilities/UtteranceTranscriptStitcher.swift b/OSGKeyboardShared/Utilities/UtteranceTranscriptStitcher.swift
new file mode 100644
index 0000000..e33b1c6
--- /dev/null
+++ b/OSGKeyboardShared/Utilities/UtteranceTranscriptStitcher.swift
@@ -0,0 +1,109 @@
+// UtteranceTranscriptStitcher.swift
+// OSGKeyboard · Shared
+//
+// Orders pipelined chunk transcripts and merges overlap at boundaries.
+
+import Foundation
+
+public struct UtteranceTranscriptStitcher: Sendable {
+ private var segments: [(index: Int, text: String)] = []
+
+ public init() {}
+
+ public mutating func append(index: Int, text: String) {
+ let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
+ guard !trimmed.isEmpty else { return }
+ if let existing = segments.firstIndex(where: { $0.index == index }) {
+ segments[existing].text = trimmed
+ } else {
+ segments.append((index, trimmed))
+ segments.sort { $0.index < $1.index }
+ }
+ }
+
+ public func composed() -> String {
+ guard let first = segments.first else { return "" }
+ var result = first.text
+ for segment in segments.dropFirst() {
+ result = Self.mergeWithOverlap(previous: result, next: segment.text)
+ }
+ return result
+ }
+
+ /// Merge `next` onto `previous`, dropping duplicated suffix/prefix overlap.
+ public static func mergeWithOverlap(previous: String, next: String) -> String {
+ let trimmedNext = next.trimmingCharacters(in: .whitespacesAndNewlines)
+ guard !trimmedNext.isEmpty else { return previous }
+ guard !previous.isEmpty else { return trimmedNext }
+
+ // Character-granular probe — works for CJK without word boundaries.
+ let prevChars = Array(previous)
+ let nextChars = Array(trimmedNext)
+ let maxProbe = min(64, prevChars.count, nextChars.count)
+ if maxProbe > 0 {
+ for length in stride(from: maxProbe, through: 1, by: -1) {
+ let suffix = prevChars.suffix(length)
+ let prefix = nextChars.prefix(length)
+ if suffix.elementsEqual(prefix) {
+ return previous + String(nextChars.dropFirst(length))
+ }
+ }
+ }
+
+ // Punctuation-insensitive CJK overlap (e.g. "很好," + "很好继续").
+ let normalizedPrev = normalizeForOverlap(previous)
+ let normalizedNext = normalizeForOverlap(trimmedNext)
+ let nPrev = Array(normalizedPrev)
+ let nNext = Array(normalizedNext)
+ let normProbe = min(64, nPrev.count, nNext.count)
+ if normProbe > 0 {
+ for length in stride(from: normProbe, through: 2, by: -1) {
+ if nPrev.suffix(length).elementsEqual(nNext.prefix(length)) {
+ // Map normalized overlap length back to raw `next` drop count.
+ let drop = overlapDropCount(in: trimmedNext, normalizedPrefixLength: length)
+ return previous + String(trimmedNext.dropFirst(drop))
+ }
+ }
+ }
+
+ // English / spaced languages.
+ let maxWordProbe = min(6, previous.split(separator: " ").count, trimmedNext.split(separator: " ").count)
+ if maxWordProbe > 0 {
+ let prevWords = previous.split(separator: " ", omittingEmptySubsequences: true)
+ let nextWords = trimmedNext.split(separator: " ", omittingEmptySubsequences: true)
+ for wordCount in stride(from: maxWordProbe, through: 1, by: -1) {
+ if prevWords.suffix(wordCount).elementsEqual(nextWords.prefix(wordCount)) {
+ let mergedPrefix = nextWords.dropFirst(wordCount).joined(separator: " ")
+ if mergedPrefix.isEmpty { return previous }
+ if previous.last == " " || previous.last == "\n" {
+ return previous + mergedPrefix
+ }
+ return previous + " " + mergedPrefix
+ }
+ }
+ }
+
+ return DictationTextComposer.compose(anchor: previous, live: trimmedNext)
+ }
+
+ private static func normalizeForOverlap(_ text: String) -> String {
+ text.unicodeScalars.filter {
+ !CharacterSet.whitespacesAndNewlines.contains($0)
+ && !CharacterSet.punctuationCharacters.contains($0)
+ }.map { Character($0) }.reduce(into: "") { $0.append($1) }
+ }
+
+ /// How many raw characters to drop from `next` given a normalized-prefix overlap length.
+ private static func overlapDropCount(in next: String, normalizedPrefixLength: Int) -> Int {
+ var normalizedCount = 0
+ var rawIndex = next.startIndex
+ while rawIndex < next.endIndex, normalizedCount < normalizedPrefixLength {
+ let scalar = next[rawIndex]
+ if !scalar.isWhitespace, !scalar.isPunctuation {
+ normalizedCount += 1
+ }
+ rawIndex = next.index(after: rawIndex)
+ }
+ return next.distance(from: next.startIndex, to: rawIndex)
+ }
+}
diff --git a/OSGKeyboardShared/en.lproj/Shared.strings b/OSGKeyboardShared/en.lproj/Shared.strings
new file mode 100644
index 0000000..a4814e6
--- /dev/null
+++ b/OSGKeyboardShared/en.lproj/Shared.strings
@@ -0,0 +1,30 @@
+/* Engine status labels */
+"engine.summary.local" = "On-device · %@";
+"engine.summary.cloud" = "Active: %@";
+"engine.summary.cloudWithModel" = "Active: %1$@ · %2$@";
+"engine.asr.appleSpeech" = "Apple SpeechAnalyzer";
+"model.qwen3asr.name" = "Qwen3-ASR 0.6B (CoreML)";
+
+/* LLM providers */
+"provider.openai" = "OpenAI";
+"provider.deepseek" = "DeepSeek";
+"provider.qwen" = "Qwen (DashScope)";
+"provider.zhipu" = "Zhipu GLM";
+"provider.moonshot" = "Moonshot";
+"provider.custom" = "Custom";
+
+/* LLM errors */
+"error.llm.invalidURL" = "Invalid API URL. Check Base URL in Settings.";
+"error.llm.noAPIKey" = "API Key is missing.";
+"error.llm.http" = "API returned HTTP %lld. Try again later or contact the provider.";
+"error.llm.decoding" = "Failed to parse the API response.";
+"error.llm.transport" = "Network error. Check your connection and try again.";
+"error.llm.rateLimited" = "Too many API requests. Please wait and try again.";
+"error.llm.cancelled" = "Request cancelled.";
+
+/* ASR errors */
+"error.asr.localeUnsupported" = "Speech language assets are unavailable. Try again later or switch the recognition language.";
+"error.asr.assetsNotReady" = "Speech language assets are not ready. Try again later.";
+"error.asr.formatUnsupported" = "This device does not support the required audio format.";
+"error.asr.noSpeech" = "No speech detected. Please try again.";
+"error.asr.chunkFailed" = "Segment %lld failed: %@";
diff --git a/OSGKeyboardShared/zh-Hans.lproj/Shared.strings b/OSGKeyboardShared/zh-Hans.lproj/Shared.strings
new file mode 100644
index 0000000..8763c4b
--- /dev/null
+++ b/OSGKeyboardShared/zh-Hans.lproj/Shared.strings
@@ -0,0 +1,30 @@
+/* Engine status labels */
+"engine.summary.local" = "本地 · %@";
+"engine.summary.cloud" = "当前:%@";
+"engine.summary.cloudWithModel" = "当前:%1$@ · %2$@";
+"engine.asr.appleSpeech" = "Apple 语音识别";
+"model.qwen3asr.name" = "Qwen3-ASR 0.6B (CoreML)";
+
+/* LLM providers */
+"provider.openai" = "OpenAI";
+"provider.deepseek" = "DeepSeek";
+"provider.qwen" = "通义千问";
+"provider.zhipu" = "智谱 GLM";
+"provider.moonshot" = "月之暗面";
+"provider.custom" = "自定义";
+
+/* LLM errors */
+"error.llm.invalidURL" = "API 地址无效。请在设置中检查 Base URL。";
+"error.llm.noAPIKey" = "未填写 API Key。";
+"error.llm.http" = "API 返回 HTTP %lld。请稍后重试或联系服务方。";
+"error.llm.decoding" = "解析 API 响应失败。";
+"error.llm.transport" = "网络错误,请检查连接后重试。";
+"error.llm.rateLimited" = "API 调用过于频繁,请稍候再试。";
+"error.llm.cancelled" = "请求已取消。";
+
+/* ASR errors */
+"error.asr.localeUnsupported" = "当前系统未分配可用语音语言模型,请稍后重试或切换语言。";
+"error.asr.assetsNotReady" = "语音语言资源未就绪,请稍后重试。";
+"error.asr.formatUnsupported" = "当前设备不支持该语音输入格式。";
+"error.asr.noSpeech" = "未识别到语音内容,请重试。";
+"error.asr.chunkFailed" = "第 %lld 段识别失败:%@";
diff --git a/OSGKeyboardTests/ChunkedUtterancePipelineTests.swift b/OSGKeyboardTests/ChunkedUtterancePipelineTests.swift
new file mode 100644
index 0000000..f6f7ad9
--- /dev/null
+++ b/OSGKeyboardTests/ChunkedUtterancePipelineTests.swift
@@ -0,0 +1,114 @@
+// ChunkedUtterancePipelineTests.swift
+// OSGKeyboardTests
+
+import XCTest
+@testable import OSGKeyboardShared
+
+private struct StubChunkASR: ASRService, @unchecked Sendable {
+ let labels: @Sendable ([Float]) -> String
+
+ func transcribe(
+ stream: AsyncStream,
+ locale: Locale
+ ) -> AsyncStream {
+ AsyncStream { $0.finish() }
+ }
+
+ func cancel() {}
+
+ func transcribeChunk(samples: [Float], locale: Locale) async -> ASRChunkResult {
+ _ = locale
+ return .success(labels(samples))
+ }
+}
+
+final class ChunkedUtterancePipelineTests: XCTestCase {
+
+ func testPipelineStitchesQueuedChunks() async {
+ let config = FlowUtteranceChunkConfig(
+ maxChunkDurationSeconds: 0.05,
+ overlapDurationSeconds: 0,
+ pauseExtensionMaxSeconds: 0,
+ pauseRMSThreshold: 0.02,
+ sampleRate: 1_000
+ )
+ let asr = StubChunkASR { samples in
+ samples.isEmpty ? "" : "seg\(samples.count)"
+ }
+ let pipeline = ChunkedUtterancePipeline(
+ asr: asr,
+ locale: Locale(identifier: "zh-Hans"),
+ config: config
+ )
+
+ let (stream, continuation) = AsyncStream.makeStream()
+ continuation.yield(AudioBufferSnapshot(samples: [Float](repeating: 0.1, count: 80), sampleRate: 1_000))
+ continuation.yield(AudioBufferSnapshot(samples: [Float](repeating: 0.1, count: 80), sampleRate: 1_000))
+ continuation.finish()
+
+ var partials: [String] = []
+ let outcome = await pipeline.transcribe(stream: stream) { partial in
+ partials.append(partial)
+ }
+
+ guard case .success(let success) = outcome else {
+ return XCTFail("expected success, got \(outcome)")
+ }
+ XCTAssertTrue(success.text.contains("seg"))
+ XCTAssertFalse(partials.isEmpty)
+ }
+
+ func testPipelineDeliversPartialSuccessWhenOneChunkFails() async {
+ let config = FlowUtteranceChunkConfig(
+ maxChunkDurationSeconds: 0.05,
+ overlapDurationSeconds: 0,
+ pauseExtensionMaxSeconds: 0,
+ pauseRMSThreshold: 0.02,
+ sampleRate: 1_000
+ )
+ let pipeline = ChunkedUtterancePipeline(
+ asr: FailingSecondChunkASR(),
+ locale: Locale(identifier: "zh-Hans"),
+ config: config
+ )
+
+ let (stream, continuation) = AsyncStream.makeStream()
+ continuation.yield(AudioBufferSnapshot(samples: [Float](repeating: 0.1, count: 80), sampleRate: 1_000))
+ continuation.yield(AudioBufferSnapshot(samples: [Float](repeating: 0.1, count: 80), sampleRate: 1_000))
+ continuation.finish()
+
+ let outcome = await pipeline.transcribe(stream: stream) { _ in }
+
+ guard case .success(let success) = outcome else {
+ return XCTFail("expected partial success, got \(outcome)")
+ }
+ XCTAssertFalse(success.text.isEmpty)
+ XCTAssertEqual(success.chunkWarnings.count, 1)
+ }
+}
+
+private struct FailingSecondChunkASR: ASRService, @unchecked Sendable {
+ private let lock = OSAllocatedUnfairLock()
+ private var index = 0
+
+ func transcribe(
+ stream: AsyncStream,
+ locale: Locale
+ ) -> AsyncStream {
+ AsyncStream { $0.finish() }
+ }
+
+ func cancel() {}
+
+ func transcribeChunk(samples: [Float], locale: Locale) async -> ASRChunkResult {
+ _ = locale
+ let current = lock.withLock {
+ defer { index += 1 }
+ return index
+ }
+ if current == 1 {
+ return .failure("simulated chunk error")
+ }
+ return .success("seg\(samples.count)")
+ }
+}
diff --git a/OSGKeyboardTests/FlowSessionBridgeTests.swift b/OSGKeyboardTests/FlowSessionBridgeTests.swift
index c568e48..f09be5a 100644
--- a/OSGKeyboardTests/FlowSessionBridgeTests.swift
+++ b/OSGKeyboardTests/FlowSessionBridgeTests.swift
@@ -13,14 +13,25 @@ final class FlowSessionBridgeTests: XCTestCase {
return defaults
}
- func testSessionActiveRequiresFreshHeartbeat() {
+ func testSessionActiveSurvivesStaleHeartbeatWhileNotExpired() {
let defaults = makeDefaults()
FlowSessionBridge.markSessionActive(duration: 60, defaults: defaults)
XCTAssertTrue(FlowSessionBridge.isSessionActive(defaults: defaults))
+ XCTAssertTrue(FlowSessionBridge.isHostReachable(defaults: defaults))
let staleHeartbeat = Date().timeIntervalSince1970 - 10
defaults.set(staleHeartbeat, forKey: FlowSessionKeys.flowHeartbeat)
+ XCTAssertTrue(FlowSessionBridge.isSessionActive(defaults: defaults))
+ XCTAssertFalse(FlowSessionBridge.isHostReachable(defaults: defaults))
+ }
+
+ func testSessionInactiveWhenExpired() {
+ let defaults = makeDefaults()
+ FlowSessionBridge.markSessionActive(duration: 1, defaults: defaults)
+ let expired = Date().timeIntervalSince1970 - 5
+ defaults.set(expired, forKey: FlowSessionKeys.flowSessionExpires)
XCTAssertFalse(FlowSessionBridge.isSessionActive(defaults: defaults))
+ XCTAssertFalse(FlowSessionBridge.isHostReachable(defaults: defaults))
}
func testRecordingStateRoundTrip() {
@@ -39,6 +50,19 @@ final class FlowSessionBridgeTests: XCTestCase {
XCTAssertNil(FlowSessionBridge.consumeTranscriptionResult(defaults: defaults))
}
+ func testConsumeTranscriptionDeliveryIncludesPolishWarning() {
+ let defaults = makeDefaults()
+ FlowSessionBridge.storeTranscriptionResult(
+ "raw text",
+ polishWarning: "polish failed",
+ defaults: defaults
+ )
+ let delivery = FlowSessionBridge.consumeTranscriptionDelivery(defaults: defaults)
+ XCTAssertEqual(delivery?.text, "raw text")
+ XCTAssertEqual(delivery?.polishWarning, "polish failed")
+ XCTAssertNil(FlowSessionBridge.consumeTranscriptionDelivery(defaults: defaults))
+ }
+
func testClearFlowStateRemovesSessionKeys() {
let defaults = makeDefaults()
FlowSessionBridge.markSessionActive(defaults: defaults)
diff --git a/OSGKeyboardTests/LLMClientTests.swift b/OSGKeyboardTests/LLMClientTests.swift
index e5fd1c9..65cd5a2 100644
--- a/OSGKeyboardTests/LLMClientTests.swift
+++ b/OSGKeyboardTests/LLMClientTests.swift
@@ -232,8 +232,7 @@ final class LLMClientTests: XCTestCase {
}
/// Cross-process App Group contract: what `ProviderConfig` writes must
- /// be readable through `AppGroupStore` (and vice-versa) on the same
- /// suite, and `mode == .off` short-circuits before any network call.
+ /// be readable through `AppGroupStore` on the same suite.
func testAppGroupCrossProcessAndOffModeShortCircuit() async {
let suiteName = "group.com.osgkeyboard.shared.tests.\(UUID().uuidString)"
let defaults = UserDefaults(suiteName: suiteName)!
@@ -253,14 +252,6 @@ final class LLMClientTests: XCTestCase {
XCTAssertEqual(store.apiKey, "sk-test-1234", "API key did not survive the cross-process boundary")
XCTAssertEqual(store.modeId, "off")
XCTAssertEqual(store.model, "gpt-4o-mini")
-
- // mode == .off must short-circuit (the keyboard extension never
- // even calls `polisher.polish` in this mode, so no LLMClient is
- // constructed and no network request happens). We model the
- // short-circuit on the read side: the persisted mode is "off" and
- // any upstream caller checking `state.mode == .off` would skip
- // the LLM. The guarantee is the persistence + the literal value.
- XCTAssertEqual(store.modeId, "off")
}
func testAppGroupStoreNoAPIKeySurfacesAsLLMError() async {
@@ -288,30 +279,25 @@ final class LLMClientTests: XCTestCase {
}
}
- // MARK: - TEST-2: mode = .off short-circuits PolishingService
+ // MARK: - TEST-2: cloud always polishes (legacy modeId ignored)
- /// `PolishingService.polish()` must not invoke the underlying
- /// `LLMClient` when the App Group store reports `modeId == "off"`.
- /// We verify both halves of that contract:
- /// 1. The return value is the trimmed input (not a polished round-trip).
- /// 2. The `LLMClient` is never asked to talk to the network.
- func testPolisherSkipsNetworkWhenModeOff() async throws {
+ /// Cloud engine must invoke the LLM even when a legacy `modeId == "off"`
+ /// value is still present in the App Group suite.
+ func testPolisherPolishesWhenCloudEvenIfModeOffLegacy() async throws {
let suiteName = "group.com.osgkeyboard.shared.tests.\(UUID().uuidString)"
let defaults = UserDefaults(suiteName: suiteName)!
defaults.removePersistentDomain(forName: suiteName)
defer { defaults.removePersistentDomain(forName: suiteName) }
- // modeId = "off" — this is the switch we care about.
defaults.set("off", forKey: "config.modeId")
+ defaults.set("cloud", forKey: "config.engineMode")
defaults.set("https://example.com/v1", forKey: "config.baseURL")
- defaults.set("sk-should-not-be-used", forKey: "config.apiKey")
+ defaults.set("sk-test", forKey: "config.apiKey")
defaults.set("gpt-4o-mini", forKey: "config.model")
- // Counter LLMClient: if `polish()` is ever called, this trips.
let counter = CallCounter()
- let countingClient = CountingLLMClient(counter: counter) { _, _ in
- XCTFail("LLMClient.polish was invoked under mode=off — short-circuit failed")
- return ""
+ let countingClient = CountingLLMClient(counter: counter) { raw, _ in
+ "POLISHED: \(raw)"
}
let store = AppGroupStore(defaults: defaults)
@@ -322,9 +308,38 @@ final class LLMClientTests: XCTestCase {
)
let result = try await polisher.polish(" hello world ")
- XCTAssertEqual(result, "hello world", "mode=off must return trimmed input, not polished output")
+ XCTAssertEqual(result, "POLISHED: hello world")
let calls = await counter.value()
- XCTAssertEqual(calls, 0, "LLMClient.polish must not be called when modeId == \"off\"")
+ XCTAssertEqual(calls, 1, "cloud engine must polish even with legacy modeId=off")
+ }
+
+ /// Local engine is ASR-only and never calls the cloud `LLMClient`.
+ func testPolisherReturnsRawWhenEngineLocal() async throws {
+ let suiteName = "group.com.osgkeyboard.shared.tests.\(UUID().uuidString)"
+ let defaults = UserDefaults(suiteName: suiteName)!
+ defaults.removePersistentDomain(forName: suiteName)
+ defer { defaults.removePersistentDomain(forName: suiteName) }
+
+ defaults.set("local", forKey: "config.engineMode")
+ defaults.set("off", forKey: "config.modeId")
+
+ let counter = CallCounter()
+ let countingClient = CountingLLMClient(counter: counter) { _, _ in
+ XCTFail("cloud LLMClient must not run under local engine")
+ return ""
+ }
+
+ let store = AppGroupStore(defaults: defaults)
+ let polisher = PolishingService(
+ store: store,
+ client: countingClient,
+ timeout: 1
+ )
+
+ let result = try await polisher.polish(" hello ")
+ XCTAssertEqual(result, "hello")
+ let calls = await counter.value()
+ XCTAssertEqual(calls, 0)
}
}
diff --git a/OSGKeyboardTests/ProgressiveDictationTranscriptAccumulatorTests.swift b/OSGKeyboardTests/ProgressiveDictationTranscriptAccumulatorTests.swift
new file mode 100644
index 0000000..bf8bc60
--- /dev/null
+++ b/OSGKeyboardTests/ProgressiveDictationTranscriptAccumulatorTests.swift
@@ -0,0 +1,44 @@
+// ProgressiveDictationTranscriptAccumulatorTests.swift
+// OSGKeyboardTests
+
+import CoreMedia
+import XCTest
+@testable import OSGKeyboardShared
+
+final class ProgressiveDictationTranscriptAccumulatorTests: XCTestCase {
+
+ private func range(start: Double, duration: Double = 30) -> CMTimeRange {
+ CMTimeRange(
+ start: CMTime(seconds: start, preferredTimescale: 600),
+ duration: CMTime(seconds: duration, preferredTimescale: 600)
+ )
+ }
+
+ func testCumulativePartialsWithinSameRangeUpdateSegment() {
+ var acc = ProgressiveDictationTranscriptAccumulator()
+ let r0 = range(start: 0)
+
+ XCTAssertEqual(acc.ingest(range: r0, text: "今天天气"), "今天天气")
+ XCTAssertEqual(acc.ingest(range: r0, text: "今天天气很好"), "今天天气很好")
+ XCTAssertEqual(acc.finalize(), "今天天气很好")
+ }
+
+ func testNewRangeAppendsInsteadOfReplacingEarlierSpeech() {
+ var acc = ProgressiveDictationTranscriptAccumulator()
+ let r0 = range(start: 0)
+ let r30 = range(start: 30)
+
+ _ = acc.ingest(range: r0, text: "前三十秒的内容")
+ _ = acc.ingest(range: r30, text: "后二十秒的内容")
+
+ XCTAssertEqual(acc.finalize(), "前三十秒的内容 后二十秒的内容")
+ }
+
+ func testDuplicateEmissionIsSuppressed() {
+ var acc = ProgressiveDictationTranscriptAccumulator()
+ let r0 = range(start: 0)
+
+ XCTAssertNotNil(acc.ingest(range: r0, text: "hello"))
+ XCTAssertNil(acc.ingest(range: r0, text: "hello"))
+ }
+}
diff --git a/OSGKeyboardTests/UtteranceStreamChunkerTests.swift b/OSGKeyboardTests/UtteranceStreamChunkerTests.swift
new file mode 100644
index 0000000..b1de506
--- /dev/null
+++ b/OSGKeyboardTests/UtteranceStreamChunkerTests.swift
@@ -0,0 +1,42 @@
+// UtteranceStreamChunkerTests.swift
+// OSGKeyboardTests
+
+import XCTest
+@testable import OSGKeyboardShared
+
+final class UtteranceStreamChunkerTests: XCTestCase {
+
+ private let config = FlowUtteranceChunkConfig(
+ maxChunkDurationSeconds: 1,
+ overlapDurationSeconds: 0.1,
+ pauseExtensionMaxSeconds: 0.2,
+ pauseRMSThreshold: 0.02,
+ sampleRate: 1_000
+ )
+
+ func testPauseAwareSplitPrefersSilenceNearWindowEnd() {
+ var buffer = [Float](repeating: 0.2, count: 900)
+ buffer.append(contentsOf: [Float](repeating: 0.001, count: 50))
+ buffer.append(contentsOf: [Float](repeating: 0.2, count: 100))
+
+ let split = UtteranceStreamChunker.pauseAwareSplitIndex(in: buffer, config: config)
+ XCTAssertGreaterThanOrEqual(split, config.maxChunkSamples)
+ XCTAssertLessThanOrEqual(split, config.maxChunkSamples + config.pauseExtensionSamples)
+ }
+
+ func testChunksEmitMultipleSegmentsForLongStream() async {
+ let sampleCount = config.maxChunkSamples * 2 + 100
+ let samples = [Float](repeating: 0.05, count: sampleCount)
+ let (stream, continuation) = AsyncStream.makeStream()
+ continuation.yield(AudioBufferSnapshot(samples: samples, sampleRate: Double(config.sampleRate)))
+ continuation.finish()
+
+ var received: [UtteranceAudioChunk] = []
+ for await chunk in UtteranceStreamChunker.chunks(from: stream, config: config) {
+ received.append(chunk)
+ }
+
+ XCTAssertGreaterThanOrEqual(received.count, 2)
+ XCTAssertTrue(received.last?.isLast == true)
+ }
+}
diff --git a/OSGKeyboardTests/UtteranceTranscriptStitcherTests.swift b/OSGKeyboardTests/UtteranceTranscriptStitcherTests.swift
new file mode 100644
index 0000000..85ab076
--- /dev/null
+++ b/OSGKeyboardTests/UtteranceTranscriptStitcherTests.swift
@@ -0,0 +1,23 @@
+// UtteranceTranscriptStitcherTests.swift
+// OSGKeyboardTests
+
+import XCTest
+@testable import OSGKeyboardShared
+
+final class UtteranceTranscriptStitcherTests: XCTestCase {
+
+ func testMergeWithOverlapRemovesDuplicatedSuffixPrefix() {
+ let merged = UtteranceTranscriptStitcher.mergeWithOverlap(
+ previous: "今天天气很好",
+ next: "很好我们继续"
+ )
+ XCTAssertEqual(merged, "今天天气很好我们继续")
+ }
+
+ func testStitcherOrdersChunksByIndex() {
+ var stitcher = UtteranceTranscriptStitcher()
+ stitcher.append(index: 1, text: "第二段")
+ stitcher.append(index: 0, text: "第一段")
+ XCTAssertEqual(stitcher.composed(), "第一段 第二段")
+ }
+}
diff --git a/Scripts/generate-xcodeproj.sh b/Scripts/generate-xcodeproj.sh
new file mode 100755
index 0000000..6806720
--- /dev/null
+++ b/Scripts/generate-xcodeproj.sh
@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+# Generate OSGKeyboard.xcodeproj from project.yml and apply the Icon Composer
+# patch required by XcodeGen 2.43.
+set -euo pipefail
+
+ROOT="$(cd "$(dirname "$0")/.." && pwd)"
+cd "$ROOT"
+
+xcodegen generate
+"$ROOT/Scripts/patch-icon-composer.sh"
diff --git a/Scripts/patch-icon-composer.sh b/Scripts/patch-icon-composer.sh
new file mode 100755
index 0000000..5eb17a1
--- /dev/null
+++ b/Scripts/patch-icon-composer.sh
@@ -0,0 +1,138 @@
+#!/usr/bin/env bash
+# XcodeGen 2.43 expands AppIcon.icon into a PBXGroup and adds icon.json / svg
+# files to Copy Bundle Resources. Icon Composer bundles must be a single
+# PBXFileReference (folder.iconcomposer.icon) linked to the target so actool
+# compiles them together with Assets.xcassets.
+set -euo pipefail
+
+ROOT="$(cd "$(dirname "$0")/.." && pwd)"
+PBXPROJ="$ROOT/OSGKeyboard.xcodeproj/project.pbxproj"
+ICON_PATH="$ROOT/OSGKeyboard/AppIcon.icon"
+
+if [[ ! -f "$PBXPROJ" ]]; then
+ echo "error: $PBXPROJ not found (run xcodegen generate first)" >&2
+ exit 1
+fi
+
+if [[ ! -d "$ICON_PATH" ]]; then
+ echo "error: $ICON_PATH not found" >&2
+ exit 1
+fi
+
+python3 - "$PBXPROJ" <<'PY'
+import re
+import sys
+import uuid
+from pathlib import Path
+
+pbxproj = Path(sys.argv[1])
+text = pbxproj.read_text()
+
+if "/* AppIcon.icon in Resources */" in text and "folder.iconcomposer.icon" in text:
+ print("AppIcon.icon already patched")
+ sys.exit(0)
+
+icon_uuid_match = re.search(
+ r"([A-F0-9]{24}) /\* AppIcon\.icon \*/ = \{",
+ text,
+)
+if not icon_uuid_match:
+ print("error: AppIcon.icon not found in project.pbxproj", file=sys.stderr)
+ sys.exit(1)
+icon_uuid = icon_uuid_match.group(1)
+
+group_pattern = re.compile(
+ rf"(?P\t\t){icon_uuid} /\* AppIcon\.icon \*/ = \{{\n"
+ r"\t\t\tisa = PBXGroup;\n"
+ r"\t\t\tchildren = \(\n"
+ r"(?:\t\t\t\t[A-F0-9]{24} /\* .* \*/,\n)*"
+ r"\t\t\t\);\n"
+ r"\t\t\tpath = AppIcon\.icon;\n"
+ r"\t\t\tsourceTree = \"\";\n"
+ r"\t\t\};",
+ re.MULTILINE,
+)
+
+group_match = group_pattern.search(text)
+if group_match:
+ indent = group_match.group("indent")
+ replacement = (
+ f"{indent}{icon_uuid} /* AppIcon.icon */ = {{\n"
+ f"\t\t\tisa = PBXFileReference;\n"
+ f"\t\t\tlastKnownFileType = folder.iconcomposer.icon;\n"
+ f"\t\t\tpath = AppIcon.icon;\n"
+ f"\t\t\tsourceTree = \"\";\n"
+ f"\t\t}};"
+ )
+ text = text[: group_match.start()] + replacement + text[group_match.end() :]
+else:
+ file_ref_pattern = re.compile(
+ rf"\t\t{icon_uuid} /\* AppIcon\.icon \*/ = \{{\n"
+ r"\t\t\tisa = PBXFileReference;\n"
+ r"\t\t\tlastKnownFileType = [^;]+;\n"
+ r"\t\t\tpath = AppIcon\.icon;\n"
+ r"\t\t\tsourceTree = \"\";\n"
+ r"\t\t\};",
+ re.MULTILINE,
+ )
+ file_ref_match = file_ref_pattern.search(text)
+ if not file_ref_match:
+ print("error: AppIcon.icon entry has unexpected shape", file=sys.stderr)
+ sys.exit(1)
+ replacement = (
+ f"\t\t{icon_uuid} /* AppIcon.icon */ = {{\n"
+ f"\t\t\tisa = PBXFileReference;\n"
+ f"\t\t\tlastKnownFileType = folder.iconcomposer.icon;\n"
+ f"\t\t\tpath = AppIcon.icon;\n"
+ f"\t\t\tsourceTree = \"\";\n"
+ f"\t\t}};"
+ )
+ text = text[: file_ref_match.start()] + replacement + text[file_ref_match.end() :]
+
+nested_resource_names = ("icon.json", "App Icon Template.svg")
+lines = text.splitlines(keepends=True)
+filtered: list[str] = []
+for line in lines:
+ if " in Resources */ = {isa = PBXBuildFile;" in line and any(
+ name in line for name in nested_resource_names
+ ):
+ continue
+ if " in Resources */," in line and any(name in line for name in nested_resource_names):
+ continue
+ filtered.append(line)
+text = "".join(filtered)
+
+build_uuid = uuid.uuid4().hex[:24].upper()
+build_entry = (
+ f"\t\t{build_uuid} /* AppIcon.icon in Resources */ = "
+ f"{{isa = PBXBuildFile; fileRef = {icon_uuid} /* AppIcon.icon */; }};\n"
+)
+text = text.replace("/* Begin PBXBuildFile section */\n", "/* Begin PBXBuildFile section */\n" + build_entry, 1)
+
+resources_phase = re.search(
+ r"\t\t(?P[A-F0-9]{24}) /\* Resources \*/ = \{\n"
+ r"\t\t\tisa = PBXResourcesBuildPhase;\n"
+ r"\t\t\tbuildActionMask = 2147483647;\n"
+ r"\t\t\tfiles = \(\n"
+ r"(?P.*?Assets\.xcassets in Resources.*?\n)"
+ r"(?P.*?)"
+ r"\t\t\t\);\n"
+ r"\t\t\trunOnlyForDeploymentPostprocessing = 0;\n"
+ r"\t\t\};",
+ text,
+ re.DOTALL,
+)
+if not resources_phase:
+ print("error: OSGKeyboard Resources phase not found", file=sys.stderr)
+ sys.exit(1)
+
+insert_at = resources_phase.end("body")
+text = (
+ text[:insert_at]
+ + f"\t\t\t\t{build_uuid} /* AppIcon.icon in Resources */,\n"
+ + text[insert_at:]
+)
+
+pbxproj.write_text(text)
+print("Patched AppIcon.icon -> folder.iconcomposer.icon (target Resources)")
+PY
diff --git a/project.yml b/project.yml
index 8d8c6f0..7affb80 100644
--- a/project.yml
+++ b/project.yml
@@ -1,10 +1,16 @@
# XcodeGen configuration for OSGKeyboard
-# Run `xcodegen generate` to create OSGKeyboard.xcodeproj
+# Run `./Scripts/generate-xcodeproj.sh` to create OSGKeyboard.xcodeproj.
+# App home-screen icon: OSGKeyboard/AppIcon.icon (Icon Composer, iOS 26).
# Repo only tracks project.yml; .xcodeproj is gitignored.
name: OSGKeyboard
options:
bundleIdPrefix: com.osgkeyboard
+ # Xcode 26 Icon Composer bundles must be treated as a single file,
+ # not expanded into icon.json + SVG children (see XcodeGen #1556).
+ fileTypes:
+ icon:
+ file: true
# Minimum OS: iOS 26. We dropped older iOS support so the
# legacy speech + AVAudioSession branching could be removed in
# favour of iOS 26's `SpeechAnalyzer` (always on-device) and
@@ -17,6 +23,23 @@ options:
generateEmptyDirectories: true
groupSortPosition: top
+# Local package: `Qwen3Speech`. Vendored fork of
+# `soniqo/speech-swift`'s Qwen3-ASR + Qwen3-Chat modules (plus the
+# AudioCommon / MLXCommon / SpeechVAD slices they depend on).
+# Upstream's `Package.swift` references a `CSpeechCore` binary
+# target whose URL doesn't match its declared filename (the file
+# is `SpeechCore.xcframework.zip`, the target is `CSpeechCore`),
+# which breaks SwiftPM resolve on a clean checkout. Vendoring
+# lets us ship the four targets we actually link without waiting
+# for the upstream fix; the `audio` / `audio-server` CLIs and the
+# `AudioServer` HTTP backend that needed `CSpeechCore` aren't part
+# of OSGKeyboard's build graph. Sources carry their original
+# Apache-2.0 copyright headers; the `Package.swift` here adds
+# only what we need.
+packages:
+ Qwen3Speech:
+ path: OSGKeyboard/ThirdParty/Qwen3Speech
+
settings:
base:
SWIFT_VERSION: "6.0"
@@ -25,8 +48,9 @@ settings:
SWIFT_STRICT_CONCURRENCY: complete
GENERATE_INFOPLIST_FILE: NO
ENABLE_MODULE_VERIFIER: YES
- MARKETING_VERSION: "0.1.2"
- CURRENT_PROJECT_VERSION: "3"
+ CLANG_CXX_LANGUAGE_STANDARD: c++17
+ MARKETING_VERSION: "0.2.0"
+ CURRENT_PROJECT_VERSION: "4"
# 签名配置来自 Signing.local.xcconfig(gitignored,不会被覆盖)
# 项目级签名 xcconfig,适用于所有 target
@@ -43,10 +67,28 @@ targets:
type: application
platform: iOS
sources:
+ - path: OSGKeyboard/AppIcon.icon
+ type: file
+ buildPhase: resources
- path: OSGKeyboard
excludes:
- "AppIcon.icon"
- "OSGKeyboard.icon"
+ # Legacy PNG app icons must not coexist with AppIcon.icon — actool
+ # crashes when both are passed. iOS 26 uses Icon Composer only.
+ - "Assets.xcassets/AppIcon.appiconset"
+ # `ThirdParty/Qwen3Speech` is a local SPM package pulled in
+ # by `packages:` below — it builds as its own targets
+ # (`Qwen3ASR`, …) and we link the products.
+ # If we don't exclude the source folder, xcodegen walks
+ # into the package's `Sources/*/` directories and pulls
+ # the swift files straight into the host app target,
+ # which collides on duplicate type names (e.g. two
+ # `Configuration.swift` files in `Qwen3ASR` and
+ # `SpeechVAD`) and on multiple `@main` / top-level
+ # declarations. Keep the folder as a package; don't
+ # flatten it into the app.
+ - "ThirdParty/Qwen3Speech"
entitlements:
path: OSGKeyboard/OSGKeyboard.entitlements
properties:
@@ -65,6 +107,7 @@ targets:
- path: OSGKeyboard/en.lproj
- path: OSGKeyboard/zh-Hans.lproj
- path: OSGKeyboard/Resources/Fonts/MaterialIcons-Regular.ttf
+ - path: OSGKeyboard/Resources/PrivacyPolicy.html
- path: OSGKeyboard/PrivacyInfo.xcprivacy
info:
path: OSGKeyboard/Info.plist
@@ -106,14 +149,20 @@ targets:
SUPPORTS_MACCATALYST: NO
SUPPORTS_MAC_DESIGNED_FOR_IPHONE_IPAD: NO
SUPPORTS_XR_DESIGNED_FOR_IPHONE_IPAD: NO
- # Use AppIcon.appiconset only. AppIcon.icon (Icon Composer) must not
- # be passed to actool alongside the asset catalog — it crashes actool.
+ ASSETCATALOG_COMPILER_APPICON_NAME: AppIcon
dependencies:
- target: OSGKeyboardShared
embed: true
- target: OSGKeyboardExt
# Keyboard Extension is a plugin of the main App; embed it.
- sdk: Speech.framework
+ # Local fork of soniqo/speech-swift. See the `packages:`
+ # block at the top of this file for the why-fork rationale.
+ # Only the two products we actually link are listed; AudioServer
+ # / AudioCLI (which is what the upstream binary target was
+ # gating) aren't in our build graph.
+ - package: Qwen3Speech
+ product: Qwen3ASR
# =========================================================
# Keyboard Extension
@@ -184,6 +233,13 @@ targets:
platform: iOS
sources:
- path: OSGKeyboardShared
+ excludes:
+ - "en.lproj"
+ - "zh-Hans.lproj"
+ - path: OSGKeyboardShared/en.lproj/Shared.strings
+ buildPhase: resources
+ - path: OSGKeyboardShared/zh-Hans.lproj/Shared.strings
+ buildPhase: resources
info:
path: OSGKeyboardShared/Info.plist
settings: