File size: 8,224 Bytes
64f7370 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 | import CoreML
import Foundation
import SpeechCore
/// The primary entry point for on-device speech recognition.
///
/// ```swift
/// let transcriber = try ASRTranscriber() // ASRModels.bundle in main bundle
/// let text = try transcriber.transcribe(samples) // 16 kHz mono Float32
/// ```
///
/// ## Threading
/// All public methods are thread-safe. Inference is serialized on an internal
/// queue: concurrent `transcribe` calls run one at a time in call order.
/// Synchronous methods must not be called from the main thread (they block);
/// use the `async` variants from UI code.
///
/// ## Lifecycle
/// Initialization compiles/loads Core ML models on first use per install
/// (tens of seconds on-device); subsequent launches hit the system cache.
/// Call `warmUp()` after init to move that cost off the first transcription.
public final class ASRTranscriber {
public struct Options: Sendable {
/// Cap on generated text tokens (clamped internally to decoder budget).
public var maxNewTokens = 128
/// Compute units for the audio tower. `.cpuAndNeuralEngine` is the
/// validated production path.
public var computeUnits: MLComputeUnits = .cpuAndNeuralEngine
/// Verify SHA-256 of every asset at load (~1-2 s). Enable for the
/// first launch after install/update.
public var verifyAssets = false
public init() {}
}
public struct Timings: Sendable {
public var mel: Double = 0
public var tower: Double = 0
public var decode: Double = 0
public var total: Double = 0
}
public struct Result: Sendable {
public let text: String
public let tokenIDs: [Int]
public let hitStop: Bool
public let timings: Timings
}
/// Cancellation handle for an in-flight transcription.
public final class CancellationToken: @unchecked Sendable {
private let lock = NSLock()
private var _cancelled = false
public var isCancelled: Bool {
lock.lock(); defer { lock.unlock() }
return _cancelled
}
public func cancel() {
lock.lock(); _cancelled = true; lock.unlock()
}
public init() {}
}
package let store: AssetStore
package let mel: MelExtractor
package let tower: AudioTower
package let decoder: LMDecoder
private let queue = DispatchQueue(label: "com.speechkit.asr.inference", qos: .userInitiated)
/// Minimum audio length accepted (seconds).
public static let minimumAudioSeconds = 0.25
// MARK: - Init
/// Opens `ASRModels.bundle` from the host bundle's resources.
public convenience init(options: Options = Options()) throws {
let bundle = try AssetBundle.named("ASRModels", verifyHashes: options.verifyAssets)
try self.init(assetBundle: bundle, options: options)
}
/// Opens a bundle at an explicit directory URL (e.g. downloaded models).
public convenience init(bundleURL: URL, options: Options = Options()) throws {
let bundle = try AssetBundle(at: bundleURL, verifyHashes: options.verifyAssets)
try self.init(assetBundle: bundle, options: options)
}
public init(assetBundle: AssetBundle, options: Options = Options()) throws {
self.options = options
store = try AssetStore(bundle: assetBundle)
mel = MelExtractor(hannWindow: store.hannWindow, melFilters: store.melFilters)
tower = try AudioTower(unifiedModelURL: assetBundle.url("tower"),
store: store, computeUnits: options.computeUnits)
decoder = try LMDecoder(prefillURL: assetBundle.url("lm_prefill"),
decodeURL: assetBundle.url("lm_decode"),
maskGenURL: assetBundle.url("mask_gen"),
store: store)
}
private let options: Options
/// Pre-loads ANE programs so the first transcription is fast.
/// Synchronous; call off the main thread.
public func warmUp() {
queue.sync { tower.warmUp() }
}
// MARK: - Transcription
/// Transcribes 16 kHz mono Float32 samples. Blocks until complete.
public func transcribe(_ samples: [Float],
cancellation: CancellationToken? = nil) throws -> Result {
try queue.sync {
try autoreleasepool {
try transcribeLocked(samples, cancellation: cancellation)
}
}
}
/// Async variant for Swift Concurrency callers.
public func transcribe(_ samples: [Float],
cancellation: CancellationToken? = nil) async throws -> Result {
try await withCheckedThrowingContinuation { cont in
queue.async {
do {
let r = try autoreleasepool {
try self.transcribeLocked(samples, cancellation: cancellation)
}
cont.resume(returning: r)
} catch {
cont.resume(throwing: error)
}
}
}
}
/// Transcribes an audio file (any format/rate AVFoundation can read).
public func transcribeFile(_ url: URL,
cancellation: CancellationToken? = nil) throws -> Result {
let samples = try AudioResampler.loadFile(url)
return try transcribe(samples, cancellation: cancellation)
}
// MARK: - Internal
private func transcribeLocked(_ samples: [Float],
cancellation: CancellationToken?) throws -> Result {
guard Double(samples.count) / 16_000 >= Self.minimumAudioSeconds else {
throw SpeechError.audioTooShort(minimumSeconds: Self.minimumAudioSeconds)
}
func checkCancel() throws {
if cancellation?.isCancelled == true { throw SpeechError.cancelled }
}
var t = Timings()
let tAll = Date()
var t0 = Date()
let (melFeature, encLen, bucketFrames) = mel.extract(samples)
t.mel = -t0.timeIntervalSinceNow
try checkCancel()
t0 = Date()
let masks: (attn: [Float], valid: [Bool])
let audioEmbeds: [[Float]]
do {
masks = try decoder.generateMasks(encLen: encLen)
audioEmbeds = try tower.embed(
mel: melFeature, bucketFrames: bucketFrames,
attnMask390: masks.attn, validMask390: masks.valid,
sampleCount: min(samples.count, MelExtractor.maxSamples))
} catch let e as SpeechError {
throw e
} catch {
throw SpeechError.inferenceFailed(stage: "audio-tower", underlying: error.localizedDescription)
}
t.tower = -t0.timeIntervalSinceNow
try checkCancel()
let p = store.manifest.prompt_ids
var promptIDs: [Int] = [p.user, p.bos_audio]
promptIDs.append(contentsOf: Array(repeating: p.audio, count: audioEmbeds.count))
promptIDs.append(p.eos_audio)
promptIDs.append(contentsOf: p.text)
promptIDs.append(p.assistant)
var embeds: [[Float]] = []
var cursor = 0
for id in promptIDs {
if id == p.audio {
embeds.append(audioEmbeds[cursor]); cursor += 1
} else {
embeds.append(store.embedding(for: id))
}
}
t0 = Date()
let generation: LMDecoder.GenerationResult
do {
generation = try decoder.generate(
promptEmbeds: embeds, promptIDs: promptIDs,
maxNewTokens: options.maxNewTokens,
isCancelled: { cancellation?.isCancelled == true })
} catch let e as SpeechError {
throw e
} catch {
throw SpeechError.inferenceFailed(stage: "decoder", underlying: error.localizedDescription)
}
t.decode = -t0.timeIntervalSinceNow
t.total = -tAll.timeIntervalSinceNow
return Result(text: generation.text,
tokenIDs: generation.tokenIDs,
hitStop: generation.hitStop,
timings: t)
}
}
|