File size: 5,652 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 | import Accelerate
import CoreML
import Foundation
import SpeechCore
/// Unified audio tower (multifunction Core ML model, ANE):
/// mel bucket -> function tower_5s / tower_10s / tower_30s -> hidden
/// -> projector (CPU). Masks come from LMDecoder.maskGen (390-sized, sliced).
package final class AudioTower {
private let modelURL: URL
private let computeUnits: MLComputeUnits
private var modelsByBucket: [Int: MLModel] = [:]
private let store: AssetStore
private static let functionByBucket = [500: "tower_5s", 1000: "tower_10s", 3000: "tower_30s"]
package static let tokensByBucket = [500: 65, 1000: 130, 3000: 390]
package init(unifiedModelURL: URL, store: AssetStore,
computeUnits: MLComputeUnits = .cpuAndNeuralEngine) throws {
self.computeUnits = computeUnits
self.store = store
if unifiedModelURL.pathExtension == "mlpackage" {
modelURL = try MLModel.compileModel(at: unifiedModelURL)
} else {
modelURL = unifiedModelURL
}
}
/// Loads (and caches) the function for a bucket. First load per bucket
/// triggers the ANE compile; call warmUp() at app start.
private func model(forBucket bucket: Int) throws -> MLModel {
if let m = modelsByBucket[bucket] { return m }
let config = MLModelConfiguration()
config.computeUnits = computeUnits
config.functionName = Self.functionByBucket[bucket]!
let m = try MLModel(contentsOf: modelURL, configuration: config)
modelsByBucket[bucket] = m
return m
}
package func warmUp(buckets: [Int] = [500, 3000]) {
for b in buckets { _ = try? model(forBucket: b) }
}
private static func multiArray(_ values: [Float], shape: [NSNumber]) throws -> MLMultiArray {
let arr = try MLMultiArray(shape: shape, dataType: .float32)
values.withUnsafeBufferPointer { src in
arr.dataPointer.bindMemory(to: Float.self, capacity: values.count)
.update(from: src.baseAddress!, count: values.count)
}
return arr
}
private static func floats(_ arr: MLMultiArray) -> [Float] {
switch arr.dataType {
case .float32:
return arr.withUnsafeBufferPointer(ofType: Float.self) { Array($0) }
case .float16:
return arr.withUnsafeBufferPointer(ofType: Float16.self) { $0.map(Float.init) }
default:
return (0..<arr.count).map { arr[$0].floatValue }
}
}
/// mel (128*bucketFrames) + 390-sized masks -> projected embeddings (N, 512)
package func embed(mel: [Float], bucketFrames: Int, attnMask390: [Float],
validMask390: [Bool], sampleCount: Int) throws -> [[Float]] {
let tokens = Self.tokensByBucket[bucketFrames]!
// slice the 390x390 additive mask to tokens x tokens, fp16-safe fill
var mask = [Float](repeating: 0, count: tokens * tokens)
for q in 0..<tokens {
for k in 0..<tokens {
mask[q * tokens + k] = max(attnMask390[q * 390 + k], -3e4)
}
}
let m = try model(forBucket: bucketFrames)
let out = try m.prediction(from: MLDictionaryFeatureProvider(dictionary: [
"audios": Self.multiArray(mel, shape: [1, 128, NSNumber(value: bucketFrames)]),
"attn_mask": Self.multiArray(mask, shape: [1, 1, NSNumber(value: tokens), NSNumber(value: tokens)]),
]))
let hidden = Self.floats(out.featureValue(for: "hidden")!.multiArrayValue!) // tokens*1024
let dim = 1024
var valid: [[Float]] = []
for i in 0..<tokens where validMask390[i] {
valid.append(Array(hidden[(i * dim)..<((i + 1) * dim)]))
}
let target = Self.arkAudioTokenCount(sampleCount: sampleCount)
if valid.count != target {
valid = Self.adaptiveAvgPool(valid, outputSize: target)
}
return valid.map { row in
let normed = Self.layerNorm(row, weight: store.projNormW, bias: store.projNormB)
var out = store.projLinB
cblas_sgemv(CblasRowMajor, CblasNoTrans, 512, 1024,
1, store.projLinW, 1024, normed, 1, 1, &out, 1)
return out
}
}
package static func arkAudioTokenCount(sampleCount: Int, hop: Int = 160, mergeFactor: Int = 4) -> Int {
let melFrames = sampleCount / hop
let downsampled = (melFrames + 1) / 2
return max(downsampled / mergeFactor, 1)
}
static func adaptiveAvgPool(_ x: [[Float]], outputSize: Int) -> [[Float]] {
let inputSize = x.count
guard inputSize != outputSize, inputSize > 0 else { return x }
let dim = x[0].count
var out: [[Float]] = []
out.reserveCapacity(outputSize)
for i in 0..<outputSize {
let start = i * inputSize / outputSize
var end = (((i + 1) * inputSize) + outputSize - 1) / outputSize
end = max(end, start + 1)
var acc = [Float](repeating: 0, count: dim)
for r in start..<end { vDSP.add(acc, x[r], result: &acc) }
out.append(vDSP.divide(acc, Float(end - start)))
}
return out
}
static func layerNorm(_ x: [Float], weight: [Float], bias: [Float], eps: Float = 1e-5) -> [Float] {
let mean = vDSP.mean(x)
let centered = vDSP.add(-mean, x)
let variance = vDSP.meanSquare(centered)
let inv = 1.0 / sqrt(variance + eps)
var out = vDSP.multiply(inv, centered)
out = vDSP.multiply(out, weight)
return vDSP.add(out, bias)
}
}
|