diff --git a/Sources/FluidAudio/ASR/Paraformer/ParaformerCif.swift b/Sources/FluidAudio/ASR/Paraformer/ParaformerCif.swift index 56b81cbc..e48a4e7b 100644 --- a/Sources/FluidAudio/ASR/Paraformer/ParaformerCif.swift +++ b/Sources/FluidAudio/ASR/Paraformer/ParaformerCif.swift @@ -12,13 +12,20 @@ enum ParaformerCif { /// - encRows: encoder output rows `[T][D]` (real frames only). /// - alphas: per-frame weights `[T]` from the CifAlphas model. /// - Returns: acoustic embeddings `[L][D]`. - static func integrateAndFire(encRows: [[Float]], alphas: [Float]) -> [[Float]] { + /// CIF integrate-and-fire that also records the encoder-frame index at which + /// each acoustic token "fires" (the CIF boundary). The returned `fireFrames` + /// is the Swift port of FunASR's `pre_peak_index` and is consumed by the + /// timestamp routine (`ts_prediction_lfr6_standard`) to map tokens to time. + static func integrateAndFireWithFireFrames( + encRows: [[Float]], alphas: [Float] + ) -> (embeds: [[Float]], fireFrames: [Int]) { let threshold = ParaformerConfig.cifThreshold let tail = ParaformerConfig.cifTailThreshold let T = encRows.count let dim = encRows.first?.count ?? ParaformerConfig.encoderDim var embeds: [[Float]] = [] + var fireFrames: [Int] = [] var integrate: Float = 0 var frame = [Float](repeating: 0, count: dim) @@ -33,11 +40,17 @@ enum ParaformerCif { let used = alpha - (integrate - threshold) // portion to reach threshold for d in 0.. [[Float]] { + integrateAndFireWithFireFrames(encRows: encRows, alphas: alphas).embeds } } diff --git a/Sources/FluidAudio/ASR/Paraformer/ParaformerManager.swift b/Sources/FluidAudio/ASR/Paraformer/ParaformerManager.swift index 24f7155b..c649057f 100644 --- a/Sources/FluidAudio/ASR/Paraformer/ParaformerManager.swift +++ b/Sources/FluidAudio/ASR/Paraformer/ParaformerManager.swift @@ -1,6 +1,26 @@ @preconcurrency import CoreML +import Accelerate import Foundation +/// A recognized token together with its time span (seconds) in the source audio. +/// +/// Mirrors the per-character word-level timestamp JSON produced by FunASR's +/// Paraformer (`ts_prediction_lfr6_standard`), e.g. `m4ai/audio/x-cut.words.json`: +/// `startTime`, `endTime`, `text`. (Silence `` and the SentencePiece +/// word-boundary token `▁` are *not* emitted — segments with empty `text` are +/// skipped during decoding.) +public struct TimestampedSegment: Sendable { + public let startTime: Double + public let endTime: Double + public let text: String + + public init(startTime: Double, endTime: Double, text: String) { + self.startTime = startTime + self.endTime = endTime + self.text = text + } +} + /// Manager for Paraformer-large (zh) transcription. /// /// Pipeline: waveform -> [Preprocessor fp32/CPU] -> features @@ -51,7 +71,273 @@ public actor ParaformerManager { // 4) decoder -> logits, then greedy decode let logits = try runDecoder(encRows: encRows, validLen: T, embeds: embeds, tokenCount: L) - return decode(logits: logits, tokenCount: L, dim: dim) + return decode(logits: logits, tokenCount: L) + } + + // MARK: - Timestamps + + /// Transcribe with per-token timestamps (seconds). CIF integrate-and-fire + /// gives each token's acoustic centroid; an energy-based refinement then + /// snaps the spans to the true waveform onset/offset (see + /// `decodeWithTimestamps`). + public func transcribeWithTimestamps(audioURL: URL) throws -> [TimestampedSegment] { + let converter = AudioConverter(sampleRate: Double(ParaformerConfig.sampleRate)) + return try transcribeWithTimestamps(audio: try converter.resampleAudioFile(audioURL)) + } + + public func transcribeWithTimestamps(audio: [Float]) throws -> [TimestampedSegment] { + let dim = ParaformerConfig.encoderDim + // 1) preprocessor: waveform -> features [1, T, 560] + let features = try runPreprocessor(audio: audio) + var T = features.shape[1].intValue + if T > ParaformerConfig.decoderEncFrames { + Self.logger.warning("audio too long (\(T) frames); truncating to \(ParaformerConfig.decoderEncFrames)") + T = ParaformerConfig.decoderEncFrames + } + + // 2) encoder (padded to bucket) -> enc_out + let bucket = ParaformerConfig.pickEncoderBucket(forFrames: T) + let encOut = try runEncoder(features: features, validLen: T, bucket: bucket) + + // 3) CIF alphas (CoreML) + host integrate-and-fire (decoder embeddings) + let alphas = try runCifAlphas(encOut: encOut, validLen: T) + let encRows = rows(of: encOut, count: T, dim: dim) + let embeds = ParaformerCif.integrateAndFire(encRows: encRows, alphas: alphas) + let tokenCount = min(embeds.count, ParaformerConfig.decoderMaxTokens) + if tokenCount == 0 { return [] } + + // 4) decoder -> logits, then decode tokens + timestamps + let logits = try runDecoder( + encRows: encRows, validLen: T, + embeds: Array(embeds.prefix(tokenCount)), tokenCount: tokenCount) + return decodeWithTimestamps( + logits: logits, tokenCount: tokenCount, alphas: alphas, audio: audio) + } + + /// Greedy-decode the logits into tokens, then assign each a `[start, end]` time + /// span by faithfully porting FunASR's `ts_prediction_lfr6_standard`, with an + /// additional **energy-based boundary refinement** so the emitted spans line up + /// with the visible waveform (e.g. CapCut/剪映) instead of the CIF acoustic + /// centroids: + /// + /// * `alphas` are **upsampled 3×** (FunASR `repeat_interleave`), then + /// `cif_wo_hidden` integrates them at the 20 ms resolution that yields + /// `pre_peak_index` — these are the per-token acoustic *centroids*. + /// * CIF gives the *relative* token order/spacing (centroids); the 16 kHz + /// RMS energy envelope (smoothed) gives the *absolute* onset/offset. We + /// walk tokens left to right, each starting where the previous ended (this + /// absorbs the CIF drift), and snap its span to the energy "on" run whose + /// centre is nearest the centroid. This removes the systematic half-token + /// shift of CIF and the heuristic `force_time_shift` used by FunASR. + /// * BPE continuations (`cu@@`+`t` → `cut`) are merged; the `▁` boundary + /// and inter-token silence are NOT emitted. + private func decodeWithTimestamps( + logits: MLMultiArray, tokenCount: Int, alphas: [Float], audio: [Float] + ) -> [TimestampedSegment] { + let upsampleRate = 3 + let timeRate = 10.0 * 6.0 / 1000.0 / Double(upsampleRate) // 0.02 s + let cifThreshold: Float = 1.0 - 1e-4 + + // 1) Greedy decode every decoder position; ids stay 1:1 with the acoustic + // fire frames. Uses Accelerate (vDSP_maxvi) + stride-correct pointer read, + // matching `decode` and the SenseVoice CTC optimization. + let tokenIds = greedyArgmax(logits: logits, tokenCount: tokenCount) + + // 2) char_list: drop // (keep ▁ if present). + var charList: [String] = [] + for id in tokenIds { + if id == ParaformerConfig.blankId + || id == ParaformerConfig.sosId + || id == ParaformerConfig.eosId { continue } + guard let tok = models.vocabulary[id], !tok.isEmpty else { continue } + charList.append(tok) + } + guard charList.count >= 1 else { return [] } + + // 3) Upsample alphas 3× and run `cif_wo_hidden` to obtain the 20 ms- + // resolution fire frames (FunASR `pre_peak_index`). A CIF tail alpha is + // appended so the final boundary fire is produced. + var usAlphas: [Float] = [] + usAlphas.reserveCapacity(alphas.count * upsampleRate + 1) + for a in alphas { usAlphas.append(contentsOf: Array(repeating: a, count: upsampleRate)) } + usAlphas.append(ParaformerConfig.cifTailThreshold) + var fireIndices = Self.cifWoHiddenFireIndices(alphas: usAlphas, threshold: cifThreshold) + + // 4) Fallback (mirrors FunASR): if the fire count doesn't line up with the + // token count, renormalise the alphas so their sum equals + // len(charList)+1 and recompute. Otherwise keep the true acoustic fires + // (best match for the leading word, which carries no ▁ boundary). + if fireIndices.count != charList.count + 1 { + let target = Float(charList.count + 1) + let sum = usAlphas.reduce(0, +) + let scale = target / max(sum, 1e-6) + usAlphas = usAlphas.map { $0 * scale } + fireIndices = Self.cifWoHiddenFireIndices(alphas: usAlphas, threshold: cifThreshold) + } + guard fireIndices.count >= 2 else { return [] } + + let audioEnd = Double(audio.count) / Double(ParaformerConfig.sampleRate) + let timeOf: (Double) -> Double = { $0 * timeRate } + // Acoustic centroids in seconds (no force_time_shift — energy refinement + // replaces it). + let centroids: [Double] = fireIndices.map { timeOf(Double($0)) } + + // 5) Build the token layout via sequential energy forced-alignment. + // CIF gives the *relative* token order and spacing (centroids); the 16 kHz + // RMS envelope (smoothed) gives the *absolute* onset/offset. We walk + // tokens left to right, each starting where the previous ended (this + // absorbs the CIF drift), and snap its span to the energy "on" run whose + // centre is nearest the centroid. No silence entries are emitted. + let rawEnv = Self.energyEnvelope(audio: audio, sampleRate: Double(ParaformerConfig.sampleRate), hopSec: 0.01) + let env = Self.smooth(rawEnv, window: 3) + let floor = env.isEmpty ? 0 : Self.percentile(env, q: 0.1) + let energyThreshold = max(floor * 2.5, 1e-4) + let minRun = 3 // frames (~30 ms) — reject single-frame noise spikes + let n = min(charList.count, centroids.count - 1) + var raw: [(text: String, start: Double, end: Double)] = [] + var cursor: Double = 0.0 + for i in 0..