Rocky
537a68552a
Reduce perceived latency from key release to final text: - Adaptive chunking: 2.5s first chunk + 5s follow-ups so short utterances start on-device recognition while still recording. - Session-level ASR warmup and audio-format cache reuse to remove per-utterance cold-start of SpeechAnalyzer. - Mirror live pipelined partials to the keyboard transcript line via a new flow.transcriptionPartial App Group key + Darwin ping. Also commits the accumulated custom language model, Flow session, keyboard extension restructure, and Xiaomi MiMo provider work in progress on this branch.
375 lines
13 KiB
Swift
375 lines
13 KiB
Swift
// CustomLanguageModelManager.swift
|
|
// OSGKeyboard · Shared
|
|
//
|
|
// Prepares the bundled SFCustomLanguageModelData asset on device and shares
|
|
// the compiled LM + Vocab through the App Group container. Both the host app
|
|
// and keyboard extension read the same prepared configuration for
|
|
// DictationTranscriber content hints.
|
|
|
|
import Foundation
|
|
import Speech
|
|
import os
|
|
|
|
public final class CustomLanguageModelManager: @unchecked Sendable {
|
|
|
|
public static let shared = CustomLanguageModelManager()
|
|
|
|
public enum PrepareState: Equatable, Sendable {
|
|
case idle
|
|
case preparing
|
|
case ready
|
|
case failed(String)
|
|
}
|
|
|
|
struct BundledManifest: Decodable {
|
|
let version: String
|
|
let bin_bytes: Int
|
|
let identifier: String
|
|
}
|
|
|
|
private enum Storage {
|
|
static let subdirectory = "CustomLanguageModel/v1"
|
|
static let fingerprintKey = "customLM.preparedFingerprint"
|
|
static let preparedAtKey = "customLM.preparedAt"
|
|
static let lastFailureAtKey = "customLM.lastFailureAt"
|
|
static let attemptCountKey = "customLM.attemptCount"
|
|
static let maxRetryAttempts = 3
|
|
/// Backoff after failure attempts 1, 2, and 3 (seconds).
|
|
static let backoffIntervals: [TimeInterval] = [30, 120, 600]
|
|
}
|
|
|
|
private let lock = OSAllocatedUnfairLock()
|
|
private var cachedConfiguration: SFSpeechLanguageModel.Configuration?
|
|
private var state: PrepareState = .idle
|
|
private var prepareTask: Task<Void, Never>?
|
|
|
|
private init() {}
|
|
|
|
// MARK: - Public API
|
|
|
|
/// Returns a prepared configuration for Chinese locales when available.
|
|
public func configurationForTranscription(locale: Locale) -> SFSpeechLanguageModel.Configuration? {
|
|
guard Self.isChineseLocale(locale) else { return nil }
|
|
return lock.withLock { () -> SFSpeechLanguageModel.Configuration? in
|
|
if let cachedConfiguration {
|
|
return cachedConfiguration
|
|
}
|
|
if let loaded = Self.loadCachedConfigurationFromDisk() {
|
|
cachedConfiguration = loaded
|
|
state = .ready
|
|
return loaded
|
|
}
|
|
return nil
|
|
}
|
|
}
|
|
|
|
public func currentState() -> PrepareState {
|
|
lock.withLock { state }
|
|
}
|
|
|
|
/// Fire-and-forget preparation for the host app. Safe to call repeatedly.
|
|
/// Retries after exponential backoff when a prior attempt failed.
|
|
public func prepareInBackgroundIfNeeded() {
|
|
guard AppGroup.isAvailable else { return }
|
|
|
|
let shouldStart = lock.withLock { () -> Bool in
|
|
if case .preparing = state { return false }
|
|
if cachedConfiguration != nil { return false }
|
|
if let loaded = Self.loadCachedConfigurationFromDisk() {
|
|
cachedConfiguration = loaded
|
|
state = .ready
|
|
Self.clearRetryState()
|
|
return false
|
|
}
|
|
if prepareTask != nil { return false }
|
|
|
|
if case .failed = state {
|
|
guard Self.canRetryAfterFailure() else { return false }
|
|
} else if !Self.canRetryAfterFailure() {
|
|
return false
|
|
}
|
|
|
|
state = .preparing
|
|
return true
|
|
}
|
|
guard shouldStart else { return }
|
|
|
|
prepareTask = Task.detached(priority: .utility) { [weak self] in
|
|
guard let self else { return }
|
|
defer {
|
|
self.lock.withLock { self.prepareTask = nil }
|
|
}
|
|
do {
|
|
_ = try await self.prepareIfNeeded()
|
|
} catch {
|
|
Self.recordFailure()
|
|
self.lock.withLock {
|
|
self.state = .failed(error.localizedDescription)
|
|
}
|
|
Self.log(
|
|
"prepare failed (attempt \(Self.storedAttemptCount())): \(error.localizedDescription)"
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Prepares the bundled training asset into the App Group container.
|
|
@discardableResult
|
|
public func prepareIfNeeded() async throws -> SFSpeechLanguageModel.Configuration? {
|
|
if let existing = configurationForTranscription(locale: Locale(identifier: "zh-Hans")) {
|
|
lock.withLock { state = .ready }
|
|
Self.clearRetryState()
|
|
return existing
|
|
}
|
|
|
|
guard Self.canRetryAfterFailure() else {
|
|
throw PrepareError.retryBudgetExhausted
|
|
}
|
|
|
|
guard let manifest = Self.bundledManifest() else {
|
|
throw PrepareError.missingManifest
|
|
}
|
|
guard let assetURL = Self.bundledTrainingAssetURL() else {
|
|
throw PrepareError.missingTrainingAsset
|
|
}
|
|
guard let preparedDir = Self.preparedDirectoryURL() else {
|
|
throw PrepareError.missingAppGroupContainer
|
|
}
|
|
|
|
let fingerprint = Self.fingerprint(for: manifest)
|
|
if Self.storedFingerprint() == fingerprint,
|
|
let cached = Self.loadCachedConfigurationFromDisk() {
|
|
lock.withLock {
|
|
cachedConfiguration = cached
|
|
state = .ready
|
|
}
|
|
Self.clearRetryState()
|
|
return cached
|
|
}
|
|
|
|
lock.withLock { state = .preparing }
|
|
|
|
let languageModelURL = preparedDir.appendingPathComponent("LM")
|
|
let vocabularyURL = preparedDir.appendingPathComponent("Vocab")
|
|
try Self.removeItemIfExists(at: languageModelURL)
|
|
try Self.removeItemIfExists(at: vocabularyURL)
|
|
|
|
let configuration = SFSpeechLanguageModel.Configuration(
|
|
languageModel: languageModelURL,
|
|
vocabulary: vocabularyURL
|
|
)
|
|
|
|
Self.log("preparing custom LM (\(manifest.bin_bytes) byte asset)…")
|
|
try await Self.prepareLanguageModel(assetURL: assetURL, configuration: configuration)
|
|
|
|
guard FileManager.default.fileExists(atPath: languageModelURL.path),
|
|
FileManager.default.fileExists(atPath: vocabularyURL.path) else {
|
|
throw PrepareError.missingPreparedArtifacts
|
|
}
|
|
|
|
AppGroup.defaultsIfAvailable?.set(fingerprint, forKey: Storage.fingerprintKey)
|
|
AppGroup.defaultsIfAvailable?.set(Date().timeIntervalSince1970, forKey: Storage.preparedAtKey)
|
|
Self.clearRetryState()
|
|
|
|
lock.withLock {
|
|
cachedConfiguration = configuration
|
|
state = .ready
|
|
}
|
|
|
|
Self.log("custom LM ready at \(preparedDir.path)")
|
|
return configuration
|
|
}
|
|
|
|
// MARK: - DictationTranscriber factory
|
|
|
|
public static func makeDictationTranscriber(
|
|
locale: Locale,
|
|
lmConfiguration: SFSpeechLanguageModel.Configuration?
|
|
) -> DictationTranscriber {
|
|
let preset = DictationTranscriber.Preset.progressiveLongDictation
|
|
guard let lmConfiguration, isChineseLocale(locale) else {
|
|
return DictationTranscriber(locale: locale, preset: preset)
|
|
}
|
|
|
|
let contentHints = preset.contentHints.union([
|
|
.customizedLanguage(modelConfiguration: lmConfiguration),
|
|
])
|
|
return DictationTranscriber(
|
|
locale: locale,
|
|
contentHints: contentHints,
|
|
transcriptionOptions: preset.transcriptionOptions,
|
|
reportingOptions: preset.reportingOptions,
|
|
attributeOptions: preset.attributeOptions
|
|
)
|
|
}
|
|
|
|
// MARK: - Bundle / disk helpers
|
|
|
|
private static var resourceBundle: Bundle {
|
|
Bundle(for: CustomLanguageModelManager.self)
|
|
}
|
|
|
|
static func bundledTrainingAssetURL() -> URL? {
|
|
if let url = resourceBundle.url(
|
|
forResource: "OSGKeyboardCLM",
|
|
withExtension: "bin",
|
|
subdirectory: Storage.subdirectory
|
|
) {
|
|
return url
|
|
}
|
|
return resourceBundle.url(forResource: "OSGKeyboardCLM", withExtension: "bin")
|
|
}
|
|
|
|
static func bundledManifest() -> BundledManifest? {
|
|
let manifestURL =
|
|
resourceBundle.url(
|
|
forResource: "compiled-manifest",
|
|
withExtension: "json",
|
|
subdirectory: Storage.subdirectory
|
|
)
|
|
?? resourceBundle.url(forResource: "compiled-manifest", withExtension: "json")
|
|
guard let manifestURL,
|
|
let data = try? Data(contentsOf: manifestURL),
|
|
let manifest = try? JSONDecoder().decode(BundledManifest.self, from: data)
|
|
else {
|
|
return nil
|
|
}
|
|
return manifest
|
|
}
|
|
|
|
static func preparedDirectoryURL() -> URL? {
|
|
guard let container = FileManager.default.containerURL(
|
|
forSecurityApplicationGroupIdentifier: AppGroup.identifier
|
|
) else {
|
|
return nil
|
|
}
|
|
let directory = container.appendingPathComponent(Storage.subdirectory, isDirectory: true)
|
|
try? FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
|
|
return directory
|
|
}
|
|
|
|
static func loadCachedConfigurationFromDisk() -> SFSpeechLanguageModel.Configuration? {
|
|
guard let manifest = bundledManifest(),
|
|
storedFingerprint() == fingerprint(for: manifest),
|
|
let preparedDir = preparedDirectoryURL()
|
|
else {
|
|
return nil
|
|
}
|
|
|
|
let languageModelURL = preparedDir.appendingPathComponent("LM")
|
|
let vocabularyURL = preparedDir.appendingPathComponent("Vocab")
|
|
let fm = FileManager.default
|
|
guard fm.fileExists(atPath: languageModelURL.path),
|
|
fm.fileExists(atPath: vocabularyURL.path) else {
|
|
return nil
|
|
}
|
|
|
|
return SFSpeechLanguageModel.Configuration(
|
|
languageModel: languageModelURL,
|
|
vocabulary: vocabularyURL
|
|
)
|
|
}
|
|
|
|
static func isChineseLocale(_ locale: Locale) -> Bool {
|
|
locale.identifier(.bcp47).lowercased().hasPrefix("zh")
|
|
}
|
|
|
|
private static func fingerprint(for manifest: BundledManifest) -> String {
|
|
"\(manifest.identifier)|\(manifest.version)|\(manifest.bin_bytes)"
|
|
}
|
|
|
|
private static func storedFingerprint() -> String? {
|
|
AppGroup.defaultsIfAvailable?.string(forKey: Storage.fingerprintKey)
|
|
}
|
|
|
|
private static func removeItemIfExists(at url: URL) throws {
|
|
let fm = FileManager.default
|
|
if fm.fileExists(atPath: url.path) {
|
|
try fm.removeItem(at: url)
|
|
}
|
|
}
|
|
|
|
private static func prepareLanguageModel(
|
|
assetURL: URL,
|
|
configuration: SFSpeechLanguageModel.Configuration
|
|
) async throws {
|
|
try await withCheckedThrowingContinuation {
|
|
(continuation: CheckedContinuation<Void, Error>) in
|
|
SFSpeechLanguageModel.prepareCustomLanguageModel(
|
|
for: assetURL,
|
|
configuration: configuration
|
|
) { error in
|
|
if let error {
|
|
continuation.resume(throwing: error)
|
|
} else {
|
|
continuation.resume()
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// MARK: - Retry / backoff
|
|
|
|
private static func storedAttemptCount() -> Int {
|
|
AppGroup.defaultsIfAvailable?.integer(forKey: Storage.attemptCountKey) ?? 0
|
|
}
|
|
|
|
private static func storedLastFailureAt() -> TimeInterval? {
|
|
let value = AppGroup.defaultsIfAvailable?.double(forKey: Storage.lastFailureAtKey) ?? 0
|
|
return value > 0 ? value : nil
|
|
}
|
|
|
|
private static func recordFailure() {
|
|
guard let defaults = AppGroup.defaultsIfAvailable else { return }
|
|
let nextAttempt = storedAttemptCount() + 1
|
|
defaults.set(nextAttempt, forKey: Storage.attemptCountKey)
|
|
defaults.set(Date().timeIntervalSince1970, forKey: Storage.lastFailureAtKey)
|
|
}
|
|
|
|
private static func clearRetryState() {
|
|
guard let defaults = AppGroup.defaultsIfAvailable else { return }
|
|
defaults.removeObject(forKey: Storage.attemptCountKey)
|
|
defaults.removeObject(forKey: Storage.lastFailureAtKey)
|
|
}
|
|
|
|
/// Returns false when retry budget is exhausted or backoff has not elapsed.
|
|
private static func canRetryAfterFailure() -> Bool {
|
|
let attempts = storedAttemptCount()
|
|
guard attempts > 0 else { return true }
|
|
guard attempts <= Storage.maxRetryAttempts else { return false }
|
|
|
|
guard let lastFailureAt = storedLastFailureAt() else { return true }
|
|
let backoffIndex = min(attempts - 1, Storage.backoffIntervals.count - 1)
|
|
let requiredDelay = Storage.backoffIntervals[backoffIndex]
|
|
let elapsed = Date().timeIntervalSince1970 - lastFailureAt
|
|
return elapsed >= requiredDelay
|
|
}
|
|
|
|
private static func log(_ message: String) {
|
|
OSGLog.clm.info("\(message, privacy: .public)")
|
|
}
|
|
|
|
enum PrepareError: LocalizedError {
|
|
case missingManifest
|
|
case missingTrainingAsset
|
|
case missingAppGroupContainer
|
|
case missingPreparedArtifacts
|
|
case retryBudgetExhausted
|
|
|
|
var errorDescription: String? {
|
|
switch self {
|
|
case .missingManifest:
|
|
return "Missing bundled custom language model manifest."
|
|
case .missingTrainingAsset:
|
|
return "Missing bundled custom language model training asset."
|
|
case .missingAppGroupContainer:
|
|
return "App Group container unavailable for custom language model preparation."
|
|
case .missingPreparedArtifacts:
|
|
return "Custom language model preparation did not produce LM/Vocab artifacts."
|
|
case .retryBudgetExhausted:
|
|
return "Custom language model preparation retry budget exhausted."
|
|
}
|
|
}
|
|
}
|
|
}
|