feat: migrate on-device Qwen3 ASR to CoreML for background Flow dictation
Replace MLX GPU inference with CoreML bundles so transcription continues while the host app is backgrounded. Adds model download and warm-up, vendored Qwen3Speech, and updates onboarding, settings, and copy for the ~1.6 GB CoreML package (iOS 18+).
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>OSGKeyboard Privacy Policy</title>
|
||||
<style>
|
||||
body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; line-height: 1.6; max-width: 720px; margin: 0 auto; padding: 1rem; color: #1a1a1a; background: transparent; }
|
||||
h1, h2 { line-height: 1.3; }
|
||||
a { color: #0a7; }
|
||||
hr { margin: 2rem 0; border: none; border-top: 1px solid #ddd; }
|
||||
.lang { font-size: 0.9rem; color: #666; }
|
||||
@media (prefers-color-scheme: dark) {
|
||||
body { color: #f2f2f7; }
|
||||
hr { border-top-color: #3a3a3c; }
|
||||
.lang { color: #98989d; }
|
||||
a { color: #64d2ff; }
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<p class="lang"><a href="#zh">中文</a> · <a href="#top">English</a></p>
|
||||
<div id="top">
|
||||
<h1>OSGKeyboard Privacy Policy</h1>
|
||||
<p><strong>Last updated:</strong> June 19, 2026</p>
|
||||
<p>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.</p>
|
||||
|
||||
<h2>What we collect</h2>
|
||||
<ul>
|
||||
<li><strong>Voice audio</strong> — captured only while you actively record. On-device mode transcribes locally with Apple’s speech APIs; raw audio is not uploaded by OSGKeyboard.</li>
|
||||
<li><strong>Transcribed text</strong> — 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.</li>
|
||||
<li><strong>API credentials</strong> — stored in the iOS Keychain on your device and shared only between the main app and keyboard extension via an App Group.</li>
|
||||
<li><strong>App preferences</strong> — engine mode, language, and keyboard settings stored in App Group UserDefaults on your device.</li>
|
||||
</ul>
|
||||
|
||||
<h2>What we do not collect</h2>
|
||||
<ul>
|
||||
<li>We do <strong>not</strong> log or upload ordinary keystrokes you type with the keyboard.</li>
|
||||
<li>We do <strong>not</strong> operate analytics or advertising SDKs.</li>
|
||||
<li>We do <strong>not</strong> sell personal data.</li>
|
||||
</ul>
|
||||
|
||||
<h2>Permissions</h2>
|
||||
<ul>
|
||||
<li><strong>Microphone</strong> — required for voice input and background voice sessions.</li>
|
||||
<li><strong>Speech recognition</strong> — required for on-device transcription.</li>
|
||||
<li><strong>Full Access</strong> — 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.</li>
|
||||
</ul>
|
||||
|
||||
<h2>Third parties</h2>
|
||||
<p>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.</p>
|
||||
|
||||
<h2>Data retention</h2>
|
||||
<p>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.</p>
|
||||
<p><strong>Voice history</strong> — 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.</p>
|
||||
|
||||
<h2>Contact</h2>
|
||||
<p>Questions: open an issue at <a href="https://github.com/hkgood/OSGKeyboard">github.com/hkgood/OSGKeyboard</a>.</p>
|
||||
</div>
|
||||
|
||||
<hr id="zh">
|
||||
<h1>OSGKeyboard 隐私政策</h1>
|
||||
<p><strong>更新日期:</strong>2026 年 6 月 19 日</p>
|
||||
<p>OSGKeyboard 是一款 iOS 自定义键盘,可将语音转为文字。本政策说明应用处理哪些数据及用途。</p>
|
||||
|
||||
<h2>我们处理的数据</h2>
|
||||
<ul>
|
||||
<li><strong>语音音频</strong> — 仅在你主动录音时采集。本地模式在设备端通过 Apple 语音识别转写,OSGKeyboard 不会上传原始录音。</li>
|
||||
<li><strong>转写文字</strong> — 云端润色模式下,最终文字(非音频)可能发送到你配置的 LLM 服务商以整理标点和格式。</li>
|
||||
<li><strong>API 凭证</strong> — 保存在设备 Keychain,仅通过 App Group 在主 App 与键盘扩展间共享。</li>
|
||||
<li><strong>应用偏好</strong> — 引擎、语言等设置保存在设备 App Group 中。</li>
|
||||
</ul>
|
||||
|
||||
<h2>我们不收集的内容</h2>
|
||||
<ul>
|
||||
<li>我们<strong>不会</strong>记录或上传你平时在键盘上的击键内容。</li>
|
||||
<li>我们<strong>不会</strong>集成广告或第三方分析 SDK。</li>
|
||||
<li>我们<strong>不会</strong>出售个人数据。</li>
|
||||
</ul>
|
||||
|
||||
<h2>权限说明</h2>
|
||||
<ul>
|
||||
<li><strong>麦克风</strong> — 语音输入与后台语音会话所需。</li>
|
||||
<li><strong>语音识别</strong> — 端侧转写所需。</li>
|
||||
<li><strong>完全访问</strong> — 使键盘能使用麦克风、读取 API Key 并与主 App 通信。完全访问不代表我们会收集全部击键内容。</li>
|
||||
</ul>
|
||||
|
||||
<h2>第三方</h2>
|
||||
<p>选择云端润色时,转写文字会发往你配置的 API,该服务商的隐私政策适用于相关请求。</p>
|
||||
|
||||
<h2>数据保留</h2>
|
||||
<p>设置与 API Key 保留在设备上,直至卸载或重置。识别结果写入你正在使用的宿主 App,OSGKeyboard 不会长期存储。</p>
|
||||
<p><strong>语音历史</strong> — 成功的转写可能保存在主 App「历史」页,仅供本机查看,不会上传,可随时在历史页清空或通过重置设置清除。</p>
|
||||
|
||||
<h2>联系</h2>
|
||||
<p>问题反馈:<a href="https://github.com/hkgood/OSGKeyboard">github.com/hkgood/OSGKeyboard</a></p>
|
||||
</body>
|
||||
</html>
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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<Void, Never>?
|
||||
private var heartbeatTask: Task<Void, Never>?
|
||||
@@ -33,10 +48,13 @@ final class FlowSessionManager: ObservableObject {
|
||||
private var isUtteranceProcessing = false
|
||||
private var finalizeTask: Task<Void, Never>?
|
||||
private var asrTask: Task<Void, Never>?
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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/")
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<OnDeviceModel> = []
|
||||
|
||||
private var downloadTasks: [OnDeviceModel: Task<Void, Never>] = [:]
|
||||
|
||||
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: <reason>" 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!
|
||||
}
|
||||
}
|
||||
@@ -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<Void, Never>?
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
"""
|
||||
}
|
||||
@@ -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<Void, Never>?
|
||||
private var cancelled = false
|
||||
|
||||
private var model: CoreMLASRModel?
|
||||
private var loadError: Error?
|
||||
private var loadingTask: Task<CoreMLASRModel, Error>?
|
||||
|
||||
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<CoreMLASRModel, Error> = lock.withLock {
|
||||
if let existing = loadingTask { return existing }
|
||||
let new = Task<CoreMLASRModel, Error> { [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<AudioBufferSnapshot>,
|
||||
locale: Locale
|
||||
) -> AsyncStream<ASREvent> {
|
||||
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..<end])
|
||||
offset = end
|
||||
chunkIndex += 1
|
||||
FlowDiagnostics.log(
|
||||
"Qwen3ASR CoreML chunk \(chunkIndex) samples=\(chunk.count) " +
|
||||
"offset=\(offset - chunk.count)"
|
||||
)
|
||||
let piece = model.transcribeBackgroundSafe(
|
||||
audio: chunk,
|
||||
sampleRate: TranscribeConstants.sampleRate,
|
||||
language: language
|
||||
)
|
||||
.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
if !piece.isEmpty, !piece.hasPrefix("[CoreML error:") {
|
||||
parts.append(piece)
|
||||
}
|
||||
}
|
||||
return parts.joined(separator: " ")
|
||||
}
|
||||
|
||||
private static func languageHint(from locale: Locale) -> 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")
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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/<Target>/` 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"),
|
||||
]
|
||||
),
|
||||
]
|
||||
)
|
||||
+382
@@ -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..<sampleCount {
|
||||
// Take first channel only
|
||||
let sampleIndex = i * channels
|
||||
if sampleIndex < int16Ptr.count {
|
||||
samples[i] = Float(int16Ptr[sampleIndex]) / 32768.0
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return (samples, Int(sampleRate))
|
||||
}
|
||||
|
||||
/// Resample mono audio using `AVAudioConverter`.
|
||||
///
|
||||
/// Fully drains the converter via `.endOfStream` so the filter tail isn't
|
||||
/// truncated, and normalizes the result to the exact expected frame count.
|
||||
/// `quality` selects the SRC filter: `.standard` (default) for speech,
|
||||
/// `.mastering` for music/upsampling — see ``ResampleQuality``.
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - samples: mono PCM Float32 audio
|
||||
/// - inputRate: source sample rate in Hz
|
||||
/// - outputRate: target sample rate in Hz
|
||||
/// - quality: SRC filter quality (default `.standard`)
|
||||
/// - Returns: resampled audio at `outputRate`. On converter-setup or
|
||||
/// conversion failure, returns the original `samples` unchanged (callers
|
||||
/// never receive partial/truncated output).
|
||||
public static func resample(_ samples: [Float], from inputRate: Int, to outputRate: Int, quality: ResampleQuality = .standard) -> [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..<channels { out[c].reserveCapacity(Int(capacity)) }
|
||||
|
||||
var fed = false
|
||||
while true {
|
||||
target.frameLength = 0
|
||||
var error: NSError?
|
||||
let status = converter.convert(to: target, error: &error) { _, outStatus in
|
||||
if fed {
|
||||
// Signal end-of-stream so the converter flushes its sinc
|
||||
// tail instead of stopping short.
|
||||
outStatus.pointee = .endOfStream
|
||||
return nil
|
||||
}
|
||||
fed = true
|
||||
outStatus.pointee = .haveData
|
||||
return source
|
||||
}
|
||||
|
||||
let produced = Int(target.frameLength)
|
||||
if produced > 0, let chans = target.floatChannelData {
|
||||
for c in 0..<channels {
|
||||
out[c].append(contentsOf: UnsafeBufferPointer(start: chans[c], count: produced))
|
||||
}
|
||||
}
|
||||
|
||||
// Order matters: a clean `.endOfStream` (success) usually arrives on
|
||||
// a call that produces 0 frames, so check it before the no-progress
|
||||
// guard. `.inputRanDry` is not terminal — keep pulling until the
|
||||
// input block signals `.endOfStream`.
|
||||
if status == .error || error != nil { return nil }
|
||||
if status == .endOfStream { break }
|
||||
if produced == 0 { return nil }
|
||||
}
|
||||
|
||||
// Normalize to the exact expected length so output is a strict function
|
||||
// of input length and ratio (downstream stereo alignment and A/V sync
|
||||
// rely on this). Mastering SRC compensates filter latency, so this is a
|
||||
// 0–1 sample trim/pad in practice.
|
||||
let expected = Int((Double(inputFrames) * ratio).rounded())
|
||||
guard expected > 0, !out[0].isEmpty else { return nil }
|
||||
for c in 0..<channels {
|
||||
if out[c].count > 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)"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<Float>.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
|
||||
@@ -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)'"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<Float>, count sampleCount: Int) {
|
||||
os_unfair_lock_lock(&_lock)
|
||||
defer { os_unfair_lock_unlock(&_lock) }
|
||||
for i in 0..<sampleCount {
|
||||
if count == capacity {
|
||||
readPos = (readPos + 1) % capacity
|
||||
count -= 1
|
||||
}
|
||||
buffer[writePos] = pointer[i]
|
||||
writePos = (writePos + 1) % capacity
|
||||
count += 1
|
||||
}
|
||||
}
|
||||
|
||||
/// Called from MLX inference thread — returns zeros on underrun; never blocks.
|
||||
public func read(_ n: Int) -> [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..<available {
|
||||
result[i] = buffer[(readPos + i) % capacity]
|
||||
}
|
||||
readPos = (readPos + available) % capacity
|
||||
count -= available
|
||||
// Remaining positions stay as zero (underrun padding)
|
||||
return result
|
||||
}
|
||||
|
||||
/// Number of samples currently available to read.
|
||||
public var available: Int {
|
||||
os_unfair_lock_lock(&_lock)
|
||||
defer { os_unfair_lock_unlock(&_lock) }
|
||||
return count
|
||||
}
|
||||
}
|
||||
+49
@@ -0,0 +1,49 @@
|
||||
#if canImport(CoreML)
|
||||
import CoreML
|
||||
import Foundation
|
||||
|
||||
/// Resolves the `MLComputeUnits` a CoreML model should load with, honoring
|
||||
/// the `SPEECH_COREML_COMPUTE_UNITS` environment override.
|
||||
///
|
||||
/// **Why this exists.** A `.mlmodelc` is compiled MIL, but the device-specific
|
||||
/// program (ANE *or* GPU/Metal) is generated the *first time* the model loads.
|
||||
/// On real M-series hardware that first-load codegen takes seconds; on a
|
||||
/// virtualized GitHub `macos-15` runner it **hangs**: the runner has no usable
|
||||
/// Neural Engine (so `.cpuAndNeuralEngine`/`.all` stall attempting the ANE
|
||||
/// compile) AND its paravirtual GPU can't JIT CoreML's Metal program for a
|
||||
/// stateful graph (so `.cpuAndGPU` stalls too). Both were observed as 17-27 min
|
||||
/// silent hangs loading the Qwen3-ASR T=128 stateful decoder.
|
||||
///
|
||||
/// On-device we keep the normal default (callers pass it as `fallback`; the env
|
||||
/// is unset so this is a no-op). In CI we set `SPEECH_COREML_COMPUTE_UNITS=cpuOnly`
|
||||
/// so every loader skips ANE *and* GPU codegen — pure CPU MIL execution loads
|
||||
/// instantly, is deterministic, and yields identical text for our roundtrip
|
||||
/// assertions (~86 ms/step for the T=128 decoder, fine for correctness tests).
|
||||
public enum CoreMLComputeUnitsResolver {
|
||||
public static let envKey = "SPEECH_COREML_COMPUTE_UNITS"
|
||||
|
||||
/// Returns the env-overridden compute units, or `fallback` when unset/unrecognized.
|
||||
/// Accepted env values (case-insensitive): `ane`/`cpuAndNeuralEngine`,
|
||||
/// `gpu`/`cpuAndGPU`, `cpu`/`cpuOnly`, `all`.
|
||||
public static func resolved(default fallback: MLComputeUnits) -> 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
|
||||
@@ -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<String>()
|
||||
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() }
|
||||
}
|
||||
}
|
||||
+428
@@ -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<String> = [
|
||||
"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/<org>/<model>` 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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
+335
@@ -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<String>(["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
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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<AudioChunk, Error>
|
||||
}
|
||||
|
||||
extension SpeechGenerationModel {
|
||||
/// Default: wraps `generate()` as a single-chunk stream.
|
||||
public func generateStream(text: String, language: String?) -> AsyncThrowingStream<AudioChunk, Error> {
|
||||
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<AudioChunk, Error>
|
||||
}
|
||||
|
||||
// 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]
|
||||
}
|
||||
+182
@@ -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)"
|
||||
}
|
||||
}
|
||||
}
|
||||
+511
@@ -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<Float>
|
||||
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<Float>.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<Float>, 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..<fadeFrames {
|
||||
output[i] *= Float(i) / Float(fadeFrames)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
lock.lock()
|
||||
ringBuffer?.write(output)
|
||||
totalWritten += output.count
|
||||
isPlaying = true
|
||||
|
||||
if !playbackStarted && preBufferSamples > 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..<end]
|
||||
|
||||
// Compute frame RMS
|
||||
var sumSq: Float = 0
|
||||
for s in frame { sumSq += s * s }
|
||||
let rms = sqrt(sumSq / Float(frame.count))
|
||||
|
||||
if rms < threshold {
|
||||
silenceRun += frame.count
|
||||
// Only keep silence up to maxSilence
|
||||
if silenceRun <= maxSilence {
|
||||
result.append(contentsOf: frame)
|
||||
}
|
||||
// Else: drop this frame (compress the silence)
|
||||
} else {
|
||||
silenceRun = 0
|
||||
result.append(contentsOf: frame)
|
||||
}
|
||||
offset = end
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,335 @@
|
||||
import Foundation
|
||||
|
||||
// Library logs route to stderr so they don't corrupt a stdout-based IPC
|
||||
// channel (e.g. the speech-studio sidecar's NDJSON protocol).
|
||||
@inline(__always)
|
||||
private func logTokenizer(_ message: String) {
|
||||
FileHandle.standardError.write(Data((message + "\n").utf8))
|
||||
}
|
||||
|
||||
// MARK: - Errors
|
||||
|
||||
public enum TokenizerError: Error, LocalizedError {
|
||||
case invalidFormat(String)
|
||||
|
||||
public var errorDescription: String? {
|
||||
switch self {
|
||||
case .invalidFormat(let reason):
|
||||
return "Invalid tokenizer format: \(reason)"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple tokenizer for Qwen3 that loads from vocab.json
|
||||
/// Supports decoding (id->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 <asr_text> 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>)] = [
|
||||
(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|>", "<asr_text>", "<|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 {}
|
||||
@@ -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..<frameCount {
|
||||
let l = max(-1.0, min(1.0, left[i]))
|
||||
let r = max(-1.0, min(1.0, right[i]))
|
||||
appendInt16(&data, Int16(l * 32767.0))
|
||||
appendInt16(&data, Int16(r * 32767.0))
|
||||
}
|
||||
|
||||
try data.write(to: url)
|
||||
}
|
||||
|
||||
private static func appendUInt32(_ data: inout Data, _ value: UInt32) {
|
||||
var v = value.littleEndian
|
||||
data.append(Data(bytes: &v, count: 4))
|
||||
}
|
||||
|
||||
private static func appendUInt16(_ data: inout Data, _ value: UInt16) {
|
||||
var v = value.littleEndian
|
||||
data.append(Data(bytes: &v, count: 2))
|
||||
}
|
||||
|
||||
private static func appendInt16(_ data: inout Data, _ value: Int16) {
|
||||
var v = value.littleEndian
|
||||
data.append(Data(bytes: &v, count: 2))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
import Foundation
|
||||
import MLX
|
||||
import MLXNN
|
||||
|
||||
// MARK: - LSTM cell matching mlx-community EnCodec weight layout
|
||||
|
||||
/// Single-layer LSTM cell that consumes weight tensors stored under the keys
|
||||
/// `Wx`, `Wh`, `bias` — the exact layout used by the `mlx-community` EnCodec
|
||||
/// safetensors (24 kHz, 32 kHz, 48 kHz variants).
|
||||
///
|
||||
/// Runs the full input sequence in one call and returns the per-step hidden
|
||||
/// state stack `[B, T, hiddenSize]`. There is no streaming variant — EnCodec
|
||||
/// is used on bounded chunks so a full-sequence pass is fine and keeps the
|
||||
/// gate math straightforward.
|
||||
///
|
||||
/// Gate ordering matches Apple's mlx-examples Encodec port (i, f, g, o).
|
||||
public final class EncodecLSTMCell: Module {
|
||||
@ParameterInfo public var Wx: MLXArray
|
||||
@ParameterInfo public var Wh: MLXArray
|
||||
@ParameterInfo public var bias: MLXArray
|
||||
public let hiddenSize: Int
|
||||
|
||||
public init(inputSize: Int, hiddenSize: Int) {
|
||||
self.hiddenSize = hiddenSize
|
||||
self._Wx = ParameterInfo(wrappedValue: MLXArray.zeros([4 * hiddenSize, inputSize]))
|
||||
self._Wh = ParameterInfo(wrappedValue: MLXArray.zeros([4 * hiddenSize, hiddenSize]))
|
||||
self._bias = ParameterInfo(wrappedValue: MLXArray.zeros([4 * hiddenSize]))
|
||||
super.init()
|
||||
}
|
||||
|
||||
/// `x: [B, T, inputSize]` → `[B, T, hiddenSize]`.
|
||||
public func callAsFunction(_ x: MLXArray) -> 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..<T {
|
||||
let xt = xT[0..., t, 0...] // [B, 4*H]
|
||||
let gates = xt + matmul(h, Wh.T) // [B, 4*H]
|
||||
let iGate = sigmoid(gates[0..., 0..<H])
|
||||
let fGate = sigmoid(gates[0..., H..<(2 * H)])
|
||||
let gGate = tanh(gates[0..., (2 * H)..<(3 * H)])
|
||||
let oGate = sigmoid(gates[0..., (3 * H)..<(4 * H)])
|
||||
c = fGate * c + iGate * gGate
|
||||
h = oGate * tanh(c)
|
||||
outputs.append(h)
|
||||
}
|
||||
return stacked(outputs, axis: 1)
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Stacked LSTM with residual
|
||||
|
||||
/// EnCodec's LSTM block: `numLayers` of `EncodecLSTMCell` stacked, with a
|
||||
/// residual add of the original input to the stack's final hidden output.
|
||||
/// Weight key is `lstm.<i>.{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..<numLayers).map { _ in
|
||||
EncodecLSTMCell(inputSize: dimension, hiddenSize: dimension)
|
||||
})
|
||||
super.init()
|
||||
}
|
||||
|
||||
public func callAsFunction(_ x: MLXArray) -> MLXArray {
|
||||
var h = x
|
||||
for cell in lstm { h = cell(h) }
|
||||
return h + x
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
+50
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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<String, MLXArray>] = [:]
|
||||
|
||||
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<String, MLXArray>] = [:]
|
||||
|
||||
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<String, MLXArray>] = [:]
|
||||
|
||||
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<String, MLXArray>] = [:]
|
||||
|
||||
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<String, MLXArray>] = [:]
|
||||
|
||||
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<String, MLXArray>] = [:]
|
||||
|
||||
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<String, MLXArray>] = [:]
|
||||
|
||||
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<String, MLXArray>] = [:]
|
||||
|
||||
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)"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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..<halfDim).asType(.float32) * (-logTimescaleIncrement)
|
||||
)
|
||||
|
||||
// Position indices: [0, 1, 2, ..., seqLen-1]
|
||||
let positions = MLXArray(0..<seqLen).asType(.float32)
|
||||
|
||||
// Compute scaled_time: positions[:, None] * inv_timescales[None, :]
|
||||
// [seqLen, 1] * [1, halfDim] -> [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..<config.encoderLayers).map { _ in
|
||||
AudioEncoderLayer(
|
||||
hiddenSize: config.dModel,
|
||||
numHeads: config.encoderAttentionHeads,
|
||||
ffnDim: config.encoderFFNDim,
|
||||
layerNormEps: config.layerNormEps
|
||||
)
|
||||
}
|
||||
|
||||
super.init()
|
||||
}
|
||||
|
||||
/// Calculate output length from input length using the chunking formula
|
||||
/// This matches the Python _get_feat_extract_output_lengths function
|
||||
private func getOutputLength(_ inputLength: Int) -> 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..<end {
|
||||
blockIds[pos] = Int32(i)
|
||||
}
|
||||
}
|
||||
|
||||
// Use broadcast comparison: mask[i,j] = 0 if same block, -1e9 otherwise
|
||||
let rowIds = MLXArray(blockIds).expandedDimensions(axis: 1) // [seqLen, 1]
|
||||
let colIds = MLXArray(blockIds).expandedDimensions(axis: 0) // [1, seqLen]
|
||||
|
||||
// Where block IDs match: 0 (attend), otherwise -1e9 (block)
|
||||
let mask = MLX.where(rowIds .== colIds, MLXArray(Float(0)), MLXArray(Float(-1e9)))
|
||||
|
||||
// Add batch and head dimensions: [seqLen, seqLen] -> [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 {
|
||||
if i == numChunks - 1 {
|
||||
let remainder = timeFrames % chunkSize
|
||||
chunkLengths.append(remainder == 0 ? chunkSize : remainder)
|
||||
} else {
|
||||
chunkLengths.append(chunkSize)
|
||||
}
|
||||
}
|
||||
|
||||
let maxChunkLen = chunkLengths.max() ?? chunkSize
|
||||
|
||||
// Extract and pad chunks - stack as batch dimension for parallel processing
|
||||
var paddedChunks: [MLXArray] = []
|
||||
var pos = 0
|
||||
for i in 0..<numChunks {
|
||||
let clen = chunkLengths[i]
|
||||
// Extract chunk from first (and only) batch item
|
||||
let feat = melFeatures[0, 0..., pos..<(pos + clen)] // [mel, clen]
|
||||
pos += clen
|
||||
|
||||
// Pad if needed
|
||||
var chunk: MLXArray
|
||||
if clen < maxChunkLen {
|
||||
let padWidth = maxChunkLen - clen
|
||||
chunk = padded(feat, widths: [.init((low: 0, high: 0)), .init((low: 0, high: padWidth))])
|
||||
} else {
|
||||
chunk = feat
|
||||
}
|
||||
paddedChunks.append(chunk)
|
||||
}
|
||||
|
||||
// Stack chunks as batch: [numChunks, mel, maxChunkLen]
|
||||
let paddedFeature = stacked(paddedChunks, axis: 0)
|
||||
|
||||
// Add channel dim: [numChunks, mel, time, 1] for Conv2d (NHWC)
|
||||
var x = paddedFeature.expandedDimensions(axis: -1)
|
||||
|
||||
// Process through conv layers
|
||||
x = conv2d1(x)
|
||||
x = gelu(x)
|
||||
x = conv2d2(x)
|
||||
x = gelu(x)
|
||||
x = conv2d3(x)
|
||||
x = gelu(x)
|
||||
|
||||
// Shape after conv: [numChunks, freq=16, time', channels=480]
|
||||
let numChunksBatch = x.dim(0)
|
||||
let freq = x.dim(1) // 16
|
||||
let timeAfterConv = x.dim(2) // ~13 for 100 input frames
|
||||
let channels = x.dim(3) // 480
|
||||
|
||||
// Transpose and reshape: [numChunks, freq, time, channels] -> [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..<numChunks {
|
||||
let validLen = featureLensAfterCnn[i]
|
||||
let chunkHidden = x[i, 0..<validLen, 0...] // [validLen, 896]
|
||||
hiddenList.append(chunkHidden)
|
||||
}
|
||||
|
||||
// Concatenate all valid hidden states: [totalTokens, 896]
|
||||
var hiddenStates = concatenated(hiddenList, axis: 0)
|
||||
let totalTokens = hiddenStates.dim(0)
|
||||
|
||||
// Build cumulative sequence lengths for block attention mask
|
||||
let maxLenAfterCnn = featureLensAfterCnn.max() ?? 13
|
||||
let windowAfterCnn = maxLenAfterCnn * (config.nWindowInfer / (config.nWindow * 2)) // 13 * 8 = 104
|
||||
|
||||
let totalCnnLen = getOutputLength(timeFrames) // 260
|
||||
|
||||
// Build chunk lengths for windowed attention (NOT starting with 0)
|
||||
var cuChunkLens = [Int]()
|
||||
let numFullWindows = totalCnnLen / windowAfterCnn
|
||||
for _ in 0..<numFullWindows {
|
||||
cuChunkLens.append(windowAfterCnn)
|
||||
}
|
||||
let windowRemainder = totalCnnLen % windowAfterCnn
|
||||
if windowRemainder != 0 {
|
||||
cuChunkLens.append(windowRemainder)
|
||||
}
|
||||
|
||||
// Compute cumulative sums: [0, 104, 208, 260]
|
||||
var cuSeqlens = [0]
|
||||
var cumsum = 0
|
||||
for len in cuChunkLens {
|
||||
cumsum += len
|
||||
cuSeqlens.append(cumsum)
|
||||
}
|
||||
|
||||
// Create block attention mask
|
||||
let attentionMask = createBlockAttentionMask(seqLen: totalTokens, cuSeqlens: cuSeqlens)
|
||||
|
||||
// Add batch dimension for transformer: [1, totalTokens, 896]
|
||||
hiddenStates = hiddenStates.expandedDimensions(axis: 0)
|
||||
|
||||
// Apply transformer layers with attention mask
|
||||
for layer in layers {
|
||||
hiddenStates = layer(hiddenStates, attentionMask: attentionMask)
|
||||
}
|
||||
|
||||
// Remove batch dimension: [totalTokens, 896]
|
||||
hiddenStates = hiddenStates.squeezed(axis: 0)
|
||||
|
||||
// Post processing
|
||||
hiddenStates = lnPost(hiddenStates)
|
||||
|
||||
// Project to text model dimension (GELU activation)
|
||||
hiddenStates = proj1(hiddenStates)
|
||||
hiddenStates = gelu(hiddenStates)
|
||||
hiddenStates = proj2(hiddenStates)
|
||||
|
||||
return hiddenStates
|
||||
}
|
||||
}
|
||||
+490
@@ -0,0 +1,490 @@
|
||||
import Foundation
|
||||
import Accelerate
|
||||
import MLX
|
||||
import AudioCommon
|
||||
|
||||
/// MLX-free container for mel spectrogram data.
|
||||
/// Layout: `data` is a flat [Float] in row-major [melBins, timeFrames] order.
|
||||
public struct MelFeatures {
|
||||
public let data: [Float]
|
||||
public let melBins: Int
|
||||
public let timeFrames: Int
|
||||
|
||||
public init(data: [Float], melBins: Int, timeFrames: Int) {
|
||||
self.data = data
|
||||
self.melBins = melBins
|
||||
self.timeFrames = timeFrames
|
||||
}
|
||||
}
|
||||
|
||||
/// Whisper-style feature extractor for Qwen3-ASR
|
||||
/// Converts raw audio to mel spectrograms
|
||||
/// Parameters from HuggingFace preprocessor_config.json
|
||||
public class WhisperFeatureExtractor {
|
||||
public let sampleRate: Int = 16000 // HF WhisperFeatureExtractor uses 16kHz
|
||||
public let nFFT: Int = 400 // FFT size (from config)
|
||||
public let hopLength: Int = 160 // 10ms hop at 16kHz (from config)
|
||||
public let nMels: Int = 128 // Mel filterbank bins (feature_size)
|
||||
public let chunkLength: Int = 30 // Max audio chunk in seconds
|
||||
|
||||
private var melFilterbank: [Float]?
|
||||
|
||||
// Cached Hann window and FFT setup for Accelerate
|
||||
private var hannWindow: [Float]
|
||||
// Power-of-2 FFT: zero-pad nFFT=400 to paddedFFT=512 for vDSP compatibility
|
||||
private let paddedFFT: Int = 512
|
||||
private let log2PaddedFFT: vDSP_Length = 9 // log2(512) = 9
|
||||
private var fftSetup: FFTSetup
|
||||
|
||||
public init() {
|
||||
// Precompute periodic Hann window
|
||||
hannWindow = [Float](repeating: 0, count: 400) // nFFT
|
||||
for i in 0..<400 {
|
||||
hannWindow[i] = 0.5 * (1.0 - cos(2.0 * Float.pi * Float(i) / Float(400)))
|
||||
}
|
||||
|
||||
// Create power-of-2 FFT setup (512-point)
|
||||
guard let setup = vDSP_create_fftsetup(9, FFTRadix(kFFTRadix2)) else {
|
||||
fatalError("Failed to create vDSP FFT setup for paddedFFT=512")
|
||||
}
|
||||
fftSetup = setup
|
||||
|
||||
setupMelFilterbank()
|
||||
}
|
||||
|
||||
deinit {
|
||||
vDSP_destroy_fftsetup(fftSetup)
|
||||
}
|
||||
|
||||
/// Setup mel filterbank matrix with slaney normalization
|
||||
/// Matches HuggingFace transformers.audio_utils.mel_filter_bank exactly
|
||||
private func setupMelFilterbank() {
|
||||
let fMin: Float = 0.0
|
||||
let fMax: Float = Float(sampleRate) / 2.0 // Nyquist frequency (8000 Hz for 16kHz)
|
||||
|
||||
// Slaney mel scale conversion functions (HuggingFace style)
|
||||
// This is a piecewise function: linear below 1000 Hz, logarithmic above
|
||||
let minLogHertz: Float = 1000.0
|
||||
let minLogMel: Float = 15.0
|
||||
let logstepHzToMel: Float = 27.0 / log(6.4) // For Hz->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..<nBins {
|
||||
fftFreqs[i] = Float(i) * Float(sampleRate) / Float(paddedFFT)
|
||||
}
|
||||
|
||||
// Create mel filter center frequencies
|
||||
let melMin = hzToMel(fMin)
|
||||
let melMax = hzToMel(fMax)
|
||||
|
||||
// nMels + 2 points for triangular filters (includes low and high edges)
|
||||
let nMelPoints = nMels + 2
|
||||
var melPoints = [Float](repeating: 0, count: nMelPoints)
|
||||
for i in 0..<nMelPoints {
|
||||
melPoints[i] = melMin + Float(i) * (melMax - melMin) / Float(nMelPoints - 1)
|
||||
}
|
||||
|
||||
// Convert mel points to Hz - these are the filter edge frequencies
|
||||
let filterFreqs = melPoints.map { melToHz($0) }
|
||||
|
||||
// Calculate filter frequency differences for normalization
|
||||
var filterDiff = [Float](repeating: 0, count: nMelPoints - 1)
|
||||
for i in 0..<(nMelPoints - 1) {
|
||||
filterDiff[i] = filterFreqs[i + 1] - filterFreqs[i]
|
||||
}
|
||||
|
||||
// Create filterbank using HuggingFace's _create_triangular_filter_bank approach
|
||||
// This creates smooth triangular filters in frequency space
|
||||
// Output shape: [nBins, nMels] - we'll transpose later for our use
|
||||
var filterbank = [Float](repeating: 0, count: nBins * nMels)
|
||||
|
||||
for bin in 0..<nBins {
|
||||
let fftFreq = fftFreqs[bin]
|
||||
|
||||
for mel in 0..<nMels {
|
||||
// Filter edges: filterFreqs[mel], filterFreqs[mel+1], filterFreqs[mel+2]
|
||||
let lowFreq = filterFreqs[mel]
|
||||
let highFreq = filterFreqs[mel + 2]
|
||||
|
||||
// Calculate slopes (HuggingFace formula)
|
||||
// slopes = filter_freqs - fft_freqs (broadcast)
|
||||
// down_slopes = -slopes[:, :-2] / filter_diff[:-1]
|
||||
// up_slopes = slopes[:, 2:] / filter_diff[1:]
|
||||
|
||||
let downSlope = (fftFreq - lowFreq) / filterDiff[mel] // Rising edge
|
||||
let upSlope = (highFreq - fftFreq) / filterDiff[mel + 1] // Falling edge
|
||||
|
||||
// Triangular filter: max(0, min(down_slope, up_slope))
|
||||
let filterValue = max(0.0, min(downSlope, upSlope))
|
||||
|
||||
// Store in [nBins, nMels] layout
|
||||
filterbank[bin * nMels + mel] = filterValue
|
||||
}
|
||||
}
|
||||
|
||||
// Apply slaney normalization: 2.0 / (high_freq - low_freq) for each mel filter
|
||||
for mel in 0..<nMels {
|
||||
let enorm = 2.0 / (filterFreqs[mel + 2] - filterFreqs[mel])
|
||||
for bin in 0..<nBins {
|
||||
filterbank[bin * nMels + mel] *= enorm
|
||||
}
|
||||
}
|
||||
|
||||
// Transpose to [nMels, nBins] for our matrix multiplication
|
||||
var filterbankTransposed = [Float](repeating: 0, count: nMels * nBins)
|
||||
for mel in 0..<nMels {
|
||||
for bin in 0..<nBins {
|
||||
filterbankTransposed[mel * nBins + bin] = filterbank[bin * nMels + mel]
|
||||
}
|
||||
}
|
||||
|
||||
self.melFilterbank = filterbankTransposed
|
||||
}
|
||||
|
||||
/// Extract mel spectrogram features from audio samples
|
||||
/// - Parameter audio: Raw audio samples (Float array, mono, at sampleRate)
|
||||
/// - Returns: Mel spectrogram [mel_bins, time_frames]
|
||||
public func extractFeatures(_ audio: [Float]) -> 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..<padLength {
|
||||
let srcIdx = min(padLength - i, audio.count - 1)
|
||||
paddedAudio[i] = audio[max(0, srcIdx)]
|
||||
}
|
||||
|
||||
// Copy original audio
|
||||
for i in 0..<audio.count {
|
||||
paddedAudio[padLength + i] = audio[i]
|
||||
}
|
||||
|
||||
// Reflect pad right side
|
||||
for i in 0..<padLength {
|
||||
let srcIdx = audio.count - 2 - i
|
||||
paddedAudio[padLength + audio.count + i] = audio[max(0, srcIdx)]
|
||||
}
|
||||
|
||||
// Calculate number of frames
|
||||
let nFrames = (paddedAudio.count - nFFT) / hopLength + 1
|
||||
|
||||
// --- Accelerate FFT-based STFT using vDSP_fft_zrip ---
|
||||
// Zero-pad nFFT=400 samples to paddedFFT=512 for power-of-2 FFT
|
||||
// vDSP_fft_zrip uses split-complex in-place: even-indexed in realp, odd-indexed in imagp
|
||||
// Output packing: DC in realp[0], Nyquist in imagp[0], bins 1..N/2-1 in realp[k]+j*imagp[k]
|
||||
|
||||
// Preallocate buffers outside the frame loop
|
||||
var splitReal = [Float](repeating: 0, count: halfPadded)
|
||||
var splitImag = [Float](repeating: 0, count: halfPadded)
|
||||
var paddedFrame = [Float](repeating: 0, count: paddedFFT)
|
||||
|
||||
var magnitude = [Float](repeating: 0, count: nFrames * nBins)
|
||||
|
||||
for frame in 0..<nFrames {
|
||||
let start = frame * hopLength
|
||||
|
||||
// Apply window and write into paddedFrame (first nFFT elements)
|
||||
paddedAudio.withUnsafeBufferPointer { buf in
|
||||
vDSP_vmul(buf.baseAddress! + start, 1, hannWindow, 1, &paddedFrame, 1, vDSP_Length(nFFT))
|
||||
}
|
||||
// Zero-pad the rest (nFFT..<paddedFFT)
|
||||
for i in nFFT..<paddedFFT {
|
||||
paddedFrame[i] = 0
|
||||
}
|
||||
|
||||
// Pack into split-complex: realp[i] = frame[2*i], imagp[i] = frame[2*i+1]
|
||||
for i in 0..<halfPadded {
|
||||
splitReal[i] = paddedFrame[2 * i]
|
||||
splitImag[i] = paddedFrame[2 * i + 1]
|
||||
}
|
||||
|
||||
// Execute in-place real FFT
|
||||
splitReal.withUnsafeMutableBufferPointer { realBuf in
|
||||
splitImag.withUnsafeMutableBufferPointer { imagBuf in
|
||||
var splitComplex = DSPSplitComplex(
|
||||
realp: realBuf.baseAddress!,
|
||||
imagp: imagBuf.baseAddress!)
|
||||
vDSP_fft_zrip(fftSetup, &splitComplex, 1, log2PaddedFFT, FFTDirection(kFFTDirection_Forward))
|
||||
}
|
||||
}
|
||||
|
||||
// Extract power spectrum
|
||||
let baseIdx = frame * nBins
|
||||
|
||||
// DC component: realp[0]^2 (DC is purely real, stored in realp[0])
|
||||
magnitude[baseIdx] = splitReal[0] * splitReal[0]
|
||||
|
||||
// Nyquist component: imagp[0]^2 (Nyquist is purely real, packed in imagp[0])
|
||||
magnitude[baseIdx + halfPadded] = splitImag[0] * splitImag[0]
|
||||
|
||||
// Bins 1 to N/2-1: realp[k]^2 + imagp[k]^2
|
||||
for k in 1..<halfPadded {
|
||||
magnitude[baseIdx + k] = splitReal[k] * splitReal[k] + splitImag[k] * splitImag[k]
|
||||
}
|
||||
}
|
||||
|
||||
// --- BLAS sgemm for mel filterbank ---
|
||||
guard let filterbank = melFilterbank else {
|
||||
fatalError("Mel filterbank not initialized")
|
||||
}
|
||||
|
||||
var melSpec = [Float](repeating: 0, count: nFrames * nMels)
|
||||
|
||||
// melSpec[nFrames, nMels] = magnitude[nFrames, nBins] * filterbankT[nBins, nMels]
|
||||
// filterbank is [nMels, nBins]. We need A * B^T.
|
||||
// vDSP_mmul computes C = A * B, so we need to pre-transpose filterbank.
|
||||
// Transpose filterbank [nMels, nBins] -> 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..<padLength {
|
||||
let srcIdx = min(padLength - i, audio.count - 1)
|
||||
paddedAudio[i] = audio[max(0, srcIdx)]
|
||||
}
|
||||
|
||||
// Copy original audio
|
||||
for i in 0..<audio.count {
|
||||
paddedAudio[padLength + i] = audio[i]
|
||||
}
|
||||
|
||||
// Reflect pad right side
|
||||
for i in 0..<padLength {
|
||||
let srcIdx = audio.count - 2 - i
|
||||
paddedAudio[padLength + audio.count + i] = audio[max(0, srcIdx)]
|
||||
}
|
||||
|
||||
// Calculate number of frames
|
||||
let nFrames = (paddedAudio.count - nFFT) / hopLength + 1
|
||||
|
||||
// --- Accelerate FFT-based STFT using vDSP_fft_zrip ---
|
||||
var splitReal = [Float](repeating: 0, count: halfPadded)
|
||||
var splitImag = [Float](repeating: 0, count: halfPadded)
|
||||
var paddedFrame = [Float](repeating: 0, count: paddedFFT)
|
||||
|
||||
var magnitude = [Float](repeating: 0, count: nFrames * nBins)
|
||||
|
||||
for frame in 0..<nFrames {
|
||||
let start = frame * hopLength
|
||||
|
||||
paddedAudio.withUnsafeBufferPointer { buf in
|
||||
vDSP_vmul(buf.baseAddress! + start, 1, hannWindow, 1, &paddedFrame, 1, vDSP_Length(nFFT))
|
||||
}
|
||||
for i in nFFT..<paddedFFT {
|
||||
paddedFrame[i] = 0
|
||||
}
|
||||
|
||||
for i in 0..<halfPadded {
|
||||
splitReal[i] = paddedFrame[2 * i]
|
||||
splitImag[i] = paddedFrame[2 * i + 1]
|
||||
}
|
||||
|
||||
splitReal.withUnsafeMutableBufferPointer { realBuf in
|
||||
splitImag.withUnsafeMutableBufferPointer { imagBuf in
|
||||
var splitComplex = DSPSplitComplex(
|
||||
realp: realBuf.baseAddress!,
|
||||
imagp: imagBuf.baseAddress!)
|
||||
vDSP_fft_zrip(fftSetup, &splitComplex, 1, log2PaddedFFT, FFTDirection(kFFTDirection_Forward))
|
||||
}
|
||||
}
|
||||
|
||||
let baseIdx = frame * nBins
|
||||
magnitude[baseIdx] = splitReal[0] * splitReal[0]
|
||||
magnitude[baseIdx + halfPadded] = splitImag[0] * splitImag[0]
|
||||
for k in 1..<halfPadded {
|
||||
magnitude[baseIdx + k] = splitReal[k] * splitReal[k] + splitImag[k] * splitImag[k]
|
||||
}
|
||||
}
|
||||
|
||||
// --- BLAS sgemm for mel filterbank ---
|
||||
guard let filterbank = melFilterbank else {
|
||||
fatalError("Mel filterbank not initialized")
|
||||
}
|
||||
|
||||
var melSpec = [Float](repeating: 0, count: nFrames * nMels)
|
||||
|
||||
var filterbankT = [Float](repeating: 0, count: nBins * nMels)
|
||||
vDSP_mtrans(filterbank, 1, &filterbankT, 1, vDSP_Length(nBins), vDSP_Length(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)
|
||||
|
||||
var epsilon: Float = 1e-10
|
||||
vDSP_vclip(melSpec, 1, &epsilon, [Float.greatestFiniteMagnitude], &melSpec, 1, vDSP_Length(count))
|
||||
|
||||
vvlog10f(&melSpec, melSpec, &countN)
|
||||
|
||||
var maxVal: Float = -Float.infinity
|
||||
vDSP_maxv(melSpec, 1, &maxVal, vDSP_Length(count))
|
||||
|
||||
var minClamp = maxVal - 8.0
|
||||
var maxClamp = Float.greatestFiniteMagnitude
|
||||
vDSP_vclip(melSpec, 1, &minClamp, &maxClamp, &melSpec, 1, vDSP_Length(count))
|
||||
|
||||
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
|
||||
let trimmedFrames = nFrames - 1
|
||||
let trimmedMelSpec = Array(melSpec.prefix(trimmedFrames * nMels))
|
||||
|
||||
let maxFrames = 1200 * sampleRate / hopLength // 120000 frames
|
||||
let finalFrames: Int
|
||||
let finalMelSpec: [Float]
|
||||
if trimmedFrames > 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..<finalFrames {
|
||||
for m in 0..<nMels {
|
||||
transposed[m * finalFrames + t] = finalMelSpec[t * nMels + m]
|
||||
}
|
||||
}
|
||||
return MelFeatures(data: transposed, melBins: nMels, timeFrames: finalFrames)
|
||||
}
|
||||
|
||||
/// Process audio for Qwen3-ASR without MLXArray dependency.
|
||||
/// MLX-free equivalent of `process(_:sampleRate:)`.
|
||||
/// - Parameter audio: Raw audio samples (any sample rate)
|
||||
/// - Parameter inputSampleRate: Sample rate of input audio
|
||||
/// - Returns: `MelFeatures` with layout `[melBins, timeFrames]`
|
||||
public func processRaw(_ audio: [Float], sampleRate inputSampleRate: Int) -> 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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -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..<maxTokens {
|
||||
if nextToken == imEndId { break }
|
||||
let embedding = try decoder.embed(tokenId: nextToken)
|
||||
logits = try decoder.decoderStep(embedding: embedding)
|
||||
nextToken = decoder.argmax(logits: logits)
|
||||
generatedTokens.append(nextToken)
|
||||
}
|
||||
if profile {
|
||||
let audioDuration = Double(audio.count) / Double(sampleRate)
|
||||
print("[COREML-ASR-PROFILE] audio=\(audioDuration)s gen=\(generatedTokens.count)")
|
||||
}
|
||||
let t6 = CFAbsoluteTimeGetCurrent()
|
||||
|
||||
if profile {
|
||||
let ms = { (a: CFAbsoluteTime, b: CFAbsoluteTime) in (b - a) * 1000 }
|
||||
print(String(format: "[COREML-ASR-PROFILE] mel=%.0fms encoder=%.0fms prefix=%.0fms audio_prefill=%.0fms(%dtok→%.1fms/tok) suffix=%.0fms gen=%.0fms(%dtok→%.1fms/tok) total=%.0fms",
|
||||
ms(t0, t1), ms(t1, t2), ms(t2, t3),
|
||||
ms(t3, t4), numAudioTokens, ms(t3, t4) / Double(max(numAudioTokens, 1)),
|
||||
ms(t4, t5),
|
||||
ms(t5, t6), generatedTokens.count, ms(t5, t6) / Double(max(generatedTokens.count, 1)),
|
||||
ms(t0, t6)))
|
||||
}
|
||||
|
||||
// Decode tokens
|
||||
if let tokenizer = tokenizer {
|
||||
let rawText = tokenizer.decode(tokens: generatedTokens.map { Int($0) })
|
||||
if let range = rawText.range(of: "<asr_text>") {
|
||||
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..<numAudioTokens {
|
||||
let audioEmbed = try decoder.audioEmbeddingFromMultiArray(audioEmbeds, at: i)
|
||||
lastLogits = try decoder.decoderStep(embedding: audioEmbed)
|
||||
}
|
||||
|
||||
// Prefill: process suffix tokens
|
||||
for token in suffixTokens {
|
||||
let embedding = try decoder.embed(tokenId: token)
|
||||
lastLogits = try decoder.decoderStep(embedding: embedding)
|
||||
}
|
||||
|
||||
// 7. Autoregressive generation (same EOS note as `transcribe()` —
|
||||
// see that path for the background).
|
||||
guard var logits = lastLogits else {
|
||||
return "[CoreML decoder: no output]"
|
||||
}
|
||||
|
||||
var generatedTokens: [Int32] = []
|
||||
var nextToken = decoder.argmax(logits: logits)
|
||||
if nextToken == imEndId {
|
||||
nextToken = decoder.argmax(logits: logits, skipping: imEndId)
|
||||
}
|
||||
generatedTokens.append(nextToken)
|
||||
|
||||
for _ in 1..<maxTokens {
|
||||
if nextToken == imEndId { break }
|
||||
let embedding = try decoder.embed(tokenId: nextToken)
|
||||
logits = try decoder.decoderStep(embedding: embedding)
|
||||
nextToken = decoder.argmax(logits: logits)
|
||||
generatedTokens.append(nextToken)
|
||||
}
|
||||
|
||||
// Decode tokens
|
||||
if let tokenizer = tokenizer {
|
||||
let rawText = tokenizer.decode(tokens: generatedTokens.map { Int($0) })
|
||||
if let range = rawText.range(of: "<asr_text>") {
|
||||
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
|
||||
@@ -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..<melBins {
|
||||
let src = bin * realFrames
|
||||
let dst = bin * padded
|
||||
for t in 0..<realFrames { mptr[dst + t] = melData[src + t] }
|
||||
for t in realFrames..<padded { mptr[dst + t] = 0 }
|
||||
}
|
||||
// mel_length input: [1] int32 with the real (un-padded) frame count.
|
||||
let lengthArray = try MLMultiArray(shape: [1], dataType: .int32)
|
||||
lengthArray[0] = NSNumber(value: Int32(realFrames))
|
||||
|
||||
let input = try MLDictionaryFeatureProvider(dictionary: [
|
||||
"mel": MLFeatureValue(multiArray: melArray),
|
||||
"mel_length": MLFeatureValue(multiArray: lengthArray),
|
||||
])
|
||||
let output = try model.prediction(from: input)
|
||||
|
||||
guard let embeddings = output.featureValue(for: "audio_embeddings")?.multiArrayValue else {
|
||||
throw AudioModelError.inferenceFailed(
|
||||
operation: "CoreML encoder", reason: "Missing audio_embeddings output")
|
||||
}
|
||||
guard let lengthOut = output.featureValue(for: "output_length")?.multiArrayValue else {
|
||||
throw AudioModelError.inferenceFailed(
|
||||
operation: "CoreML encoder", reason: "Missing output_length output (encoder may be an older export without the chunked-attention mask)")
|
||||
}
|
||||
let outLen = max(0, Int(lengthOut[0].int32Value))
|
||||
return EncodedAudio(embeddings: embeddings, outputLength: outLen)
|
||||
}
|
||||
|
||||
private func multiArrayToMLXArray(_ array: MLMultiArray) -> 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..<count { floats[i] = Float(src[i]) }
|
||||
return MLXArray(floats, shape)
|
||||
case .float32:
|
||||
let src = array.dataPointer.assumingMemoryBound(to: Float.self)
|
||||
return MLXArray(Array(UnsafeBufferPointer(start: src, count: count)), shape)
|
||||
default:
|
||||
let src = array.dataPointer.assumingMemoryBound(to: Float.self)
|
||||
return MLXArray(Array(UnsafeBufferPointer(start: src, count: count)), shape)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Qwen3ASRModel Integration
|
||||
|
||||
extension Qwen3ASRModel {
|
||||
/// Transcribe audio using CoreML encoder + MLX text decoder.
|
||||
///
|
||||
/// This hybrid approach runs the encoder on Neural Engine (CoreML) and the
|
||||
/// text decoder on GPU (MLX), combining the power efficiency of ANE with
|
||||
/// the flexibility of MLX for autoregressive decoding.
|
||||
public func transcribe(
|
||||
audio: [Float],
|
||||
sampleRate: Int = 16000,
|
||||
language: String? = nil,
|
||||
maxTokens: Int = 448,
|
||||
coremlEncoder: CoreMLASREncoder
|
||||
) throws -> 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..<audioLength, 0...]
|
||||
return generateText(
|
||||
audioEmbeds: realEmbeds,
|
||||
textDecoder: textDecoder,
|
||||
language: language,
|
||||
maxTokens: maxTokens
|
||||
)
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,566 @@
|
||||
#if canImport(CoreML)
|
||||
import CoreML
|
||||
import Foundation
|
||||
import MLX
|
||||
import AudioCommon
|
||||
|
||||
/// CoreML text decoder for Qwen3-ASR with MLState KV cache.
|
||||
///
|
||||
/// Runs the full text decoder on Neural Engine via CoreML instead of GPU via MLX.
|
||||
/// Requires macOS 15+ / iOS 18+ for MLState support.
|
||||
///
|
||||
/// Architecture (three CoreML models for ANE compile-budget reasons):
|
||||
/// - **embedding**: Token ID → embedding vector lookup
|
||||
/// - **decoder_part1**: Layers 0..(split-1), embedding-in → hidden-out
|
||||
/// - **decoder_part2**: Layers split..N-1 + norm + lm_head, hidden-in → logits
|
||||
///
|
||||
/// Each part keeps its own ``MLState`` pool of KV caches. Hidden state
|
||||
/// flows part1 → part2.
|
||||
///
|
||||
/// **Batched dispatch.** Both decoder parts are converted with a fixed
|
||||
/// ``T`` (per ``config.json: enumerated_t`` — typically 128). One ANE
|
||||
/// dispatch processes T tokens at once: prefill chunks a contiguous run
|
||||
/// of new positions through a single call; single-token generation
|
||||
/// reserves indices ``[0, T-2]`` for *scratch* positions (the last T-1
|
||||
/// slots of the KV cache, addresses ``maxSeqLength - (T-1)..maxSeqLength``)
|
||||
/// whose writes are discarded by attention masking. This collapses ~250
|
||||
/// single-token audio-prefill dispatches into ~2 batched calls.
|
||||
///
|
||||
/// EnumeratedShapes(T) is **not** ANE-compatible at this layer count,
|
||||
/// so the model is fixed-T and we wrap step calls with scratch padding
|
||||
/// rather than running a smaller variant.
|
||||
public class CoreMLTextDecoder {
|
||||
private let embeddingModel: MLModel
|
||||
private let decoderPart1Model: MLModel
|
||||
private let decoderPart2Model: MLModel
|
||||
private let maxSeqLength: Int
|
||||
private let vocabSize: Int
|
||||
private let hiddenSize: Int
|
||||
|
||||
/// Fixed batch size used by both decoder parts. Loaded from
|
||||
/// ``config.json: enumerated_t``. Single-token decode pads to this.
|
||||
private let batchSize: Int
|
||||
|
||||
/// One MLState per decoder part, each holds that part's KV caches.
|
||||
private var part1State: MLState
|
||||
private var part2State: MLState
|
||||
|
||||
/// Current real position in the KV cache (incremented per real token).
|
||||
private var currentPosition: Int = 0
|
||||
|
||||
/// First slot index reserved for scratch writes by partial / step calls.
|
||||
/// Cache writes at these positions are garbage; attention masks them.
|
||||
/// Range: ``[scratchStart, maxSeqLength)``, length ``batchSize - 1``.
|
||||
private var scratchStart: Int { maxSeqLength - (batchSize - 1) }
|
||||
|
||||
public static let defaultModelId = "aufklarer/Qwen3-ASR-CoreML"
|
||||
|
||||
public init(
|
||||
embeddingModel: MLModel,
|
||||
decoderPart1Model: MLModel,
|
||||
decoderPart2Model: MLModel,
|
||||
maxSeqLength: Int = 1024,
|
||||
vocabSize: Int = 151936,
|
||||
hiddenSize: Int = 1024,
|
||||
batchSize: Int = 128
|
||||
) {
|
||||
self.embeddingModel = embeddingModel
|
||||
self.decoderPart1Model = decoderPart1Model
|
||||
self.decoderPart2Model = decoderPart2Model
|
||||
self.maxSeqLength = maxSeqLength
|
||||
self.vocabSize = vocabSize
|
||||
self.hiddenSize = hiddenSize
|
||||
self.batchSize = batchSize
|
||||
self.part1State = decoderPart1Model.makeState()
|
||||
self.part2State = decoderPart2Model.makeState()
|
||||
}
|
||||
|
||||
/// Load decoder models from a directory containing
|
||||
/// ``embedding.mlmodelc``, ``decoder_part1.mlmodelc`` and ``decoder_part2.mlmodelc``.
|
||||
public static func load(
|
||||
from directory: URL,
|
||||
computeUnits: MLComputeUnits = CoreMLComputeUnitsResolver.resolved(default: .cpuAndNeuralEngine)
|
||||
) throws -> 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..<batchSize { dummyPositions[i] = NSNumber(value: Int32(i)) }
|
||||
let dummyMask = try MLMultiArray(shape: [1, 1, batchSize as NSNumber, maxSeqLength as NSNumber],
|
||||
dataType: .float32)
|
||||
let mptr = dummyMask.dataPointer.assumingMemoryBound(to: Float.self)
|
||||
for i in 0..<(batchSize * maxSeqLength) { mptr[i] = -1e4 }
|
||||
|
||||
let warmInputs = try MLDictionaryFeatureProvider(dictionary: [
|
||||
"input_embeds": MLFeatureValue(multiArray: dummyEmbeds),
|
||||
"positions": MLFeatureValue(multiArray: dummyPositions),
|
||||
"attention_mask": MLFeatureValue(multiArray: dummyMask),
|
||||
])
|
||||
let p1Out = try decoderPart1Model.prediction(from: warmInputs, using: warmP1)
|
||||
guard let hidden = p1Out.featureValue(for: "hidden_state")?.multiArrayValue else {
|
||||
throw AudioModelError.inferenceFailed(
|
||||
operation: "CoreML decoder part1 warmup",
|
||||
reason: "Missing hidden_state output")
|
||||
}
|
||||
let p2Inputs = try MLDictionaryFeatureProvider(dictionary: [
|
||||
"input_embeds": MLFeatureValue(multiArray: hidden),
|
||||
"positions": MLFeatureValue(multiArray: dummyPositions),
|
||||
"attention_mask": MLFeatureValue(multiArray: dummyMask),
|
||||
])
|
||||
_ = try decoderPart2Model.prediction(from: p2Inputs, using: warmP2)
|
||||
}
|
||||
|
||||
/// Reset the KV caches in both parts for a new transcription.
|
||||
public func resetCache() {
|
||||
currentPosition = 0
|
||||
part1State = decoderPart1Model.makeState()
|
||||
part2State = decoderPart2Model.makeState()
|
||||
}
|
||||
|
||||
// MARK: - Token Operations
|
||||
|
||||
/// Look up the embedding vector for a single token id.
|
||||
/// Returns shape ``[1, 1, hidden_size]``.
|
||||
public func embed(tokenId: Int32) throws -> 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<Float>, 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..<hidden { dst[dstBase + j] = Float(p[base + j * lastStride]) }
|
||||
case .float32:
|
||||
let p = src.dataPointer.assumingMemoryBound(to: Float.self)
|
||||
for j in 0..<hidden { dst[dstBase + j] = p[base + j * lastStride] }
|
||||
default:
|
||||
let p = src.dataPointer.assumingMemoryBound(to: Float.self)
|
||||
for j in 0..<hidden { dst[dstBase + j] = p[base + j * lastStride] }
|
||||
}
|
||||
}
|
||||
|
||||
/// Run a chunked multi-token prefill, source from MLMultiArray rows.
|
||||
///
|
||||
/// ``embeddings`` is shape ``[1, N, hidden]`` containing the real
|
||||
/// next-N tokens to commit to positions ``[currentPosition,
|
||||
/// currentPosition + N)``. ``N`` must be ``<= batchSize``. The
|
||||
/// returned MLMultiArray is the logits for the last real position.
|
||||
@discardableResult
|
||||
public func decoderPrefill(embeddings: MLMultiArray, realCount n: Int) throws -> 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..<n {
|
||||
let srcOff = t * self.hiddenSize
|
||||
let dstOff = (slot + t) * self.hiddenSize
|
||||
for j in 0..<self.hiddenSize {
|
||||
dstPtr[dstOff + j] = srcPtr[srcOff + j]
|
||||
}
|
||||
}
|
||||
}, realCount: n)
|
||||
return try runParts(embeds: bufs.embeds, positions: bufs.positions, mask: bufs.mask)
|
||||
}
|
||||
|
||||
/// Embed a contiguous run of token ids and prefill them in batched
|
||||
/// chunks of ``batchSize``. Used for the chat-template prefix/suffix
|
||||
/// runs, replacing per-token ``embed`` + ``decoderStep`` loops with
|
||||
/// one ANE dispatch per chunk. Returns the logits for the last token.
|
||||
@discardableResult
|
||||
public func decoderPrefillTokens(_ tokenIds: [Int32]) throws -> 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..<n {
|
||||
let emb = try embed(tokenId: tokenIds[consumed + k])
|
||||
Self.copyRow(from: emb, sourceRow: 0, hidden: hiddenSize,
|
||||
to: pptr, destSlot: k)
|
||||
}
|
||||
lastLogits = try decoderPrefill(embeddings: packed, realCount: n)
|
||||
consumed += n
|
||||
}
|
||||
return lastLogits
|
||||
}
|
||||
|
||||
/// Run a chunked prefill where embeddings come from a Float buffer
|
||||
/// (typical: bulk-extracted MLX audio embeddings).
|
||||
@discardableResult
|
||||
public func decoderPrefill(flatEmbeddings: [Float], offset: Int, realCount n: Int) throws -> 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..<n {
|
||||
let srcOff = (offset + t) * self.hiddenSize
|
||||
let dstOff = (slot + t) * self.hiddenSize
|
||||
for j in 0..<self.hiddenSize {
|
||||
dstPtr[dstOff + j] = src[srcOff + j]
|
||||
}
|
||||
}
|
||||
}
|
||||
}, realCount: n)
|
||||
return try runParts(embeds: bufs.embeds, positions: bufs.positions, mask: bufs.mask)
|
||||
}
|
||||
|
||||
// MARK: - Internal: chunked dispatch primitives
|
||||
|
||||
/// Set up the input buffers (positions, mask, embeds) for a chunk of
|
||||
/// ``realCount`` real tokens placed in the LAST ``realCount`` input
|
||||
/// slots, with the remaining slots filled by scratch positions. The
|
||||
/// ``realEmbeddingsSource`` callback is given the slot where real
|
||||
/// data should land and a pointer into the embeds buffer.
|
||||
///
|
||||
/// Allocates FRESH MLMultiArrays each call. Reusing buffers across
|
||||
/// calls was tried and produces wrong outputs — CoreML appears to
|
||||
/// hold references into the input buffer past prediction return,
|
||||
/// so mutating it before the next call corrupts the stored KV state.
|
||||
/// Per-call allocation costs ~0.5 ms (mask is 64 KB) which is well
|
||||
/// below the dispatch budget.
|
||||
private func writeChunk(
|
||||
realEmbeddingsSource: (Int, UnsafeMutablePointer<Float>) -> 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..<firstRealSlot {
|
||||
positions[i] = NSNumber(value: Int32(scratchStart + i))
|
||||
}
|
||||
for i in 0..<n {
|
||||
positions[firstRealSlot + i] = NSNumber(value: Int32(currentPosition + i))
|
||||
}
|
||||
|
||||
// Mask: shape [1, 1, T, MAX_SEQ]
|
||||
let mptr = mask.dataPointer.assumingMemoryBound(to: Float.self)
|
||||
let rowSize = maxSeqLength
|
||||
let scratchStartLocal = scratchStart
|
||||
for t in 0..<T {
|
||||
let rowBase = t * rowSize
|
||||
if t < firstRealSlot {
|
||||
for j in 0..<rowSize {
|
||||
mptr[rowBase + j] = -1e4
|
||||
}
|
||||
} else {
|
||||
let realIdxInChunk = t - firstRealSlot
|
||||
let absPosition = currentPosition + realIdxInChunk
|
||||
for j in 0..<rowSize {
|
||||
if j <= absPosition && j < scratchStartLocal {
|
||||
mptr[rowBase + j] = 0
|
||||
} else {
|
||||
mptr[rowBase + j] = -1e4
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Embeddings: zero the scratch rows + populate the real rows.
|
||||
let eptr = embeds.dataPointer.assumingMemoryBound(to: Float.self)
|
||||
for s in 0..<firstRealSlot {
|
||||
for j in 0..<hiddenSize { eptr[s * hiddenSize + j] = 0 }
|
||||
}
|
||||
realEmbeddingsSource(firstRealSlot, eptr)
|
||||
|
||||
currentPosition += n
|
||||
return (embeds, positions, mask)
|
||||
}
|
||||
|
||||
/// Dispatch part1 then part2 over the given input buffers. Returns
|
||||
/// the part2 logits MLMultiArray (shape ``[1, 1, vocab]`` — the last
|
||||
/// real position's next-token distribution).
|
||||
private func runParts(embeds: MLMultiArray, positions: MLMultiArray, mask: MLMultiArray) throws -> 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..<vocab where i != skipIdx {
|
||||
let val = Float(ptr[i * lastStride])
|
||||
if val.isNaN { nanCount += 1; continue }
|
||||
if val > maxVal {
|
||||
maxVal = val
|
||||
maxIdx = Int32(i)
|
||||
}
|
||||
}
|
||||
case .float32:
|
||||
let ptr = logits.dataPointer.assumingMemoryBound(to: Float.self)
|
||||
for i in 0..<vocab where i != skipIdx {
|
||||
let val = ptr[i * lastStride]
|
||||
if val.isNaN { nanCount += 1; continue }
|
||||
if val > maxVal {
|
||||
maxVal = val
|
||||
maxIdx = Int32(i)
|
||||
}
|
||||
}
|
||||
default:
|
||||
let ptr = logits.dataPointer.assumingMemoryBound(to: Float.self)
|
||||
for i in 0..<vocab where i != skipIdx {
|
||||
let val = ptr[i * lastStride]
|
||||
if val.isNaN { nanCount += 1; continue }
|
||||
if val > 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..<hidden {
|
||||
ptr[i] = data[i]
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
/// Extract audio embedding at index from MLMultiArray (no MLX dependency).
|
||||
public func audioEmbeddingFromMultiArray(_ embeddings: MLMultiArray, at index: Int) throws -> 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..<hidden {
|
||||
dstPtr[i] = srcPtr[offset + i]
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// MARK: - Helpers
|
||||
|
||||
private static func findModel(named name: String, in directory: URL) -> URL? {
|
||||
let compiled = directory.appendingPathComponent("\(name).mlmodelc", isDirectory: true)
|
||||
if FileManager.default.fileExists(atPath: compiled.path) {
|
||||
return compiled
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -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
|
||||
@@ -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..<config.numLayers).map { _ in FloatTextDecoderLayer(config: config) }
|
||||
self._norm.wrappedValue = RMSNorm(dimensions: config.hiddenSize, eps: config.rmsNormEps)
|
||||
super.init()
|
||||
}
|
||||
|
||||
public func callAsFunction(
|
||||
inputIds: MLXArray? = nil,
|
||||
inputsEmbeds: MLXArray? = nil,
|
||||
attentionMask: MLXArray? = nil,
|
||||
cache: [(MLXArray, MLXArray)]? = nil
|
||||
) -> (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..<Int32(seqLen)) + Int32(cacheLen)).expandedDimensions(axis: 1)
|
||||
let cols = MLXArray(0..<Int32(totalLen)).expandedDimensions(axis: 0)
|
||||
mask = MLX.where(cols .> 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)
|
||||
}
|
||||
}
|
||||
+9
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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..<n).reversed() {
|
||||
let dt = abs(aligned[i].startTime - aligned[i - 1].startTime)
|
||||
if dt < tolerance {
|
||||
plateauStart = i - 1
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
return (n - plateauStart) >= 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..<audioStartIndex, 0...]
|
||||
let afterAudio = inputEmbeds[0..., audioEndIndex..., 0...]
|
||||
inputEmbeds = concatenated([beforeAudio, audioEmbedsTyped, afterAudio], axis: 1)
|
||||
|
||||
// 5. Single forward pass through decoder (no cache, no autoregressive loop)
|
||||
let (hiddenStates, _) = textDecoder.decode(inputsEmbeds: inputEmbeds, attentionMask: nil, cache: nil)
|
||||
|
||||
// 6. Apply classify head to ALL hidden states → logits [1, seqLen, classifyNum]
|
||||
let logits = classifyHead(hiddenStates)
|
||||
|
||||
// 7. Extract logits at timestamp positions → argmax → raw indices
|
||||
// Adjust timestamp positions to account for the chat template prefix
|
||||
let absoluteTimestampPositions = slotted.timestampPositions.map { $0 + slottedTextStart }
|
||||
|
||||
var rawIndices: [Int] = []
|
||||
for pos in absoluteTimestampPositions {
|
||||
let posLogits = logits[0, pos, 0...] // [classifyNum]
|
||||
let idx = argMax(posLogits).item(Int32.self)
|
||||
rawIndices.append(Int(idx))
|
||||
}
|
||||
|
||||
// 8. Apply LIS monotonicity correction
|
||||
let correctedIndices = TimestampCorrection.enforceMonotonicity(rawIndices)
|
||||
|
||||
// Optional raw/corrected dump for bug-triage of misaligned timestamps.
|
||||
if ProcessInfo.processInfo.environment["ALIGN_DEBUG"] == "1" {
|
||||
print("[align-debug] indices=\(rawIndices.count) numAudioTokens=\(numAudioTokens)")
|
||||
print("[align-debug] first 10 raw: \(Array(rawIndices.prefix(10)))")
|
||||
print("[align-debug] first 10 corrected: \(Array(correctedIndices.prefix(10)))")
|
||||
print("[align-debug] last 60 raw: \(Array(rawIndices.suffix(60)))")
|
||||
print("[align-debug] last 60 corrected: \(Array(correctedIndices.suffix(60)))")
|
||||
}
|
||||
|
||||
// 9. Convert to seconds and pair as (start, end)
|
||||
let segmentTime = config.timestampSegmentTime
|
||||
var alignedWords: [AlignedWord] = []
|
||||
|
||||
for (wordIdx, word) in slotted.words.enumerated() {
|
||||
let startIdx = wordIdx * 2 // even indices are start timestamps
|
||||
let endIdx = wordIdx * 2 + 1 // odd indices are end timestamps
|
||||
|
||||
guard endIdx < correctedIndices.count else { break }
|
||||
|
||||
let startTime = Float(correctedIndices[startIdx]) * segmentTime
|
||||
let endTime = Float(correctedIndices[endIdx]) * segmentTime
|
||||
|
||||
alignedWords.append(AlignedWord(
|
||||
text: word,
|
||||
startTime: startTime,
|
||||
endTime: max(endTime, startTime) // ensure end >= 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..<numAudioTokens {
|
||||
ids.append(audioPadId)
|
||||
}
|
||||
|
||||
// <|audio_end|><|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
|
||||
}
|
||||
}
|
||||
}
|
||||
+252
@@ -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..<config.numLayers).map { _ in
|
||||
QuantizedTextDecoderLayer(config: config)
|
||||
}
|
||||
self._norm.wrappedValue = RMSNorm(dimensions: config.hiddenSize, eps: config.rmsNormEps)
|
||||
|
||||
super.init()
|
||||
}
|
||||
|
||||
/// Forward pass through text decoder
|
||||
public func callAsFunction(
|
||||
inputIds: MLXArray? = nil,
|
||||
inputsEmbeds: MLXArray? = nil,
|
||||
attentionMask: MLXArray? = nil,
|
||||
cache: [(MLXArray, MLXArray)]? = nil
|
||||
) -> (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..<Int32(seqLen)) + Int32(cacheLen)).expandedDimensions(axis: 1)
|
||||
let cols = MLXArray(0..<Int32(totalLen)).expandedDimensions(axis: 0)
|
||||
mask = MLX.where(cols .> 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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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..<numAudioTokens {
|
||||
inputIds.append(Int32(audioPadId))
|
||||
}
|
||||
let audioEndIndex = inputIds.count
|
||||
|
||||
// <|audio_end|><|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..<audioStartIndex, 0...]
|
||||
let afterAudio = inputEmbeds[0..., audioEndIndex..., 0...]
|
||||
|
||||
inputEmbeds = concatenated([beforeAudio, audioEmbedsTyped, afterAudio], axis: 1)
|
||||
|
||||
// Initialize KV cache
|
||||
var cache: [(MLXArray, MLXArray)]? = nil
|
||||
|
||||
// First pass: process the full input embeddings
|
||||
let (hiddenStates, newCache) = textDecoder(inputsEmbeds: inputEmbeds, cache: cache)
|
||||
cache = newCache
|
||||
|
||||
// Get logits from the last position using embedding as LM head (tied weights)
|
||||
let seqLen = hiddenStates.dim(1)
|
||||
let lastHidden = hiddenStates[0..., (seqLen-1)..<seqLen, 0...]
|
||||
let logits = textDecoder.embedTokens.asLinear(lastHidden)
|
||||
|
||||
// Greedy fast path uses a double-buffered asyncEval decode loop that
|
||||
// overlaps the GPU forward pass for token N+1 with the host-side
|
||||
// bookkeeping (EOS check + Swift array append) for token N. The
|
||||
// legacy slow path stays on the per-token CPU sync because it pulls
|
||||
// the full logits tensor to CPU for repetition / n-gram / temperature
|
||||
// manipulation, which would defeat the overlap.
|
||||
let generatedTokens: [Int32]
|
||||
if Self.isGreedyFastPath(decodingOptions) {
|
||||
generatedTokens = Self.generateGreedyAsyncEval(
|
||||
textDecoder: textDecoder,
|
||||
initialLogits: logits,
|
||||
cache: cache!,
|
||||
maxTokens: maxTokens
|
||||
)
|
||||
} else {
|
||||
generatedTokens = Self.generateSlow(
|
||||
textDecoder: textDecoder,
|
||||
initialLogits: logits,
|
||||
cache: cache!,
|
||||
maxTokens: maxTokens,
|
||||
options: decodingOptions
|
||||
)
|
||||
}
|
||||
|
||||
// Decode tokens to text
|
||||
if let tokenizer = tokenizer {
|
||||
let rawText = tokenizer.decode(tokens: generatedTokens.map { Int($0) })
|
||||
// Strip "language XX<asr_text>" prefix if present (auto-detection output)
|
||||
if let range = rawText.range(of: "<asr_text>") {
|
||||
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..<maxTokens {
|
||||
// Stage N+1's graph BEFORE syncing N. embedTokens expects a
|
||||
// [batch, seq] int32 tensor; nextTokenArr is 0-D so we expand
|
||||
// twice to [1, 1].
|
||||
//
|
||||
// Skip the speculative graph build on the LAST iteration —
|
||||
// we'd never consume it, and on tiny `maxTokens` the saved
|
||||
// GPU/host work matters.
|
||||
var nextTokenArrN1: MLXArray? = nil
|
||||
var cacheN1: [(MLXArray, MLXArray)]? = nil
|
||||
if step + 1 < maxTokens {
|
||||
let nextEmbed = textDecoder.embedTokens(
|
||||
nextTokenArr.expandedDimensions(axis: 0).expandedDimensions(axis: 0)
|
||||
)
|
||||
let (hiddenN1, newCacheN1) = textDecoder(inputsEmbeds: nextEmbed, cache: cache)
|
||||
let lastHiddenN1 = hiddenN1[0..., (-1)..., .ellipsis]
|
||||
let logitsN1 = textDecoder.embedTokens.asLinear(lastHiddenN1)
|
||||
let argN1 = argMax(logitsN1, axis: -1).squeezed().asType(.int32)
|
||||
// Kick GPU on N+1 (chains after the asyncEval that is
|
||||
// still computing N).
|
||||
asyncEval(argN1, newCacheN1)
|
||||
nextTokenArrN1 = argN1
|
||||
cacheN1 = newCacheN1
|
||||
}
|
||||
|
||||
// Now sync N. By the time we reach `.item()` the GPU has
|
||||
// very likely finished N already, and the cost of this call
|
||||
// shrinks to a host-side memcpy of one int32 — which itself
|
||||
// overlaps with the in-flight N+1 forward pass.
|
||||
let nextToken = nextTokenArr.item(Int32.self)
|
||||
// Match the legacy loop semantics exactly: append the token
|
||||
// first, *then* break on EOS. The legacy loop emitted EOS
|
||||
// into `generatedTokens` whenever EOS was the most recent
|
||||
// pick, so greedy stays bit-identical.
|
||||
generatedTokens.append(nextToken)
|
||||
if nextToken == eosToken { break }
|
||||
|
||||
guard let advancedCache = cacheN1, let advancedToken = nextTokenArrN1 else {
|
||||
// Final iteration without speculative work — nothing to
|
||||
// advance to; the loop will exit naturally next.
|
||||
break
|
||||
}
|
||||
cache = advancedCache
|
||||
nextTokenArr = advancedToken
|
||||
}
|
||||
return generatedTokens
|
||||
}
|
||||
|
||||
/// Legacy decode loop kept verbatim for the non-greedy slow path.
|
||||
/// `pickNextToken` here pulls logits to CPU for repetition penalty,
|
||||
/// n-gram masking, and temperature sampling, so there's no benefit
|
||||
/// from `asyncEval` overlap.
|
||||
static func generateSlow(
|
||||
textDecoder: QuantizedTextModel,
|
||||
initialLogits: MLXArray,
|
||||
cache initialCache: [(MLXArray, MLXArray)],
|
||||
maxTokens: Int,
|
||||
options: Qwen3DecodingOptions
|
||||
) -> [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..<maxTokens {
|
||||
if nextToken == Int32(Qwen3ASRTokens.eosTokenId) { break }
|
||||
|
||||
let tokenEmbeds = textDecoder.embedTokens(
|
||||
MLXArray([nextToken]).expandedDimensions(axis: 0)
|
||||
)
|
||||
let (hiddenStates, newCache) = textDecoder(inputsEmbeds: tokenEmbeds, cache: cache)
|
||||
cache = newCache
|
||||
|
||||
let lastHiddenNext = hiddenStates[0..., (-1)..., .ellipsis]
|
||||
let logits = textDecoder.embedTokens.asLinear(lastHiddenNext)
|
||||
nextToken = Self.pickNextToken(
|
||||
logits: logits,
|
||||
generatedSoFar: generatedTokens,
|
||||
options: options
|
||||
)
|
||||
generatedTokens.append(nextToken)
|
||||
}
|
||||
return generatedTokens
|
||||
}
|
||||
|
||||
// MARK: - Decoder knobs
|
||||
|
||||
/// Pick the next token from a logits tensor, applying repetition
|
||||
/// penalty, no-repeat n-gram masking, and optional temperature sampling.
|
||||
///
|
||||
/// With default options (repetition=1.0, noRepeat=0, temperature=0) the
|
||||
/// result is the same `argMax` the decoder used pre-refactor.
|
||||
/// Implementation pulls logits to CPU (a 1-D Float array of vocab size)
|
||||
/// so we can manipulate entries in-place without fighting MLX indexing.
|
||||
///
|
||||
/// Access is `internal static` (not `private`) so
|
||||
/// ``Qwen3DecodingOptionsTests`` can exercise the sampler directly via
|
||||
/// ``@testable import Qwen3ASR`` — there is no GPU or model download
|
||||
/// involved so the path is trivially unit-testable once reachable.
|
||||
static func pickNextToken(
|
||||
logits: MLXArray,
|
||||
generatedSoFar: [Int32],
|
||||
options: Qwen3DecodingOptions
|
||||
) -> 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..<vocabSize {
|
||||
let u = Float.random(in: 1e-6...1.0)
|
||||
scores[i] = scores[i] / t - Float.log(-Float.log(u))
|
||||
}
|
||||
}
|
||||
|
||||
// Argmax of the adjusted scores.
|
||||
var bestIdx = 0
|
||||
var bestScore = -Float.infinity
|
||||
for i in 0..<vocabSize where scores[i] > 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")
|
||||
}
|
||||
}
|
||||
@@ -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<TranscriptionSegment, Error> {
|
||||
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..<end])
|
||||
let events = processor.process(samples: chunk)
|
||||
|
||||
for event in events {
|
||||
switch event {
|
||||
case .speechStarted(let time):
|
||||
speechStartSample = Int(time * 16000)
|
||||
lastPartialTime = time
|
||||
|
||||
case .speechEnded(let segment):
|
||||
if let startSample = speechStartSample {
|
||||
let endSample = min(Int(segment.endTime * 16000), samples.count)
|
||||
// After a force-split, startSample may equal endSample — skip empty spans
|
||||
guard startSample < endSample else {
|
||||
speechStartSample = nil
|
||||
continue
|
||||
}
|
||||
let segmentAudio = Array(samples[startSample..<endSample])
|
||||
let text = asrModel.transcribe(
|
||||
audio: segmentAudio, sampleRate: 16000,
|
||||
language: config.language, maxTokens: config.maxTokens,
|
||||
context: config.context)
|
||||
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
if !trimmed.isEmpty {
|
||||
continuation.yield(TranscriptionSegment(
|
||||
text: trimmed,
|
||||
startTime: segment.startTime,
|
||||
endTime: segment.endTime,
|
||||
isFinal: true,
|
||||
segmentIndex: segmentIndex))
|
||||
segmentIndex += 1
|
||||
}
|
||||
speechStartSample = nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 2: emit partial results during speech
|
||||
if config.emitPartialResults, let startSample = speechStartSample {
|
||||
let currentTime = processor.currentTime
|
||||
let speechStart = Float(startSample) / 16000
|
||||
let speechDuration = currentTime - speechStart
|
||||
|
||||
if currentTime - lastPartialTime >= config.partialResultInterval {
|
||||
let endSample = min(Int(currentTime * 16000), samples.count)
|
||||
guard startSample < endSample else {
|
||||
lastPartialTime = currentTime
|
||||
continue
|
||||
}
|
||||
let segmentAudio = Array(samples[startSample..<endSample])
|
||||
let text = asrModel.transcribe(
|
||||
audio: segmentAudio, sampleRate: 16000,
|
||||
language: config.language, maxTokens: config.maxTokens,
|
||||
context: config.context)
|
||||
let currentWords = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
.split(separator: " ").map(String.init)
|
||||
|
||||
if !currentWords.isEmpty {
|
||||
continuation.yield(TranscriptionSegment(
|
||||
text: currentWords.joined(separator: " "),
|
||||
startTime: speechStart,
|
||||
endTime: currentTime,
|
||||
isFinal: false,
|
||||
segmentIndex: segmentIndex))
|
||||
}
|
||||
lastPartialTime = currentTime
|
||||
}
|
||||
|
||||
// Force-split if speech exceeds maxSegmentDuration
|
||||
if speechDuration >= 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..<endSample])
|
||||
let text = asrModel.transcribe(
|
||||
audio: segmentAudio, sampleRate: 16000,
|
||||
language: config.language, maxTokens: config.maxTokens,
|
||||
context: config.context)
|
||||
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
if !trimmed.isEmpty {
|
||||
continuation.yield(TranscriptionSegment(
|
||||
text: trimmed,
|
||||
startTime: speechStart,
|
||||
endTime: currentTime,
|
||||
isFinal: true,
|
||||
segmentIndex: segmentIndex))
|
||||
segmentIndex += 1
|
||||
}
|
||||
speechStartSample = Int(currentTime * 16000)
|
||||
lastPartialTime = currentTime
|
||||
}
|
||||
} else if !config.emitPartialResults, let startSample = speechStartSample {
|
||||
// Force-split without partial results
|
||||
let currentTime = processor.currentTime
|
||||
let speechStart = Float(startSample) / 16000
|
||||
let speechDuration = currentTime - speechStart
|
||||
|
||||
if speechDuration >= config.maxSegmentDuration {
|
||||
let endSample = min(Int(currentTime * 16000), samples.count)
|
||||
guard startSample < endSample else {
|
||||
speechStartSample = Int(currentTime * 16000)
|
||||
continue
|
||||
}
|
||||
let segmentAudio = Array(samples[startSample..<endSample])
|
||||
let text = asrModel.transcribe(
|
||||
audio: segmentAudio, sampleRate: 16000,
|
||||
language: config.language, maxTokens: config.maxTokens,
|
||||
context: config.context)
|
||||
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
if !trimmed.isEmpty {
|
||||
continuation.yield(TranscriptionSegment(
|
||||
text: trimmed,
|
||||
startTime: speechStart,
|
||||
endTime: currentTime,
|
||||
isFinal: true,
|
||||
segmentIndex: segmentIndex))
|
||||
segmentIndex += 1
|
||||
}
|
||||
speechStartSample = Int(currentTime * 16000)
|
||||
}
|
||||
}
|
||||
|
||||
offset = end
|
||||
}
|
||||
|
||||
// Flush any remaining speech
|
||||
let flushEvents = processor.flush()
|
||||
for event in flushEvents {
|
||||
if case .speechEnded(let segment) = event, let startSample = speechStartSample {
|
||||
let endSample = min(Int(segment.endTime * 16000), samples.count)
|
||||
guard startSample < endSample else { continue }
|
||||
let segmentAudio = Array(samples[startSample..<endSample])
|
||||
let text = asrModel.transcribe(
|
||||
audio: segmentAudio, sampleRate: 16000,
|
||||
language: config.language, maxTokens: config.maxTokens,
|
||||
context: config.context)
|
||||
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
if !trimmed.isEmpty {
|
||||
continuation.yield(TranscriptionSegment(
|
||||
text: trimmed,
|
||||
startTime: segment.startTime,
|
||||
endTime: segment.endTime,
|
||||
isFinal: true,
|
||||
segmentIndex: segmentIndex))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
continuation.finish()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - LocalAgreement Helper
|
||||
|
||||
/// Returns the longest common prefix of two word arrays (case-insensitive).
|
||||
public func longestCommonPrefix(_ a: [String], _ b: [String]) -> [String] {
|
||||
var result: [String] = []
|
||||
for i in 0..<min(a.count, b.count) {
|
||||
if a[i].lowercased() == b[i].lowercased() {
|
||||
result.append(b[i])
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -0,0 +1,308 @@
|
||||
import Foundation
|
||||
import NaturalLanguage
|
||||
import AudioCommon
|
||||
|
||||
/// Result of preprocessing text for forced alignment
|
||||
public struct SlottedText: Sendable {
|
||||
/// Token IDs with timestamp tokens inserted around each word
|
||||
public let tokenIds: [Int]
|
||||
/// Indices within tokenIds that are timestamp tokens
|
||||
public let timestampPositions: [Int]
|
||||
/// Surface forms of the words (one per timestamp pair). Adjacent
|
||||
/// punctuation is preserved here so callers can reconstruct sentences;
|
||||
/// the model itself only sees the punctuation-stripped form.
|
||||
public let words: [String]
|
||||
}
|
||||
|
||||
/// Surface + cleaned form of a single word emitted by tokenization.
|
||||
/// `surface` keeps adjacent punctuation; `cleaned` is what the model tokenizer
|
||||
/// sees (letters, numbers, and combining marks only).
|
||||
struct WordPair: Sendable {
|
||||
var surface: String
|
||||
let cleaned: String
|
||||
}
|
||||
|
||||
/// Language-specific text preprocessing for forced alignment.
|
||||
///
|
||||
/// Three paths:
|
||||
/// - Japanese → morpheme-level segmentation
|
||||
/// - Korean → word-level segmentation
|
||||
/// - Other → whitespace split + per-Han-ideograph break
|
||||
///
|
||||
/// Tokens are filtered to keep only Unicode Letters / Numbers and the ASCII
|
||||
/// apostrophe — punctuation, symbols, and marks are stripped before
|
||||
/// timestamp slots are inserted. Han-ideograph splitting is restricted to
|
||||
/// CJK Unified + Extensions A–E + Compatibility; hiragana, katakana, and
|
||||
/// Hangul are NOT split per character (the model emits timestamp slots
|
||||
/// between morphemes for those scripts, not between every kana or jamo).
|
||||
///
|
||||
/// Punctuation that surrounds a word (commas, periods, brackets, CJK
|
||||
/// `,。!?` etc.) is preserved on the `surface` form so subtitle / SRT
|
||||
/// pipelines can still split on sentence boundaries after alignment.
|
||||
public enum TextPreprocessor {
|
||||
|
||||
/// Split text into words and insert timestamp slots for alignment.
|
||||
///
|
||||
/// For each word, inserts `<timestamp><timestamp>` 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<String.Index>] = []
|
||||
tokenizer.enumerateTokens(in: text.startIndex..<text.endIndex) { range, _ in
|
||||
ranges.append(range)
|
||||
return true
|
||||
}
|
||||
|
||||
var pairs: [WordPair] = []
|
||||
for (i, range) in ranges.enumerated() {
|
||||
let cleaned = cleanToken(String(text[range]))
|
||||
guard !cleaned.isEmpty else { continue }
|
||||
var surface = String(text[range])
|
||||
|
||||
let nextStart = i + 1 < ranges.count ? ranges[i + 1].lowerBound : text.endIndex
|
||||
var idx = range.upperBound
|
||||
while idx < nextStart {
|
||||
let ch = text[idx]
|
||||
if ch.isWhitespace { break }
|
||||
let allPunct = String(ch).unicodeScalars.allSatisfy { !isKeptScalar($0) }
|
||||
if !allPunct { break }
|
||||
surface.append(ch)
|
||||
idx = text.index(after: idx)
|
||||
}
|
||||
pairs.append(WordPair(surface: surface, cleaned: cleaned))
|
||||
}
|
||||
return pairs
|
||||
}
|
||||
|
||||
// MARK: - Default path (whitespace + per-Han break)
|
||||
|
||||
/// Whitespace split, then per-Han-ideograph break inside each segment.
|
||||
/// Default path for languages with reliable word-level whitespace —
|
||||
/// English, European languages, Hindi, Arabic, Vietnamese, Mongolian,
|
||||
/// Indonesian, etc. Chinese works because most Chinese text has no
|
||||
/// whitespace, so each ideograph becomes its own token via the per-Han
|
||||
/// break inside the (single) whitespace-bounded segment.
|
||||
static func tokenizeSpaceLangPairs(_ text: String) -> [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
|
||||
}
|
||||
}
|
||||
+145
@@ -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..<corrected.count {
|
||||
if corrected[i] < corrected[i - 1] {
|
||||
corrected[i] = corrected[i - 1]
|
||||
}
|
||||
}
|
||||
|
||||
return corrected
|
||||
}
|
||||
|
||||
/// Find positions of the Longest Increasing Subsequence (O(n log n))
|
||||
static func longestIncreasingSubsequencePositions(_ arr: [Int]) -> [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..<n {
|
||||
// Binary search for position to insert arr[i]
|
||||
var lo = 0, hi = tails.count
|
||||
while lo < hi {
|
||||
let mid = (lo + hi) / 2
|
||||
if tails[mid] < arr[i] {
|
||||
lo = mid + 1
|
||||
} else {
|
||||
hi = mid
|
||||
}
|
||||
}
|
||||
|
||||
if lo == tails.count {
|
||||
tails.append(arr[i])
|
||||
tailIndices.append(i)
|
||||
} else {
|
||||
tails[lo] = arr[i]
|
||||
tailIndices[lo] = i
|
||||
}
|
||||
|
||||
parent[i] = lo > 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
|
||||
}
|
||||
}
|
||||
@@ -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<String, MLXArray>] = [:]
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -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..<logits.count {
|
||||
if logits[i] > maxVal {
|
||||
maxVal = logits[i]
|
||||
maxIdx = i
|
||||
}
|
||||
}
|
||||
return maxIdx
|
||||
}
|
||||
|
||||
// Temperature scaling
|
||||
if config.temperature != 1.0 {
|
||||
for i in 0..<logits.count {
|
||||
logits[i] /= config.temperature
|
||||
}
|
||||
}
|
||||
|
||||
// Softmax
|
||||
let maxLogit = logits.max() ?? 0
|
||||
var probs = logits.map { exp($0 - maxLogit) }
|
||||
let sum = probs.reduce(0, +)
|
||||
probs = probs.map { $0 / sum }
|
||||
|
||||
// Top-K filtering
|
||||
if config.topK > 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..<probs.count {
|
||||
if !mask[i] { probs[i] = 0 }
|
||||
}
|
||||
let filteredSum = probs.reduce(0, +)
|
||||
if filteredSum > 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
|
||||
}
|
||||
}
|
||||
@@ -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 // <think>
|
||||
static let thinkEndId = 248069 // </think>
|
||||
static let newlineId = 198 // \n
|
||||
|
||||
/// Strip thinking block from generated tokens.
|
||||
///
|
||||
/// Removes tokens from `<think>` through `</think>` (inclusive)
|
||||
/// and any trailing newlines, returning only the response content.
|
||||
static func stripThinking(from tokens: [Int]) -> [Int] {
|
||||
let thinkTokens: Set<Int> = [thinkStartId, thinkEndId]
|
||||
let newlines: Set<Int> = [newlineId, 271] // 198 = \n, 271 = \n\n
|
||||
|
||||
guard let startIdx = tokens.firstIndex(where: { thinkTokens.contains($0) && $0 == thinkStartId }) else {
|
||||
// No <think> — strip any leading </think> + 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..<startIdx]) + Array(tokens[afterThink...])
|
||||
}
|
||||
return Array(tokens[0..<startIdx])
|
||||
}
|
||||
|
||||
/// Encode a conversation into token IDs using Qwen3.5 chat template.
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - config: Model config (for future extensibility)
|
||||
/// - enableThinking: If false, injects empty think block to skip reasoning
|
||||
static func encode(
|
||||
messages: [ChatMessage],
|
||||
tokenizer: ChatTokenizer,
|
||||
config: Qwen3ChatConfig? = nil,
|
||||
addGenerationPrompt: Bool = true,
|
||||
enableThinking: Bool = true
|
||||
) -> [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
|
||||
}
|
||||
}
|
||||
@@ -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<UInt8>] = [
|
||||
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("|>")
|
||||
}
|
||||
}
|
||||
@@ -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<Int> = [
|
||||
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<mach_task_basic_info>.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<String, Error> {
|
||||
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..<sampling.maxTokens {
|
||||
let nextToken = ChatSampler.sample(
|
||||
logits: logits,
|
||||
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)
|
||||
}
|
||||
|
||||
let tokenArr = MLXArray([Int32(nextToken)])
|
||||
.expandedDimensions(axis: 0)
|
||||
let (stepLogits, newState) = self.model.forward(
|
||||
inputIds: tokenArr, state: self.state)
|
||||
eval(stepLogits)
|
||||
self.state = newState
|
||||
logits = self.extractLastPositionLogits(stepLogits)
|
||||
}
|
||||
|
||||
continuation.finish()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Helpers
|
||||
|
||||
/// Extract logits for the last sequence position as a Float array.
|
||||
private func extractLastPositionLogits(_ logits: MLXArray) -> [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()
|
||||
}
|
||||
}
|
||||
@@ -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<mach_task_basic_info>.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..<sampling.maxTokens {
|
||||
let nextToken = ChatSampler.sample(
|
||||
logits: lastLogits, config: sampling,
|
||||
previousTokens: promptTokens + generatedTokens)
|
||||
|
||||
if nextToken == config.eosTokenId { break }
|
||||
if nextToken == ChatTemplate.imEndId { break }
|
||||
|
||||
generatedTokens.append(nextToken)
|
||||
if generatedTokens.count <= 10 {
|
||||
print("[CoreML] token[\(generatedTokens.count-1)]: \(nextToken) '\(tokenizer.decodeToken(nextToken) ?? "?")'")
|
||||
}
|
||||
lastLogits = try forwardStep(tokenId: nextToken)
|
||||
}
|
||||
|
||||
_decodeMs = (CFAbsoluteTimeGetCurrent() - decodeStart) * 1000
|
||||
_decodeTokens = generatedTokens.count
|
||||
|
||||
let tps = _decodeMs > 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<String, Error> {
|
||||
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<Int> = [
|
||||
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..<maxSeqLen {
|
||||
maskPtr[i] = i <= position ? 0 : -Float.greatestFiniteMagnitude
|
||||
}
|
||||
|
||||
// Position
|
||||
let posArray = try MLMultiArray(shape: [1], dataType: .int32)
|
||||
posArray[0] = NSNumber(value: Int32(position))
|
||||
|
||||
// Decoder forward
|
||||
let decoderInput = try MLDictionaryFeatureProvider(dictionary: [
|
||||
"input_embeds": MLFeatureValue(multiArray: embedding),
|
||||
"position": MLFeatureValue(multiArray: posArray),
|
||||
"attention_mask": MLFeatureValue(multiArray: mask),
|
||||
])
|
||||
|
||||
let result = try decoderModel.prediction(
|
||||
from: decoderInput, using: decoderState)
|
||||
|
||||
position += 1
|
||||
|
||||
// Extract logits
|
||||
guard let logitsArray = result.featureValue(for: "logits")?.multiArrayValue else {
|
||||
throw ChatModelError.inferenceFailed("Decoder output missing")
|
||||
}
|
||||
|
||||
let vocabSize = config.vocabSize
|
||||
|
||||
// Debug: log dtype and shape on first few calls
|
||||
if position <= 3 {
|
||||
print("[CoreML] pos=\(position) logits shape=\(logitsArray.shape) dtype=\(logitsArray.dataType.rawValue) count=\(logitsArray.count)")
|
||||
}
|
||||
|
||||
// Handle both Float16 and Float32 output
|
||||
let logits: [Float]
|
||||
if logitsArray.dataType == .float32 {
|
||||
let ptr = logitsArray.dataPointer.bindMemory(to: Float.self, capacity: vocabSize)
|
||||
logits = Array(UnsafeBufferPointer(start: ptr, count: vocabSize))
|
||||
} else {
|
||||
let ptr = logitsArray.dataPointer.bindMemory(to: Float16.self, capacity: vocabSize)
|
||||
logits = (0..<vocabSize).map { Float(ptr[$0]) }
|
||||
}
|
||||
return logits
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Memory Management
|
||||
|
||||
extension Qwen35CoreMLChat: ModelMemoryManageable {
|
||||
public var isLoaded: Bool { true }
|
||||
public func unload() { /* CoreML manages its own memory */ }
|
||||
public var memoryFootprint: Int { 500 * 1024 * 1024 } // ~500 MB estimate
|
||||
}
|
||||
@@ -0,0 +1,710 @@
|
||||
import Foundation
|
||||
import MLXCommon
|
||||
import MLX
|
||||
import MLXNN
|
||||
import MLXFast
|
||||
|
||||
// MARK: - DeltaNet Linear Attention
|
||||
|
||||
/// DeltaNet linear attention layer for Qwen3.5 hybrid model.
|
||||
///
|
||||
/// Uses linear attention (no softmax) with a recurrent state matrix S of shape [B, H, D, D].
|
||||
/// The state evolves per-token: S = alpha * S + beta * (v outer k), where alpha/beta
|
||||
/// are learned per-head scalar gates derived from the input via softplus/sigmoid.
|
||||
///
|
||||
/// A causal conv1d (kernel=4) provides short-range local context before the attention.
|
||||
/// Output is gated: `o_proj(attention_output * silu(z))` where both are 2*hiddenSize = numHeads*headDim.
|
||||
///
|
||||
/// Weight shapes (HuggingFace safetensors):
|
||||
/// - in_proj_qkv.weight: [6144, 1024] (3 * 16 * 128)
|
||||
/// - in_proj_z.weight: [2048, 1024] (2 * hiddenSize, for gate)
|
||||
/// - in_proj_b.weight: [16, 1024] (beta gate, per head)
|
||||
/// - in_proj_a.weight: [16, 1024] (alpha gate, per head)
|
||||
/// - conv1d.weight: [6144, 4, 1] (depthwise causal conv, MLX [C, K, 1] format)
|
||||
/// - dt_bias: [16] (time-step bias)
|
||||
/// - A_log: [16] (log of decay rate)
|
||||
/// - norm.weight: [128] (per-head RMSNorm)
|
||||
/// - out_proj.weight: [1024, 2048] (gated output projection)
|
||||
public final class DeltaNetLayer: Module {
|
||||
let numHeads: Int
|
||||
let headDim: Int
|
||||
let hiddenSize: Int
|
||||
let convKernel: Int
|
||||
let qkvDim: Int
|
||||
|
||||
@ModuleInfo(key: "in_proj_qkv") var inProjQKV: QuantizedLinear
|
||||
@ModuleInfo(key: "in_proj_z") var inProjZ: QuantizedLinear
|
||||
@ModuleInfo(key: "in_proj_b") var inProjB: QuantizedLinear
|
||||
@ModuleInfo(key: "in_proj_a") var inProjA: QuantizedLinear
|
||||
|
||||
/// Conv1d weight: [C, 1, K] depthwise convolution applied to QKV before attention.
|
||||
/// Stored under a flat key to avoid nested key path issues in MLXNN module traversal.
|
||||
/// Weight loading applies this directly via `layer.convWeight = ...`.
|
||||
@ParameterInfo(key: "conv1d_weight") var convWeight: MLXArray
|
||||
|
||||
@ParameterInfo(key: "dt_bias") var dtBias: MLXArray
|
||||
@ParameterInfo(key: "A_log") var aLog: MLXArray
|
||||
|
||||
@ModuleInfo var norm: RMSNorm
|
||||
|
||||
@ModuleInfo(key: "out_proj") var outProj: QuantizedLinear
|
||||
|
||||
public init(config: Qwen3ChatConfig) {
|
||||
self.numHeads = config.linearNumKeyHeads ?? 16
|
||||
self.headDim = config.linearKeyHeadDim ?? 128
|
||||
self.hiddenSize = config.hiddenSize
|
||||
self.convKernel = config.linearConvKernelDim ?? 4
|
||||
self.qkvDim = 3 * numHeads * headDim
|
||||
|
||||
let groupSize = 64
|
||||
let bits = 4
|
||||
self._inProjQKV = ModuleInfo(wrappedValue: QuantizedLinear(hiddenSize, qkvDim, bias: false, groupSize: groupSize, bits: bits))
|
||||
self._inProjZ = ModuleInfo(wrappedValue: QuantizedLinear(hiddenSize, 2 * hiddenSize, bias: false, groupSize: groupSize, bits: bits))
|
||||
self._inProjB = ModuleInfo(wrappedValue: QuantizedLinear(hiddenSize, numHeads, bias: false, groupSize: groupSize, bits: bits))
|
||||
self._inProjA = ModuleInfo(wrappedValue: QuantizedLinear(hiddenSize, numHeads, bias: false, groupSize: groupSize, bits: bits))
|
||||
|
||||
self._convWeight = ParameterInfo(
|
||||
wrappedValue: MLXArray.zeros([qkvDim, convKernel, 1]))
|
||||
|
||||
self._dtBias = ParameterInfo(wrappedValue: MLXArray.zeros([numHeads]))
|
||||
self._aLog = ParameterInfo(wrappedValue: MLXArray.zeros([numHeads]))
|
||||
|
||||
self._norm = ModuleInfo(
|
||||
wrappedValue: RMSNorm(dimensions: headDim, eps: Float(config.rmsNormEps)))
|
||||
|
||||
self._outProj = ModuleInfo(wrappedValue: QuantizedLinear(2 * hiddenSize, hiddenSize, bias: false, groupSize: groupSize, bits: bits))
|
||||
|
||||
super.init()
|
||||
}
|
||||
|
||||
/// Recurrent state for a DeltaNet layer.
|
||||
public struct State {
|
||||
/// Recurrent state matrix [B, H, D, D]
|
||||
var s: MLXArray
|
||||
/// Conv1d ring buffer [B, C, K-1] storing last K-1 inputs
|
||||
var convState: MLXArray
|
||||
|
||||
public static func initial(
|
||||
batchSize: Int, numHeads: Int, headDim: Int,
|
||||
qkvDim: Int, convKernel: Int, dtype: DType = .float32
|
||||
) -> 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..., ..<hd].reshaped(b, t, numHeads, headDim)
|
||||
var k = qkvActivated[0..., 0..., hd..<(2 * hd)].reshaped(b, t, numHeads, headDim)
|
||||
let v = qkvActivated[0..., 0..., (2 * hd)...].reshaped(b, t, numHeads, headDim)
|
||||
|
||||
// Q/K normalization (reference: inv_scale * rms_norm, different scaling for Q and K)
|
||||
// rms_norm(x, None, eps) = x / sqrt(mean(x^2) + eps)
|
||||
// q = inv_scale^2 * rms_norm(q) where inv_scale = head_dim^(-0.5)
|
||||
// k = inv_scale * rms_norm(k)
|
||||
let invScale = Float(1.0) / sqrt(Float(headDim))
|
||||
q = MLXArray(invScale * invScale) * rmsNormNoWeight(q)
|
||||
k = MLXArray(invScale) * rmsNormNoWeight(k)
|
||||
|
||||
// Compute gating: g = exp(-exp(A_log) * softplus(a + dt_bias))
|
||||
let g = computeDecayGate(aRaw: aRaw) // [B, T, H]
|
||||
// beta = sigmoid(b_raw) (independent learned gate, NOT 1-alpha)
|
||||
let beta = sigmoid(bRaw) // [B, T, H]
|
||||
|
||||
// Sequential gated delta rule recurrence
|
||||
var currentS: MLXArray
|
||||
if let s = state {
|
||||
currentS = s.s
|
||||
} else {
|
||||
currentS = MLXArray.zeros([b, numHeads, headDim, headDim], dtype: x.dtype)
|
||||
}
|
||||
|
||||
var outputSteps: [MLXArray] = []
|
||||
outputSteps.reserveCapacity(t)
|
||||
|
||||
for step in 0..<t {
|
||||
// Extract step: [B, H, D] or [B, H]
|
||||
let qStep = q[0..., step..<(step + 1), 0..., 0...].squeezed(axis: 1) // [B, H, D]
|
||||
let kStep = k[0..., step..<(step + 1), 0..., 0...].squeezed(axis: 1) // [B, H, D]
|
||||
let vStep = v[0..., step..<(step + 1), 0..., 0...].squeezed(axis: 1) // [B, H, D]
|
||||
let gStep = g[0..., step..<(step + 1), 0...].squeezed(axis: 1) // [B, H]
|
||||
let betaStep = beta[0..., step..<(step + 1), 0...].squeezed(axis: 1) // [B, H]
|
||||
|
||||
// 1. Decay: S = g * S (g is scalar per-head: [B, H, 1, 1])
|
||||
let decay = gStep.reshaped(b, numHeads, 1, 1)
|
||||
currentS = currentS * decay
|
||||
|
||||
// 2. Error correction:
|
||||
// kv_mem = (S * k[..., None, :]).sum(-1) → [B, H, Dv]
|
||||
// delta = (v - kv_mem) * beta[..., None] → [B, H, Dv]
|
||||
let kExpanded = kStep.expandedDimensions(axis: -2) // [B, H, 1, Dk]
|
||||
let kvMem = (currentS * kExpanded).sum(axis: -1) // [B, H, Dv]
|
||||
let delta = (vStep - kvMem) * betaStep.expandedDimensions(axis: -1) // [B, H, Dv]
|
||||
|
||||
// 3. Update: S = S + k[..., None, :] * delta[..., None]
|
||||
// k: [B, H, Dk] → [B, H, 1, Dk], delta: [B, H, Dv] → [B, H, Dv, 1]
|
||||
currentS = currentS + kExpanded * delta.expandedDimensions(axis: -1)
|
||||
|
||||
// 4. Output: y = (S * q[..., None, :]).sum(-1) → [B, H, Dv]
|
||||
let qExpanded = qStep.expandedDimensions(axis: -2) // [B, H, 1, Dk]
|
||||
let oStep = (currentS * qExpanded).sum(axis: -1) // [B, H, Dv]
|
||||
outputSteps.append(oStep)
|
||||
}
|
||||
|
||||
// Stack: [B, T, H, D]
|
||||
let output = stacked(outputSteps, axis: 1)
|
||||
|
||||
// RMSNormGated: norm(output) * silu(z)
|
||||
// z has shape [B, T, 2*hiddenSize=2048], reshape to [B, T, H, D] for per-head norm
|
||||
let zReshaped = zRaw.reshaped(b, t, numHeads, headDim)
|
||||
let normedOutput = norm(output) // per-head RMSNorm, [B, T, H, D]
|
||||
let gated = normedOutput * silu(zReshaped) // [B, T, H, D]
|
||||
|
||||
// Reshape to [B, T, H*D=2048] and project to hiddenSize
|
||||
let result = outProj(gated.reshaped(b, t, numHeads * headDim)) // [B, T, 1024]
|
||||
|
||||
return (result, State(s: currentS, convState: newConvState))
|
||||
}
|
||||
|
||||
/// Compute decay gate: g = exp(-exp(A_log) * softplus(a + dt_bias))
|
||||
private func computeDecayGate(aRaw: MLXArray) -> 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..<outputLen {
|
||||
windows.append(input[0..., 0..., t..<(t + k)]) // [B, C, K]
|
||||
}
|
||||
let unfolded = stacked(windows, axis: 2) // [B, C, T, K]
|
||||
|
||||
// Kernel: [C, K, 1] -> 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..<sampling.maxTokens {
|
||||
let input = MLXArray([Int32(token)]).expandedDimensions(axis: 0)
|
||||
let (logits, newState) = forward(inputIds: input, state: state)
|
||||
state = newState
|
||||
eval(logits)
|
||||
|
||||
token = sampleFromLogits(logits, at: 0,
|
||||
config: sampling, history: promptIds + generated)
|
||||
if token == config.eosTokenId { break }
|
||||
generated.append(token)
|
||||
}
|
||||
|
||||
return generated
|
||||
}
|
||||
|
||||
// MARK: - Helpers
|
||||
|
||||
private func sampleFromLogits(
|
||||
_ logits: MLXArray, at position: Int,
|
||||
config: ChatSamplingConfig, history: [Int]
|
||||
) -> 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)
|
||||
}
|
||||
}
|
||||
@@ -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<String, Error>
|
||||
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
|
||||
}
|
||||
}
|
||||
+226
@@ -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..<numLayers {
|
||||
let prefix = "layers.\(i)"
|
||||
let layer = model.layers[i]
|
||||
|
||||
// Layer norms
|
||||
CommonWeightLoader.applyRMSNormWeights(
|
||||
to: layer.inputLayerNorm,
|
||||
prefix: "\(prefix).input_layernorm",
|
||||
from: modelWeights)
|
||||
CommonWeightLoader.applyRMSNormWeights(
|
||||
to: layer.postAttentionLayerNorm,
|
||||
prefix: "\(prefix).post_attention_layernorm",
|
||||
from: modelWeights)
|
||||
|
||||
// MLP
|
||||
applyQuantizedMLPWeights(
|
||||
to: layer.mlp,
|
||||
prefix: "\(prefix).mlp",
|
||||
from: modelWeights)
|
||||
|
||||
// Attention (type-specific, different key prefix per HuggingFace convention)
|
||||
if layer.layerType == "linear_attention" {
|
||||
try applyDeltaNetWeights(
|
||||
to: layer.deltaNet!,
|
||||
prefix: "\(prefix).linear_attn",
|
||||
from: modelWeights)
|
||||
} else {
|
||||
applyGatedAttentionWeights(
|
||||
to: layer.gatedAttn!,
|
||||
prefix: "\(prefix).self_attn",
|
||||
from: modelWeights)
|
||||
}
|
||||
|
||||
let pct = 0.5 + 0.45 * Double(i + 1) / Double(numLayers)
|
||||
progressHandler?(pct, "Layer \(i + 1)/\(numLayers)")
|
||||
}
|
||||
|
||||
// Evaluate all parameters
|
||||
eval(model)
|
||||
progressHandler?(1.0, "Weights loaded")
|
||||
}
|
||||
|
||||
// MARK: - DeltaNet Weight Loading
|
||||
|
||||
/// Apply weights to a DeltaNet (linear attention) layer.
|
||||
///
|
||||
/// All DeltaNet projections are quantized INT4 (matching mlx-community format).
|
||||
/// The conv1d weight and scalar parameters (dt_bias, A_log) are loaded directly.
|
||||
private static func applyDeltaNetWeights(
|
||||
to layer: DeltaNetLayer,
|
||||
prefix: String,
|
||||
from weights: [String: MLXArray]
|
||||
) throws {
|
||||
// Quantized projections
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: layer.inProjQKV, prefix: "\(prefix).in_proj_qkv", from: weights)
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: layer.inProjZ, prefix: "\(prefix).in_proj_z", from: weights)
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: layer.inProjB, prefix: "\(prefix).in_proj_b", from: weights)
|
||||
CommonWeightLoader.applyQuantizedLinearWeights(
|
||||
to: layer.inProjA, prefix: "\(prefix).in_proj_a", from: weights)
|
||||
|
||||
// Raw parameters: conv1d weight, dt_bias, A_log
|
||||
// Note: ParameterInfo key from property declaration is lost when reassigning
|
||||
// in init(), so use Swift property names (convWeight, dtBias, aLog).
|
||||
var rawParams: [String: NestedItem<String, MLXArray>] = [:]
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
)
|
||||
}
|
||||
+64
@@ -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..<ptr.count {
|
||||
ptr[i] = 0
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
+76
@@ -0,0 +1,76 @@
|
||||
#if canImport(CoreML)
|
||||
import AudioCommon
|
||||
import CoreML
|
||||
import Foundation
|
||||
|
||||
extension WeSpeakerModel {
|
||||
|
||||
/// Run CoreML inference to extract a 256-dim speaker embedding.
|
||||
///
|
||||
/// Pads the mel spectrogram to the nearest enumerated shape (required by
|
||||
/// the CoreML model which uses EnumeratedShapes for the time dimension),
|
||||
/// runs prediction, and extracts the embedding.
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - melSpec: flat Float array of log-mel features `[nFrames * 80]`
|
||||
/// - nFrames: number of mel frames
|
||||
/// - Returns: 256-dim L2-normalized speaker embedding
|
||||
func embedCoreML(melSpec: [Float], nFrames: Int) throws -> [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..<copyCount {
|
||||
melPtr[i] = Float16(melSpec[i])
|
||||
}
|
||||
// Zero-pad remaining frames
|
||||
let totalElements = targetLength * 80
|
||||
for i in copyCount..<totalElements {
|
||||
melPtr[i] = 0
|
||||
}
|
||||
|
||||
let input = try MLDictionaryFeatureProvider(dictionary: [
|
||||
"mel": MLFeatureValue(multiArray: melArray),
|
||||
])
|
||||
|
||||
let result = try model.prediction(from: input)
|
||||
|
||||
// Extract "embedding" output: [1, 256]
|
||||
guard let embArray = result.featureValue(for: "embedding")?.multiArrayValue else {
|
||||
throw AudioModelError.inferenceFailed(
|
||||
operation: "SpeakerEmbedding", reason: "Missing 'embedding' output")
|
||||
}
|
||||
|
||||
// Read 256 float16 values
|
||||
var embedding = [Float](repeating: 0, count: 256)
|
||||
let embPtr = embArray.dataPointer.assumingMemoryBound(to: Float16.self)
|
||||
for i in 0..<256 {
|
||||
embedding[i] = Float(embPtr[i])
|
||||
}
|
||||
|
||||
// L2 normalize
|
||||
let norm = sqrt(embedding.reduce(Float(0)) { $0 + $1 * $1 })
|
||||
if norm > 1e-10 {
|
||||
for i in 0..<256 { embedding[i] /= norm }
|
||||
}
|
||||
|
||||
return embedding
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -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 <file> 1 <start> <dur> <NA> <NA> <speaker> <NA> <NA>`
|
||||
public var rttmLine: String {
|
||||
let s = String(format: "%.3f", startTime)
|
||||
let d = String(format: "%.3f", duration)
|
||||
return "SPEAKER \(filename) 1 \(s) \(d) <NA> <NA> \(speakerLabel) <NA> <NA>"
|
||||
}
|
||||
}
|
||||
|
||||
/// 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)..<min(numFrames, startFrame + collarFrames) {
|
||||
collarMask[f] = true
|
||||
}
|
||||
for f in max(0, endFrame - collarFrames)..<min(numFrames, endFrame + collarFrames) {
|
||||
collarMask[f] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build per-frame speaker sets for reference and hypothesis
|
||||
// Use sorted arrays of speaker IDs (faster than Set for small N)
|
||||
let refSpeakers = buildFrameSpeakers(segments: reference, numFrames: numFrames, resolution: resolution)
|
||||
let hypSpeakers = buildFrameSpeakers(segments: hypothesis, numFrames: numFrames, resolution: resolution)
|
||||
|
||||
// Score each frame
|
||||
var totalSpeech: Float = 0
|
||||
var falseAlarm: Float = 0
|
||||
var missedSpeech: Float = 0
|
||||
var confusion: Float = 0
|
||||
|
||||
for f in 0..<numFrames {
|
||||
if collarMask[f] { continue }
|
||||
|
||||
let refCount = refSpeakers[f].count
|
||||
let hypCount = hypSpeakers[f].count
|
||||
|
||||
if refCount == 0 && hypCount == 0 { continue }
|
||||
|
||||
if refCount == 0 {
|
||||
// No reference speech, any hypothesis is false alarm
|
||||
falseAlarm += Float(hypCount) * resolution
|
||||
continue
|
||||
}
|
||||
|
||||
// Reference speech exists — count it
|
||||
totalSpeech += Float(refCount) * resolution
|
||||
|
||||
if hypCount == 0 {
|
||||
// All reference speech is missed
|
||||
missedSpeech += Float(refCount) * resolution
|
||||
continue
|
||||
}
|
||||
|
||||
// Both have speakers — compute exact match by speaker ID
|
||||
let matched = countExactMatched(ref: refSpeakers[f], hyp: hypSpeakers[f])
|
||||
let unmatchedRef = refCount - matched // ref speakers with no matching hyp
|
||||
let unmatchedHyp = hypCount - matched // hyp speakers with no matching ref
|
||||
// Confusion: min of unmatched ref/hyp (wrong speaker assigned)
|
||||
let conf = min(unmatchedRef, unmatchedHyp)
|
||||
// Missed: unmatched ref beyond confusion
|
||||
let missed = unmatchedRef - conf
|
||||
// False alarm: unmatched hyp beyond confusion
|
||||
let fa = unmatchedHyp - conf
|
||||
|
||||
missedSpeech += Float(missed) * resolution
|
||||
confusion += Float(conf) * resolution
|
||||
falseAlarm += Float(fa) * resolution
|
||||
}
|
||||
|
||||
return DERResult(
|
||||
totalSpeech: totalSpeech,
|
||||
falseAlarm: falseAlarm,
|
||||
missedSpeech: missedSpeech,
|
||||
confusion: confusion
|
||||
)
|
||||
}
|
||||
|
||||
// MARK: - Optimal Speaker Mapping
|
||||
|
||||
/// Compute DER with optimal 1-to-1 speaker mapping between reference and hypothesis.
|
||||
///
|
||||
/// Tries all permutations of hypothesis speaker labels to find the mapping
|
||||
/// that minimizes DER. For >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..<endFrame {
|
||||
if !result[f].contains(seg.speakerId) {
|
||||
result[f].append(seg.speakerId)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
/// Count speakers present in both ref and hyp by exact ID match.
|
||||
private func countExactMatched(ref: [Int], hyp: [Int]) -> 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..<hypSpeakers.count))
|
||||
var bestResult: DERResult?
|
||||
|
||||
for perm in permutations {
|
||||
// Build mapping: hypSpeakers[i] → refSpeakers[perm[i]] (if in range)
|
||||
var mapping = [Int: Int]()
|
||||
for (i, p) in perm.enumerated() {
|
||||
if p < refSpeakers.count {
|
||||
mapping[hypSpeakers[i]] = refSpeakers[p]
|
||||
} else {
|
||||
mapping[hypSpeakers[i]] = 1000 + i // unmapped → unique ID
|
||||
}
|
||||
}
|
||||
|
||||
let remapped = hypothesis.map { seg in
|
||||
DiarizedSegment(
|
||||
startTime: seg.startTime,
|
||||
endTime: seg.endTime,
|
||||
speakerId: mapping[seg.speakerId] ?? seg.speakerId
|
||||
)
|
||||
}
|
||||
|
||||
let result = computeDER(reference: reference, hypothesis: remapped,
|
||||
collar: collar, resolution: resolution)
|
||||
|
||||
if bestResult == nil || result.der < bestResult!.der {
|
||||
bestResult = result
|
||||
}
|
||||
}
|
||||
|
||||
return bestResult!
|
||||
}
|
||||
|
||||
private func greedyOptimalMapping(
|
||||
reference: [DiarizedSegment],
|
||||
hypothesis: [DiarizedSegment],
|
||||
refSpeakers: [Int],
|
||||
hypSpeakers: [Int],
|
||||
collar: Float,
|
||||
resolution: Float
|
||||
) -> 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..<numFrames {
|
||||
for rSpk in refFrames[f] {
|
||||
guard let ri = refIndex[rSpk] else { continue }
|
||||
for hSpk in hypFrames[f] {
|
||||
guard let hi = hypIndex[hSpk] else { continue }
|
||||
overlap[ri][hi] += 1
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Greedy match: pick highest overlap pair, assign, repeat
|
||||
var mapping = [Int: Int]()
|
||||
var usedRef = Set<Int>()
|
||||
var usedHyp = Set<Int>()
|
||||
|
||||
for _ in 0..<min(refSpeakers.count, hypSpeakers.count) {
|
||||
var bestR = -1, bestH = -1, bestOverlap = -1
|
||||
for r in 0..<refSpeakers.count where !usedRef.contains(r) {
|
||||
for h in 0..<hypSpeakers.count where !usedHyp.contains(h) {
|
||||
if overlap[r][h] > 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..<elements.count
|
||||
// For speaker mapping, we need arrangements: pick from 0..<max(ref,hyp) count
|
||||
let n = elements.count
|
||||
func permute(_ current: [Int], _ remaining: [Int]) {
|
||||
if current.count == n {
|
||||
result.append(current)
|
||||
return
|
||||
}
|
||||
for (i, elem) in remaining.enumerated() {
|
||||
var rest = remaining
|
||||
rest.remove(at: i)
|
||||
permute(current + [elem], rest)
|
||||
}
|
||||
}
|
||||
|
||||
// Permute indices 0..<n (mapping hyp speakers to ref speaker slots)
|
||||
let indices = Array(0..<max(n, n))
|
||||
permute([], indices)
|
||||
|
||||
return result
|
||||
}
|
||||
+183
@@ -0,0 +1,183 @@
|
||||
import Foundation
|
||||
import AudioCommon
|
||||
|
||||
/// Shared helpers for diarization post-processing, used by both
|
||||
/// PyannoteDiarizationPipeline and SortformerDiarizer.
|
||||
enum DiarizationHelpers {
|
||||
|
||||
/// Merge adjacent segments from the same speaker when the gap is below `minSilence`.
|
||||
///
|
||||
/// Segments are grouped per-speaker, merged within each group, then sorted globally.
|
||||
static func mergeSegments(
|
||||
_ segments: [DiarizedSegment],
|
||||
minSilence: Float
|
||||
) -> [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..<sorted.count {
|
||||
let next = sorted[i]
|
||||
if next.startTime - current.endTime < minSilence {
|
||||
current = DiarizedSegment(
|
||||
startTime: current.startTime,
|
||||
endTime: next.endTime,
|
||||
speakerId: spk
|
||||
)
|
||||
} else {
|
||||
merged.append(current)
|
||||
current = next
|
||||
}
|
||||
}
|
||||
merged.append(current)
|
||||
}
|
||||
|
||||
merged.sort { $0.startTime < $1.startTime }
|
||||
return merged
|
||||
}
|
||||
|
||||
/// Remap speaker IDs to contiguous 0-based range, preserving order of first appearance.
|
||||
static func compactSpeakerIds(_ segments: [DiarizedSegment]) -> [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..<n) // item → cluster ID
|
||||
var centroids = items.map { $0.embedding } // cluster ID → centroid
|
||||
var clusterMembers = (0..<n).map { [$0] } // cluster ID → member items
|
||||
// Window indices per cluster (for constraint checking)
|
||||
var clusterWindows = items.map { Set([$0.windowIndex]) }
|
||||
var active = Set(0..<n)
|
||||
|
||||
while active.count > 1 {
|
||||
// Find closest unconstrained pair
|
||||
var bestDist: Float = Float.greatestFiniteMagnitude
|
||||
var bestI = -1, bestJ = -1
|
||||
|
||||
let activeList = active.sorted()
|
||||
for ai in 0..<activeList.count {
|
||||
for aj in (ai + 1)..<activeList.count {
|
||||
let ci = activeList[ai], cj = activeList[aj]
|
||||
|
||||
// Same-window constraint: if clusters share any window, skip
|
||||
if !clusterWindows[ci].isDisjoint(with: clusterWindows[cj]) {
|
||||
continue
|
||||
}
|
||||
|
||||
let dist = cosineDistance(centroids[ci], centroids[cj])
|
||||
if dist < bestDist {
|
||||
bestDist = dist
|
||||
bestI = ci
|
||||
bestJ = cj
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
guard bestDist < threshold && bestI >= 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..<dim {
|
||||
newCentroid[d] = (centroids[bestI][d] * Float(sizeI) + centroids[bestJ][d] * Float(sizeJ)) / totalSize
|
||||
}
|
||||
centroids[bestI] = newCentroid
|
||||
|
||||
// Transfer members
|
||||
for member in clusterMembers[bestJ] {
|
||||
clusterOf[member] = bestI
|
||||
}
|
||||
clusterMembers[bestI].append(contentsOf: clusterMembers[bestJ])
|
||||
|
||||
// Propagate window constraints
|
||||
clusterWindows[bestI].formUnion(clusterWindows[bestJ])
|
||||
|
||||
active.remove(bestJ)
|
||||
}
|
||||
|
||||
// Build final compact assignment
|
||||
let activeSorted = active.sorted()
|
||||
var clusterMap = [Int: Int]() // old cluster ID → new compact ID
|
||||
for (newId, oldId) in activeSorted.enumerated() {
|
||||
clusterMap[oldId] = newId
|
||||
}
|
||||
|
||||
let assignment = (0..<n).map { clusterMap[clusterOf[$0]]! }
|
||||
let finalCentroids = activeSorted.map { centroids[$0] }
|
||||
|
||||
return (assignment, finalCentroids)
|
||||
}
|
||||
|
||||
/// Cosine distance between two vectors: 1 - cosine_similarity.
|
||||
/// Returns value in [0, 2].
|
||||
static func cosineDistance(_ a: [Float], _ b: [Float]) -> 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..<n {
|
||||
dot += a[i] * b[i]
|
||||
normA += a[i] * a[i]
|
||||
normB += b[i] * b[i]
|
||||
}
|
||||
|
||||
let denom = sqrt(normA) * sqrt(normB)
|
||||
guard denom > 1e-10 else { return 2.0 }
|
||||
return 1.0 - dot / denom
|
||||
}
|
||||
}
|
||||
+570
@@ -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..<end])
|
||||
if window.count < windowSamples {
|
||||
window.append(contentsOf: [Float](repeating: 0, count: windowSamples - window.count))
|
||||
}
|
||||
|
||||
let input = MLXArray(window).reshaped(1, 1, windowSamples)
|
||||
let posteriors = segmentationModel(input)
|
||||
let speakerProbs = PowersetDecoder.speakerProbabilities(from: posteriors)
|
||||
eval(speakerProbs)
|
||||
|
||||
var tracks = [[Float]]()
|
||||
for spk in 0..<3 {
|
||||
tracks.append(speakerProbs[0, 0..., spk].asArray(Float.self))
|
||||
}
|
||||
windowProbs.append(WindowProbs(startSample: start, endSample: end, tracks: tracks))
|
||||
}
|
||||
|
||||
guard !windowProbs.isEmpty else {
|
||||
return DiarizationResult(segments: [], numSpeakers: 0, speakerEmbeddings: [])
|
||||
}
|
||||
|
||||
// Step 2: Extract per-window per-speaker embeddings from non-overlapping speech
|
||||
let minEmbeddingSamples = sampleRate / 2 // 0.5s minimum for embedding
|
||||
var windowEmbeddings = [WindowSpeakerEmbedding]()
|
||||
|
||||
for (wIdx, wp) in windowProbs.enumerated() {
|
||||
completedUnits += 1
|
||||
if progressHandler?(Float(completedUnits) / Float(totalUnits), "Embedding \(wIdx + 1)/\(windowProbs.count)") == false {
|
||||
return emptyResult
|
||||
}
|
||||
|
||||
let windowStartSample = wp.startSample
|
||||
|
||||
for localSpk in 0..<3 {
|
||||
let probs = wp.tracks[localSpk]
|
||||
|
||||
// Binarize this speaker's track
|
||||
let binarySegments = PowersetDecoder.binarize(
|
||||
probs: probs, onset: config.onset,
|
||||
offset: config.offset, frameDuration: frameDuration)
|
||||
|
||||
guard !binarySegments.isEmpty else { continue }
|
||||
|
||||
// Collect audio from frames where ONLY this speaker is active (non-overlapping)
|
||||
var spkAudio = [Float]()
|
||||
|
||||
for seg in binarySegments {
|
||||
let segStartFrame = Int(seg.startTime / frameDuration)
|
||||
let segEndFrame = min(Int(seg.endTime / frameDuration), probs.count)
|
||||
|
||||
for frame in segStartFrame..<segEndFrame {
|
||||
// Check if other speakers are below offset at this frame
|
||||
var otherActive = false
|
||||
for otherSpk in 0..<3 where otherSpk != localSpk {
|
||||
if wp.tracks[otherSpk][frame] >= 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..<frameEndSample])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Need minimum 0.5s of audio for a reliable embedding
|
||||
guard spkAudio.count >= 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
|
||||
@@ -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..<chunkEnd])
|
||||
let chunkProbs = runCoreML(features: chunkFeatures, numFrames: chunkFrames)
|
||||
allProbs.append(contentsOf: chunkProbs)
|
||||
offset += chunkFrames
|
||||
}
|
||||
probs = allProbs
|
||||
}
|
||||
#else
|
||||
return []
|
||||
#endif
|
||||
|
||||
let smoothed = smoothProbabilities(probs, windowSize: smoothWindowSize)
|
||||
return extractSegments(smoothed, threshold: speechThreshold, frameShift: 0.01,
|
||||
minSpeechDuration: minSpeechDuration, minSilenceDuration: minSilenceDuration)
|
||||
}
|
||||
|
||||
/// Detect speech segments in audio.
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - audio: Float32 PCM samples
|
||||
/// - sampleRate: Sample rate of the audio
|
||||
/// - Returns: Array of speech segments with start/end times
|
||||
public func detectSpeech(
|
||||
audio: [Float],
|
||||
sampleRate: Int = 16000
|
||||
) -> [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..<chunkEnd])
|
||||
let chunkProbs = runCoreML(features: chunkFeatures, numFrames: chunkFrames)
|
||||
allProbs.append(contentsOf: chunkProbs)
|
||||
offset += chunkFrames
|
||||
}
|
||||
probs = allProbs
|
||||
}
|
||||
#else
|
||||
return []
|
||||
#endif
|
||||
|
||||
// Post-process: smooth → threshold → segment
|
||||
let smoothed = smoothProbabilities(probs, windowSize: smoothWindowSize)
|
||||
return extractSegments(
|
||||
smoothed,
|
||||
threshold: speechThreshold,
|
||||
frameShift: 0.01,
|
||||
minSpeechDuration: minSpeechDuration,
|
||||
minSilenceDuration: minSilenceDuration
|
||||
)
|
||||
}
|
||||
|
||||
#if canImport(CoreML)
|
||||
private func runCoreML(features: [Float], numFrames: Int) -> [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..<numFrames {
|
||||
probs[i] = outputArray[[0, NSNumber(value: i), 0]].floatValue
|
||||
}
|
||||
return probs
|
||||
}
|
||||
#endif
|
||||
|
||||
// MARK: - Post-processing
|
||||
|
||||
private func smoothProbabilities(_ probs: [Float], windowSize: Int) -> [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..<probs.count {
|
||||
let lo = max(0, i - half)
|
||||
let hi = min(probs.count, i + half + 1)
|
||||
var sum: Float = 0
|
||||
for j in lo..<hi { sum += probs[j] }
|
||||
smoothed[i] = sum / Float(hi - lo)
|
||||
}
|
||||
return smoothed
|
||||
}
|
||||
|
||||
private func extractSegments(
|
||||
_ probs: [Float],
|
||||
threshold: Float,
|
||||
frameShift: Float,
|
||||
minSpeechDuration: Float,
|
||||
minSilenceDuration: Float
|
||||
) -> [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..<segments.count {
|
||||
let gap = segments[i].startTime - merged.last!.endTime
|
||||
if gap < minSilenceDuration {
|
||||
let last = merged.removeLast()
|
||||
merged.append(SpeechSegment(
|
||||
startTime: last.startTime,
|
||||
endTime: segments[i].endTime))
|
||||
} else {
|
||||
merged.append(segments[i])
|
||||
}
|
||||
}
|
||||
return merged
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - CoreML Input Wrapper
|
||||
|
||||
#if canImport(CoreML)
|
||||
private class FireRedVADInput: MLFeatureProvider {
|
||||
let features: MLMultiArray
|
||||
|
||||
init(features: MLMultiArray) {
|
||||
self.features = features
|
||||
}
|
||||
|
||||
var featureNames: Set<String> { ["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..<numFrames {
|
||||
let start = frame * frameShift
|
||||
|
||||
// Scale to int16 range
|
||||
var rawFrame = [Float](repeating: 0, count: frameLength)
|
||||
for i in 0..<frameLength {
|
||||
rawFrame[i] = samples[start + i] * 32768.0
|
||||
}
|
||||
|
||||
// DC offset removal
|
||||
var dcMean: Float = 0
|
||||
vDSP_meanv(rawFrame, 1, &dcMean, vDSP_Length(frameLength))
|
||||
var negMean = -dcMean
|
||||
vDSP_vsadd(rawFrame, 1, &negMean, &rawFrame, 1, vDSP_Length(frameLength))
|
||||
|
||||
// Pre-emphasis: y[0] = x[0]*(1-coeff), y[n] = x[n] - coeff*x[n-1]
|
||||
for i in stride(from: frameLength - 1, through: 1, by: -1) {
|
||||
rawFrame[i] -= preemphCoeff * rawFrame[i - 1]
|
||||
}
|
||||
rawFrame[0] *= (1.0 - preemphCoeff)
|
||||
|
||||
// Window + zero-pad
|
||||
for i in 0..<frameLength { padded[i] = rawFrame[i] * window[i] }
|
||||
for i in frameLength..<nFFT { padded[i] = 0 }
|
||||
|
||||
// DFT via Accelerate matrix-vector multiply
|
||||
// real[k] = sum_n padded[n] * cos(2πkn/N) = dftCos[257×512] × padded[512]
|
||||
// imag[k] = sum_n padded[n] * sin(2πkn/N) = dftSin[257×512] × padded[512]
|
||||
vDSP_mmul(dftCos, 1, padded, 1, &realParts, 1,
|
||||
vDSP_Length(nBins), 1, vDSP_Length(nFFT))
|
||||
vDSP_mmul(dftSin, 1, padded, 1, &imagParts, 1,
|
||||
vDSP_Length(nBins), 1, vDSP_Length(nFFT))
|
||||
|
||||
// Power spectrum: |X(k)|² = real² + imag²
|
||||
vDSP_vsq(realParts, 1, &powerSpec, 1, vDSP_Length(nBins))
|
||||
var imagSq = [Float](repeating: 0, count: nBins)
|
||||
vDSP_vsq(imagParts, 1, &imagSq, 1, vDSP_Length(nBins))
|
||||
vDSP_vadd(powerSpec, 1, imagSq, 1, &powerSpec, 1, vDSP_Length(nBins))
|
||||
|
||||
// Mel filterbank: mel[m] = sum_b fb[m,b] * power[b]
|
||||
vDSP_mmul(melFilterbank, 1, powerSpec, 1, &melEnergies, 1,
|
||||
vDSP_Length(nMels), 1, vDSP_Length(nBins))
|
||||
|
||||
// Log with Kaldi energy floor
|
||||
for m in 0..<nMels {
|
||||
melEnergies[m] = log(max(melEnergies[m], Float.ulpOfOne))
|
||||
}
|
||||
|
||||
allFeatures.append(contentsOf: melEnergies)
|
||||
}
|
||||
|
||||
return allFeatures
|
||||
}
|
||||
|
||||
/// Build mel filterbank matrix [nMels, nBins] matching Kaldi's computation.
|
||||
///
|
||||
/// Kaldi uses Hz-domain center frequencies for triangular filter weights,
|
||||
/// not integer bin indices. Each FFT bin maps to a Hz frequency via
|
||||
/// `bin * sampleRate / nFFT`, and the weight is computed from the distance
|
||||
/// to the neighboring mel center frequencies in Hz.
|
||||
private static func buildMelFilterbank(
|
||||
nMels: Int, nBins: Int, sampleRate: Int, nFFT: Int
|
||||
) -> [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..<nMels {
|
||||
let leftHz = centerFreqs[m]
|
||||
let centerHz = centerFreqs[m + 1]
|
||||
let rightHz = centerFreqs[m + 2]
|
||||
|
||||
for b in 0..<nBins {
|
||||
let freqHz = Float(b) * binToHz
|
||||
|
||||
if freqHz >= 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
|
||||
}
|
||||
}
|
||||
+238
@@ -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..<nBins {
|
||||
fftFreqs[i] = Float(i) * Float(sampleRate) / Float(paddedFFT)
|
||||
}
|
||||
|
||||
let melMin = hzToMel(fMin)
|
||||
let melMax = hzToMel(fMax)
|
||||
|
||||
let nMelPoints = nMels + 2
|
||||
var melPoints = [Float](repeating: 0, count: nMelPoints)
|
||||
for i in 0..<nMelPoints {
|
||||
melPoints[i] = melMin + Float(i) * (melMax - melMin) / Float(nMelPoints - 1)
|
||||
}
|
||||
|
||||
let filterFreqs = melPoints.map { melToHz($0) }
|
||||
|
||||
var filterDiff = [Float](repeating: 0, count: nMelPoints - 1)
|
||||
for i in 0..<(nMelPoints - 1) {
|
||||
filterDiff[i] = filterFreqs[i + 1] - filterFreqs[i]
|
||||
}
|
||||
|
||||
// Build filterbank [nBins, nMels]
|
||||
var filterbank = [Float](repeating: 0, count: nBins * nMels)
|
||||
for bin in 0..<nBins {
|
||||
let freq = fftFreqs[bin]
|
||||
for mel in 0..<nMels {
|
||||
let lowFreq = filterFreqs[mel]
|
||||
let highFreq = filterFreqs[mel + 2]
|
||||
let downSlope = (freq - lowFreq) / filterDiff[mel]
|
||||
let upSlope = (highFreq - freq) / filterDiff[mel + 1]
|
||||
filterbank[bin * nMels + mel] = max(0.0, min(downSlope, upSlope))
|
||||
}
|
||||
}
|
||||
|
||||
// Slaney normalization
|
||||
for mel in 0..<nMels {
|
||||
let enorm = 2.0 / (filterFreqs[mel + 2] - filterFreqs[mel])
|
||||
for bin in 0..<nBins {
|
||||
filterbank[bin * nMels + mel] *= enorm
|
||||
}
|
||||
}
|
||||
|
||||
// Transpose to [nMels, nBins]
|
||||
var transposed = [Float](repeating: 0, count: nMels * nBins)
|
||||
for mel in 0..<nMels {
|
||||
for bin in 0..<nBins {
|
||||
transposed[mel * nBins + bin] = filterbank[bin * nMels + mel]
|
||||
}
|
||||
}
|
||||
|
||||
self.melFilterbank = transposed
|
||||
}
|
||||
|
||||
/// Extract raw 80-dim log-mel features as a flat Float array.
|
||||
///
|
||||
/// This avoids creating an MLXArray, useful for CoreML paths that need
|
||||
/// raw floats without an MLX round-trip.
|
||||
///
|
||||
/// - Parameter audio: PCM Float32 samples at 16kHz
|
||||
/// - Returns: `(melSpec, nFrames)` where melSpec is a flat `[nFrames * 80]` array
|
||||
func extractRaw(_ audio: [Float]) -> (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..<audio.count {
|
||||
emphasized[i] = audio[i] - preEmphasis * audio[i - 1]
|
||||
}
|
||||
}
|
||||
|
||||
// Reflect padding
|
||||
let padLength = nFFT / 2
|
||||
var paddedAudio = [Float](repeating: 0, count: padLength + emphasized.count + padLength)
|
||||
|
||||
for i in 0..<padLength {
|
||||
let srcIdx = min(padLength - i, emphasized.count - 1)
|
||||
paddedAudio[i] = emphasized[max(0, srcIdx)]
|
||||
}
|
||||
for i in 0..<emphasized.count {
|
||||
paddedAudio[padLength + i] = emphasized[i]
|
||||
}
|
||||
for i in 0..<padLength {
|
||||
let srcIdx = emphasized.count - 2 - i
|
||||
paddedAudio[padLength + emphasized.count + i] = emphasized[max(0, srcIdx)]
|
||||
}
|
||||
|
||||
let nFrames = (paddedAudio.count - nFFT) / hopLength + 1
|
||||
|
||||
var splitReal = [Float](repeating: 0, count: halfPadded)
|
||||
var splitImag = [Float](repeating: 0, count: halfPadded)
|
||||
var paddedFrame = [Float](repeating: 0, count: paddedFFT)
|
||||
var magnitude = [Float](repeating: 0, count: nFrames * nBins)
|
||||
|
||||
for frame in 0..<nFrames {
|
||||
let start = frame * hopLength
|
||||
|
||||
paddedAudio.withUnsafeBufferPointer { buf in
|
||||
vDSP_vmul(buf.baseAddress! + start, 1, window, 1, &paddedFrame, 1, vDSP_Length(nFFT))
|
||||
}
|
||||
for i in nFFT..<paddedFFT {
|
||||
paddedFrame[i] = 0
|
||||
}
|
||||
|
||||
for i in 0..<halfPadded {
|
||||
splitReal[i] = paddedFrame[2 * i]
|
||||
splitImag[i] = paddedFrame[2 * i + 1]
|
||||
}
|
||||
|
||||
splitReal.withUnsafeMutableBufferPointer { realBuf in
|
||||
splitImag.withUnsafeMutableBufferPointer { imagBuf in
|
||||
var splitComplex = DSPSplitComplex(
|
||||
realp: realBuf.baseAddress!,
|
||||
imagp: imagBuf.baseAddress!)
|
||||
vDSP_fft_zrip(fftSetup, &splitComplex, 1, log2PaddedFFT, FFTDirection(kFFTDirection_Forward))
|
||||
}
|
||||
}
|
||||
|
||||
let baseIdx = frame * nBins
|
||||
magnitude[baseIdx] = splitReal[0] * splitReal[0]
|
||||
magnitude[baseIdx + halfPadded] = splitImag[0] * splitImag[0]
|
||||
for k in 1..<halfPadded {
|
||||
magnitude[baseIdx + k] = splitReal[k] * splitReal[k] + splitImag[k] * splitImag[k]
|
||||
}
|
||||
}
|
||||
|
||||
// Mel filterbank matmul
|
||||
var melSpec = [Float](repeating: 0, count: nFrames * nMels)
|
||||
var filterbankT = [Float](repeating: 0, count: nBins * nMels)
|
||||
vDSP_mtrans(melFilterbank, 1, &filterbankT, 1, vDSP_Length(nBins), vDSP_Length(nMels))
|
||||
|
||||
vDSP_mmul(magnitude, 1, filterbankT, 1, &melSpec, 1,
|
||||
vDSP_Length(nFrames), vDSP_Length(nMels), vDSP_Length(nBins))
|
||||
|
||||
// Simple log-mel: log(max(x, 1e-10))
|
||||
let count = melSpec.count
|
||||
var countN = Int32(count)
|
||||
|
||||
var epsilon: Float = 1e-10
|
||||
vDSP_vclip(melSpec, 1, &epsilon, [Float.greatestFiniteMagnitude], &melSpec, 1, vDSP_Length(count))
|
||||
vvlogf(&melSpec, melSpec, &countN)
|
||||
|
||||
// CMN (Cepstral Mean Normalization): subtract per-bin temporal mean.
|
||||
// Matches WeSpeaker Python: feat = feat - torch.mean(feat, dim=0)
|
||||
// Removes channel-specific spectral shape, critical for speaker discrimination.
|
||||
if nFrames > 0 {
|
||||
// Compute per-bin mean
|
||||
var binMeans = [Float](repeating: 0, count: nMels)
|
||||
for frame in 0..<nFrames {
|
||||
let base = frame * nMels
|
||||
for bin in 0..<nMels {
|
||||
binMeans[bin] += melSpec[base + bin]
|
||||
}
|
||||
}
|
||||
let invFrames = 1.0 / Float(nFrames)
|
||||
vDSP_vsmul(binMeans, 1, [invFrames], &binMeans, 1, vDSP_Length(nMels))
|
||||
|
||||
// Subtract mean from each frame
|
||||
for frame in 0..<nFrames {
|
||||
let base = frame * nMels
|
||||
for bin in 0..<nMels {
|
||||
melSpec[base + bin] -= binMeans[bin]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return (melSpec, nFrames)
|
||||
}
|
||||
|
||||
/// Extract 80-dim log-mel features from audio.
|
||||
/// - Parameter audio: PCM Float32 samples at 16kHz
|
||||
/// - Returns: `MLXArray [T, 80]` — time-major mel features
|
||||
func extract(_ audio: [Float]) -> MLXArray {
|
||||
let (melSpec, nFrames) = extractRaw(audio)
|
||||
return MLXArray(melSpec, [nFrames, nMels])
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
+16
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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 }
|
||||
}
|
||||
+36
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
+432
@@ -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..<framesToCopy {
|
||||
let srcBase = (chunkStart + fi) * nMels
|
||||
let dstBase = fi * nMels
|
||||
for di in 0..<nMels {
|
||||
chunkMel[dstBase + di] = melSpec[srcBase + di]
|
||||
}
|
||||
}
|
||||
|
||||
do {
|
||||
let output = try model.predict(
|
||||
chunk: chunkMel,
|
||||
chunkLength: actualLen,
|
||||
spkcache: spkcache,
|
||||
spkcacheLength: spkcacheLength,
|
||||
fifo: fifo,
|
||||
fifoLength: fifoLength
|
||||
)
|
||||
|
||||
// Extract core predictions (skip spkcache + fifo + left context,
|
||||
// trim right context)
|
||||
let validEmbs: Int = output.validEmbFrames
|
||||
let lcFrames: Int = Int(Float(leftOffset) / Float(subFactor) + 0.5)
|
||||
let rcFrames: Int = Int(ceil(Float(rightOffset) / Float(subFactor)))
|
||||
let coreLen: Int = validEmbs - lcFrames - rcFrames
|
||||
let corePredLen = coreLen > 0 ? coreLen : 0
|
||||
|
||||
let predOffset = spkcacheLength + fifoLength + lcFrames
|
||||
let totalPredFrames = output.predsFrames
|
||||
|
||||
var chunkProbs = [Float]()
|
||||
for f in 0..<corePredLen {
|
||||
let predFrame = predOffset + f
|
||||
guard predFrame < totalPredFrames else { break }
|
||||
for s in 0..<numSpeakers {
|
||||
chunkProbs.append(output.pred(frame: predFrame, speaker: s))
|
||||
}
|
||||
}
|
||||
allChunkProbs.append(chunkProbs)
|
||||
|
||||
// Update streaming state (FIFO overflow → spkcache)
|
||||
updateState(from: output)
|
||||
} catch {
|
||||
print("Warning: Sortformer inference failed on chunk at mel frame \(sttFeat): \(error)")
|
||||
}
|
||||
|
||||
sttFeat = endFeat
|
||||
}
|
||||
|
||||
guard !allChunkProbs.isEmpty else {
|
||||
return DiarizationResult(segments: [], numSpeakers: 0, speakerEmbeddings: [])
|
||||
}
|
||||
|
||||
// Concatenate all core predictions
|
||||
let audioDuration = Float(samples.count) / Float(self.config.sampleRate)
|
||||
let segments = binarizeCorePredictions(
|
||||
allChunkProbs: allChunkProbs,
|
||||
audioDuration: audioDuration,
|
||||
numSpeakers: numSpeakers,
|
||||
onset: config.onset,
|
||||
offset: config.offset,
|
||||
minSpeechDuration: config.minSpeechDuration,
|
||||
minSilenceDuration: config.minSilenceDuration
|
||||
)
|
||||
|
||||
let usedSpeakers = Set(segments.map(\.speakerId))
|
||||
return DiarizationResult(
|
||||
segments: segments,
|
||||
numSpeakers: usedSpeakers.count,
|
||||
speakerEmbeddings: [] // End-to-end model, no separate embeddings
|
||||
)
|
||||
}
|
||||
|
||||
// MARK: - State Management (NeMo FIFO→spkcache pattern)
|
||||
|
||||
/// Update spkcache and fifo buffers from encoder embeddings.
|
||||
///
|
||||
/// Follows NeMo's streaming_update: new embeddings go into FIFO.
|
||||
/// When FIFO overflows, oldest frames move to spkcache.
|
||||
private func updateState(from output: SortformerOutput) {
|
||||
let validFrames = output.validEmbFrames
|
||||
guard validFrames > 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..<validFrames {
|
||||
let srcBase = f * dim
|
||||
let dstBase = (fifoLength + f) * dim
|
||||
for d in 0..<dim {
|
||||
fifo[dstBase + d] = output.encoderEmbs[srcBase + d]
|
||||
}
|
||||
}
|
||||
fifoLength += validFrames
|
||||
} else {
|
||||
// FIFO overflow: move oldest frames to spkcache
|
||||
let overflow = fifoLength + validFrames - fifoCapacity
|
||||
|
||||
// Move overflow frames from front of FIFO to spkcache
|
||||
if spkcacheLength + overflow <= cacheCapacity {
|
||||
// Append to spkcache
|
||||
for f in 0..<overflow {
|
||||
let srcBase = f * dim
|
||||
let dstBase = (spkcacheLength + f) * dim
|
||||
for d in 0..<dim {
|
||||
spkcache[dstBase + d] = fifo[srcBase + d]
|
||||
}
|
||||
}
|
||||
spkcacheLength += overflow
|
||||
} else {
|
||||
// Spkcache also overflows — shift left and append
|
||||
let cacheOverflow = spkcacheLength + overflow - cacheCapacity
|
||||
let keep = spkcacheLength - cacheOverflow
|
||||
if keep > 0 {
|
||||
for f in 0..<keep {
|
||||
let srcBase = (f + cacheOverflow) * dim
|
||||
let dstBase = f * dim
|
||||
for d in 0..<dim {
|
||||
spkcache[dstBase + d] = spkcache[srcBase + d]
|
||||
}
|
||||
}
|
||||
}
|
||||
for f in 0..<overflow {
|
||||
let srcBase = f * dim
|
||||
let dstBase = (keep + f) * dim
|
||||
for d in 0..<dim {
|
||||
spkcache[dstBase + d] = fifo[srcBase + d]
|
||||
}
|
||||
}
|
||||
spkcacheLength = min(cacheCapacity, keep + overflow)
|
||||
}
|
||||
|
||||
// Shift FIFO left by overflow, then append new frames
|
||||
let remaining = fifoLength - overflow
|
||||
if remaining > 0 {
|
||||
for f in 0..<remaining {
|
||||
let srcBase = (f + overflow) * dim
|
||||
let dstBase = f * dim
|
||||
for d in 0..<dim {
|
||||
fifo[dstBase + d] = fifo[srcBase + d]
|
||||
}
|
||||
}
|
||||
}
|
||||
fifoLength = remaining
|
||||
for f in 0..<validFrames {
|
||||
let srcBase = f * dim
|
||||
let dstBase = (fifoLength + f) * dim
|
||||
for d in 0..<dim {
|
||||
fifo[dstBase + d] = output.encoderEmbs[srcBase + d]
|
||||
}
|
||||
}
|
||||
fifoLength += validFrames
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Binarization
|
||||
|
||||
/// Concatenate per-chunk core predictions and binarize into segments.
|
||||
private func binarizeCorePredictions(
|
||||
allChunkProbs: [[Float]],
|
||||
audioDuration: Float,
|
||||
numSpeakers: Int,
|
||||
onset: Float,
|
||||
offset: Float,
|
||||
minSpeechDuration: Float,
|
||||
minSilenceDuration: Float
|
||||
) -> [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..<allProbs.count {
|
||||
if allProbs[i] > 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..<numSpeakers {
|
||||
var probs = [Float](repeating: 0, count: totalFrames)
|
||||
for f in 0..<totalFrames {
|
||||
probs[f] = allProbs[f * numSpeakers + spk]
|
||||
}
|
||||
|
||||
let rawSegments = PowersetDecoder.binarize(
|
||||
probs: probs,
|
||||
onset: onset,
|
||||
offset: offset,
|
||||
frameDuration: frameDuration
|
||||
)
|
||||
|
||||
for seg in rawSegments {
|
||||
let duration = seg.endTime - seg.startTime
|
||||
guard duration >= 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
|
||||
+205
@@ -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..<config.nFFT {
|
||||
window[i] = 0.5 - 0.5 * cos(2.0 * Float.pi * Float(i) / Float(config.nFFT - 1))
|
||||
}
|
||||
|
||||
guard let setup = vDSP_create_fftsetup(log2PaddedFFT, 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 = 0.0
|
||||
let fMax: Float = Float(sampleRate) / 2.0
|
||||
|
||||
// HTK mel scale
|
||||
func hzToMel(_ hz: Float) -> 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..<nBins {
|
||||
fftFreqs[i] = Float(i) * Float(sampleRate) / Float(paddedFFT)
|
||||
}
|
||||
|
||||
let melMin = hzToMel(fMin)
|
||||
let melMax = hzToMel(fMax)
|
||||
|
||||
let nMelPoints = nMels + 2
|
||||
var melPoints = [Float](repeating: 0, count: nMelPoints)
|
||||
for i in 0..<nMelPoints {
|
||||
melPoints[i] = melMin + Float(i) * (melMax - melMin) / Float(nMelPoints - 1)
|
||||
}
|
||||
|
||||
let filterFreqs = melPoints.map { melToHz($0) }
|
||||
|
||||
var filterDiff = [Float](repeating: 0, count: nMelPoints - 1)
|
||||
for i in 0..<(nMelPoints - 1) {
|
||||
filterDiff[i] = filterFreqs[i + 1] - filterFreqs[i]
|
||||
}
|
||||
|
||||
// Build filterbank [nBins, nMels]
|
||||
var filterbank = [Float](repeating: 0, count: nBins * nMels)
|
||||
for bin in 0..<nBins {
|
||||
let freq = fftFreqs[bin]
|
||||
for mel in 0..<nMels {
|
||||
let lowFreq = filterFreqs[mel]
|
||||
let highFreq = filterFreqs[mel + 2]
|
||||
let downSlope = (freq - lowFreq) / filterDiff[mel]
|
||||
let upSlope = (highFreq - freq) / filterDiff[mel + 1]
|
||||
filterbank[bin * nMels + mel] = max(0.0, min(downSlope, upSlope))
|
||||
}
|
||||
}
|
||||
|
||||
// Slaney normalization
|
||||
for mel in 0..<nMels {
|
||||
let enorm = 2.0 / (filterFreqs[mel + 2] - filterFreqs[mel])
|
||||
for bin in 0..<nBins {
|
||||
filterbank[bin * nMels + mel] *= enorm
|
||||
}
|
||||
}
|
||||
|
||||
// Transpose to [nMels, nBins]
|
||||
var transposed = [Float](repeating: 0, count: nMels * nBins)
|
||||
for mel in 0..<nMels {
|
||||
for bin in 0..<nBins {
|
||||
transposed[mel * nBins + bin] = filterbank[bin * nMels + mel]
|
||||
}
|
||||
}
|
||||
|
||||
self.melFilterbank = transposed
|
||||
}
|
||||
|
||||
/// Extract 128-dim log-mel features from audio.
|
||||
///
|
||||
/// - Parameter audio: PCM Float32 samples at 16kHz
|
||||
/// - Returns: `(melSpec, nFrames)` where melSpec is a flat `[nFrames * 128]` array
|
||||
func extract(_ audio: [Float]) -> (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..<padLength {
|
||||
let srcIdx = min(padLength - i, audio.count - 1)
|
||||
paddedAudio[i] = audio[max(0, srcIdx)]
|
||||
}
|
||||
for i in 0..<audio.count {
|
||||
paddedAudio[padLength + i] = audio[i]
|
||||
}
|
||||
for i in 0..<padLength {
|
||||
let srcIdx = audio.count - 2 - i
|
||||
paddedAudio[padLength + audio.count + i] = audio[max(0, srcIdx)]
|
||||
}
|
||||
|
||||
let nFrames = (paddedAudio.count - nFFT) / hopLength + 1
|
||||
|
||||
var splitReal = [Float](repeating: 0, count: halfPadded)
|
||||
var splitImag = [Float](repeating: 0, count: halfPadded)
|
||||
var paddedFrame = [Float](repeating: 0, count: paddedFFT)
|
||||
var powerSpec = [Float](repeating: 0, count: nFrames * nBins)
|
||||
|
||||
for frame in 0..<nFrames {
|
||||
let start = frame * hopLength
|
||||
|
||||
paddedAudio.withUnsafeBufferPointer { buf in
|
||||
vDSP_vmul(buf.baseAddress! + start, 1, window, 1, &paddedFrame, 1, vDSP_Length(nFFT))
|
||||
}
|
||||
for i in nFFT..<paddedFFT {
|
||||
paddedFrame[i] = 0
|
||||
}
|
||||
|
||||
for i in 0..<halfPadded {
|
||||
splitReal[i] = paddedFrame[2 * i]
|
||||
splitImag[i] = paddedFrame[2 * i + 1]
|
||||
}
|
||||
|
||||
splitReal.withUnsafeMutableBufferPointer { realBuf in
|
||||
splitImag.withUnsafeMutableBufferPointer { imagBuf in
|
||||
var splitComplex = DSPSplitComplex(
|
||||
realp: realBuf.baseAddress!,
|
||||
imagp: imagBuf.baseAddress!)
|
||||
vDSP_fft_zrip(fftSetup, &splitComplex, 1, log2PaddedFFT, FFTDirection(kFFTDirection_Forward))
|
||||
}
|
||||
}
|
||||
|
||||
let baseIdx = frame * nBins
|
||||
// Power spectrum: |X|^2
|
||||
powerSpec[baseIdx] = splitReal[0] * splitReal[0]
|
||||
powerSpec[baseIdx + halfPadded] = splitImag[0] * splitImag[0]
|
||||
for k in 1..<halfPadded {
|
||||
powerSpec[baseIdx + k] = splitReal[k] * splitReal[k] + splitImag[k] * splitImag[k]
|
||||
}
|
||||
}
|
||||
|
||||
// Mel filterbank matmul: [nFrames, nBins] × [nBins, nMels] = [nFrames, nMels]
|
||||
var melSpec = [Float](repeating: 0, count: nFrames * nMels)
|
||||
var filterbankT = [Float](repeating: 0, count: nBins * nMels)
|
||||
vDSP_mtrans(melFilterbank, 1, &filterbankT, 1, vDSP_Length(nBins), vDSP_Length(nMels))
|
||||
|
||||
vDSP_mmul(powerSpec, 1, filterbankT, 1, &melSpec, 1,
|
||||
vDSP_Length(nFrames), vDSP_Length(nMels), vDSP_Length(nBins))
|
||||
|
||||
// Log-mel: log(max(x, 1e-10))
|
||||
let count = melSpec.count
|
||||
var countN = Int32(count)
|
||||
|
||||
var epsilon: Float = 1e-10
|
||||
vDSP_vclip(melSpec, 1, &epsilon, [Float.greatestFiniteMagnitude], &melSpec, 1, vDSP_Length(count))
|
||||
vvlogf(&melSpec, melSpec, &countN)
|
||||
|
||||
return (melSpec, nFrames)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
#if canImport(CoreML)
|
||||
import CoreML
|
||||
import Foundation
|
||||
import AudioCommon
|
||||
|
||||
/// CoreML wrapper for the Sortformer streaming diarization model.
|
||||
///
|
||||
/// Runs on Neural Engine via CoreML. The model takes a chunk of mel features
|
||||
/// plus streaming state buffers (spkcache, fifo) and outputs per-frame
|
||||
/// speaker predictions for up to 4 speakers.
|
||||
final class SortformerCoreMLModel {
|
||||
|
||||
private let model: MLModel
|
||||
let config: SortformerConfig
|
||||
|
||||
/// Input shape constants derived from the CoreML model.
|
||||
/// chunk: [1, 112, 128], spkcache: [1, 188, 512], fifo: [1, 40, 512]
|
||||
private let chunkFrames: Int = 112
|
||||
private let spkcacheFrames: Int
|
||||
private let fifoFrames: Int
|
||||
private let featureDim: Int
|
||||
|
||||
init(model: MLModel, config: SortformerConfig = .default) {
|
||||
self.model = model
|
||||
self.config = config
|
||||
self.spkcacheFrames = config.spkcacheLen
|
||||
self.fifoFrames = config.fifoLen
|
||||
self.featureDim = config.fcDModel
|
||||
}
|
||||
|
||||
/// Run one streaming inference step.
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - chunk: Mel features for this chunk, flat `[chunkFrames * nMels]` float array.
|
||||
/// If fewer frames available, zero-pad on the right.
|
||||
/// - chunkLength: Actual number of valid mel frames in the chunk
|
||||
/// - spkcache: Speaker cache state, flat `[spkcacheFrames * fcDModel]`
|
||||
/// - spkcacheLength: Number of valid frames in speaker cache
|
||||
/// - fifo: FIFO buffer state, flat `[fifoFrames * fcDModel]`
|
||||
/// - fifoLength: Number of valid frames in FIFO
|
||||
/// - Returns: `SortformerOutput` with predictions and updated state
|
||||
func predict(
|
||||
chunk: [Float],
|
||||
chunkLength: Int,
|
||||
spkcache: [Float],
|
||||
spkcacheLength: Int,
|
||||
fifo: [Float],
|
||||
fifoLength: Int
|
||||
) throws -> 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..<predsArray.shape.count).map { predsArray.shape[$0].intValue }
|
||||
let embsShape = (0..<embsArray.shape.count).map { embsArray.shape[$0].intValue }
|
||||
|
||||
let totalPreds = predsShape.reduce(1, *)
|
||||
let totalEmbs = embsShape.reduce(1, *)
|
||||
|
||||
var preds = [Float](repeating: 0, count: totalPreds)
|
||||
let predsPtr = predsArray.dataPointer.assumingMemoryBound(to: Float.self)
|
||||
for i in 0..<totalPreds { preds[i] = predsPtr[i] }
|
||||
|
||||
var embs = [Float](repeating: 0, count: totalEmbs)
|
||||
let embsPtr = embsArray.dataPointer.assumingMemoryBound(to: Float.self)
|
||||
for i in 0..<totalEmbs { embs[i] = embsPtr[i] }
|
||||
|
||||
let embsLenPtr = embsLenArray.dataPointer.assumingMemoryBound(to: Int32.self)
|
||||
let validEmbFrames = Int(embsLenPtr[0])
|
||||
|
||||
return SortformerOutput(
|
||||
speakerPreds: preds,
|
||||
predsFrames: predsShape.count >= 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..<count {
|
||||
ptr[i] = data[i]
|
||||
}
|
||||
// Zero-fill remainder
|
||||
for i in count..<array.count {
|
||||
ptr[i] = 0
|
||||
}
|
||||
return array
|
||||
}
|
||||
|
||||
private func makeScalarInt32Array(value: Int32) throws -> 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
|
||||
+29
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
}
|
||||
+210
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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..<a.count {
|
||||
dot += a[i] * b[i]
|
||||
normA += a[i] * a[i]
|
||||
normB += b[i] * b[i]
|
||||
}
|
||||
let denom = sqrt(normA) * sqrt(normB)
|
||||
return denom > 0 ? dot / denom : 0
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
+33
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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)")
|
||||
}
|
||||
|
||||
@@ -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 }
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<Bool>) {
|
||||
_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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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) {}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user