File size: 5,292 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 | import Accelerate
import Foundation
import SpeechCore
/// 128-bin log-mel extractor matching HF WhisperFeatureExtractor
/// (n_fft=400, hop=160, periodic hann, reflect padding, slaney mel,
/// log10 -> clamp to max-8 -> (x+4)/4), then zero-padded to 3000 frames.
package final class MelExtractor {
package static let nFFT = 400
package static let hop = 160
package static let nMels = 128
package static let nFreq = 201 // nFFT/2 + 1
package static let framesPadded = 3000
package static let maxSamples = 480_000
/// mel-frame buckets matching the multifunction tower (5s / 10s / 30s)
package static let buckets = [500, 1000, 3000]
package static func bucket(forFrames frames: Int) -> Int {
for b in buckets where frames <= b { return b }
return buckets.last!
}
private let window: [Float] // 400
private let melFilters: [Float] // 201*128 row-major [freq][mel]
private let dftCos: [Float] // 400*201 [n][f]
private let dftSin: [Float] // 400*201
package init(hannWindow: [Float], melFilters: [Float]) {
precondition(hannWindow.count == Self.nFFT)
precondition(melFilters.count == Self.nFreq * Self.nMels)
self.window = hannWindow
self.melFilters = melFilters
var cosM = [Float](repeating: 0, count: Self.nFFT * Self.nFreq)
var sinM = [Float](repeating: 0, count: Self.nFFT * Self.nFreq)
for n in 0..<Self.nFFT {
for f in 0..<Self.nFreq {
let angle = 2.0 * Double.pi * Double(f) * Double(n) / Double(Self.nFFT)
cosM[n * Self.nFreq + f] = Float(cos(angle))
sinM[n * Self.nFreq + f] = Float(-sin(angle))
}
}
self.dftCos = cosM
self.dftSin = sinM
}
/// Returns (mel: 128*bucketFrames row-major [mel][frame], encLen, bucketFrames).
/// The output is zero-padded to the smallest bucket that fits the audio.
package func extract(_ samples: [Float]) -> (mel: [Float], encLen: Int, bucketFrames: Int) {
var x = samples
if x.count > Self.maxSamples { x = Array(x[0..<Self.maxSamples]) }
let n = x.count
let half = Self.nFFT / 2 // 200
// reflect pad
var padded = [Float](repeating: 0, count: n + Self.nFFT)
for i in 0..<half { padded[i] = x[half - i] }
for i in 0..<n { padded[half + i] = x[i] }
for i in 0..<half { padded[half + n + i] = x[n - 2 - i] }
// frame count: whisper computes 1 + n/hop frames then drops the last
let tFrames = max(n / Self.hop, 1)
// windowed frames (T,400)
var frames = [Float](repeating: 0, count: tFrames * Self.nFFT)
for t in 0..<tFrames {
let src = t * Self.hop
for i in 0..<Self.nFFT {
frames[t * Self.nFFT + i] = padded[src + i] * window[i]
}
}
// DFT via sgemm: R = F*C (T,201), I = F*S
var re = [Float](repeating: 0, count: tFrames * Self.nFreq)
var im = [Float](repeating: 0, count: tFrames * Self.nFreq)
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
Int32(tFrames), Int32(Self.nFreq), Int32(Self.nFFT),
1, frames, Int32(Self.nFFT), dftCos, Int32(Self.nFreq),
0, &re, Int32(Self.nFreq))
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
Int32(tFrames), Int32(Self.nFreq), Int32(Self.nFFT),
1, frames, Int32(Self.nFFT), dftSin, Int32(Self.nFreq),
0, &im, Int32(Self.nFreq))
// power spectrum in-place into re
vDSP.multiply(re, re, result: &re)
vDSP.multiply(im, im, result: &im)
vDSP.add(re, im, result: &re)
// mel projection: (T,201)x(201,128) -> (T,128)
var mel = [Float](repeating: 0, count: tFrames * Self.nMels)
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
Int32(tFrames), Int32(Self.nMels), Int32(Self.nFreq),
1, re, Int32(Self.nFreq), melFilters, Int32(Self.nMels),
0, &mel, Int32(Self.nMels))
// log10(max(x,1e-10)); clamp to globalMax-8; (x+4)/4
var clamped = vDSP.clip(mel, to: 1e-10...Float.greatestFiniteMagnitude)
var logMel = [Float](repeating: 0, count: clamped.count)
var count = Int32(clamped.count)
vvlog10f(&logMel, &clamped, &count)
let maxVal = vDSP.maximum(logMel)
logMel = vDSP.clip(logMel, to: (maxVal - 8.0)...Float.greatestFiniteMagnitude)
logMel = vDSP.add(4.0, logMel)
logMel = vDSP.divide(logMel, 4.0)
// transpose (T,128) -> (128, bucket) zero-padded
let bucketFrames = Self.bucket(forFrames: tFrames)
var out = [Float](repeating: 0, count: Self.nMels * bucketFrames)
let tUse = min(tFrames, bucketFrames)
for t in 0..<tUse {
for m in 0..<Self.nMels {
out[m * bucketFrames + t] = logMel[t * Self.nMels + m]
}
}
// encoder_feature_len = min(ceil(n/hop), tFrames) per runtime
let encLen = min(max((n + Self.hop - 1) / Self.hop, 1), tFrames)
return (out, encLen, bucketFrames)
}
}
|