diff --git a/Sources/FluidAudio/ASR/AsrManager.swift b/Sources/FluidAudio/ASR/AsrManager.swift index 0202baafc..81b67e880 100644 --- a/Sources/FluidAudio/ASR/AsrManager.swift +++ b/Sources/FluidAudio/ASR/AsrManager.swift @@ -41,6 +41,8 @@ public final class AsrManager { // Vocabulary boosting state (configured via configureVocabularyBoosting) // Internal access required for AsrTranscription extension (separate file) + internal var vadManager: VadManager? + internal var customVocabulary: CustomVocabularyContext? internal var ctcSpotter: CtcKeywordSpotter? internal var vocabularyRescorer: VocabularyRescorer? @@ -110,7 +112,9 @@ public final class AsrManager { self.jointModel = models.joint self.vocabulary = models.vocabulary - logger.info("Token duration optimization model loaded successfully") + if vadManager == nil { + vadManager = try? await VadManager() + } logger.info("AsrManager initialized successfully with provided models") } @@ -306,13 +310,42 @@ public final class AsrManager { isLastChunk: Bool = false, globalFrameOffset: Int = 0 ) async throws -> TdtHypothesis { - // Route to appropriate decoder based on model version + try await tdtDecodeWithTimings( + encoderOutput: encoderOutput, + encoderSequenceLength: encoderSequenceLength, + actualAudioFrames: actualAudioFrames, + decoderState: &decoderState, + overrideTdtConfig: nil, + contextFrameAdjustment: contextFrameAdjustment, + isLastChunk: isLastChunk, + globalFrameOffset: globalFrameOffset + ) + } + + /// Diagnostic overload: allows substituting a custom TdtConfig (e.g., with blankBiasTokens) + /// while reusing the manager's already-loaded decoder/joint models. + internal func tdtDecodeWithTimings( + encoderOutput: MLMultiArray, + encoderSequenceLength: Int, + actualAudioFrames: Int, + decoderState: inout TdtDecoderState, + overrideTdtConfig: TdtConfig?, + contextFrameAdjustment: Int = 0, + isLastChunk: Bool = false, + globalFrameOffset: Int = 0 + ) async throws -> TdtHypothesis { guard let models = asrModels, let decoder_ = decoderModel, let joint = jointModel else { throw ASRError.notInitialized } + let effectiveConfig: ASRConfig + if let override = overrideTdtConfig { + effectiveConfig = ASRConfig(tdtConfig: override) + } else { + effectiveConfig = config + } switch models.version { case .v2: - let decoder = TdtDecoderV2(config: config) + let decoder = TdtDecoderV2(config: effectiveConfig) return try await decoder.decodeWithTimings( encoderOutput: encoderOutput, encoderSequenceLength: encoderSequenceLength, @@ -325,7 +358,7 @@ public final class AsrManager { globalFrameOffset: globalFrameOffset ) case .v3: - let decoder = TdtDecoderV3(config: config) + let decoder = TdtDecoderV3(config: effectiveConfig) return try await decoder.decodeWithTimings( encoderOutput: encoderOutput, encoderSequenceLength: encoderSequenceLength, diff --git a/Sources/FluidAudio/ASR/AsrTranscription.swift b/Sources/FluidAudio/ASR/AsrTranscription.swift index 0c8c2ac51..ff13e2773 100644 --- a/Sources/FluidAudio/ASR/AsrTranscription.swift +++ b/Sources/FluidAudio/ASR/AsrTranscription.swift @@ -45,6 +45,22 @@ extension AsrManager { processingTime: Date().timeIntervalSince(startTime) ) + // Recover dictation commands (period, comma, end quote, etc.) dropped by TDT + if let repairedText = await repairDictationGaps( + hypothesis: hypothesis, audioSamples: audioSamples) + { + result = ASRResult( + text: repairedText, + confidence: result.confidence, + duration: result.duration, + processingTime: result.processingTime, + tokenTimings: result.tokenTimings, + performanceMetrics: result.performanceMetrics, + ctcDetectedTerms: result.ctcDetectedTerms, + ctcAppliedTerms: result.ctcAppliedTerms + ) + } + // Auto-apply vocabulary rescoring when configured if vocabBoostingEnabled { result = await applyVocabularyRescoring(result: result, audioSamples: audioSamples) @@ -146,6 +162,45 @@ extension AsrManager { } } + /// Run only the preprocessor + encoder stages and return the raw encoder output. + /// Used by diagnostic tests to run the decoder multiple times with different configs + /// without re-encoding the audio each time. + /// + /// The encoder output is deep-copied into a new MLMultiArray so it remains valid + /// after the CoreML prediction provider is released. + internal func runEncoderOnly( + _ audioSamples: [Float] + ) async throws -> (encoderOutput: MLMultiArray, encoderSequenceLength: Int, actualAudioFrames: Int) { + guard let preprocessorModel = preprocessorModel, let encoderModel = encoderModel else { + throw ASRError.notInitialized + } + let originalLength = audioSamples.count + let padded = padAudioIfNeeded(audioSamples, targetLength: ASRConstants.maxModelSamples) + let preprocessorInput = try await preparePreprocessorInput(padded, actualLength: originalLength) + let preprocessorOutput = try await preprocessorModel.compatPrediction( + from: preprocessorInput, options: predictionOptions) + let encoderInput = try prepareEncoderInput( + encoder: encoderModel, preprocessorOutput: preprocessorOutput, + originalInput: preprocessorInput) + let encoderOutputProvider = try await encoderModel.compatPrediction( + from: encoderInput, options: predictionOptions) + let rawEncoderOutput = try extractFeatureValue( + from: encoderOutputProvider, key: "encoder", errorMessage: "Invalid encoder output") + let encoderLength = try extractFeatureValue( + from: encoderOutputProvider, key: "encoder_length", + errorMessage: "Invalid encoder output length") + let seqLen = encoderLength[0].intValue + let actualFrames = ASRConstants.calculateEncoderFrames(from: originalLength) + + // Deep-copy the encoder output: CoreML may back the MLMultiArray with its own + // internal buffer that becomes invalid once the prediction provider is released. + let copiedOutput = try MLMultiArray(shape: rawEncoderOutput.shape, dataType: rawEncoderOutput.dataType) + let byteCount = rawEncoderOutput.count * MemoryLayout.stride + memcpy(copiedOutput.dataPointer, rawEncoderOutput.dataPointer, byteCount) + + return (copiedOutput, seqLen, actualFrames) + } + private func prepareEncoderInput( encoder: MLModel, preprocessorOutput: MLFeatureProvider, diff --git a/Sources/FluidAudio/ASR/DictationRepair.swift b/Sources/FluidAudio/ASR/DictationRepair.swift new file mode 100644 index 000000000..457c23b5d --- /dev/null +++ b/Sources/FluidAudio/ASR/DictationRepair.swift @@ -0,0 +1,127 @@ +import Foundation + +extension AsrManager { + + internal static let dictationRepairGapThreshold = 12 // 960ms at 80ms/frame + private static let frameSeconds: Double = 0.08 + private static let samplesPerFrame = 1280 // 80ms @ 16kHz + + private static let dictationCommands: Set = [ + "period", "comma", "semicolon", "colon", + "exclamation point", "question mark", + "quote", "unquote", "end quote", "open quote", "close quote", + "new line", "newline", "new paragraph", "dash", "hyphen", + ] + + /// Detects blank gaps ≥960ms with VAD-confirmed speech, re-runs TDT on each gap window, + /// and inserts recovered dictation commands (period, comma, end quote, etc.) into the text. + /// + /// Returns the repaired text, or nil if nothing was changed. + internal func repairDictationGaps( + hypothesis: TdtHypothesis, + audioSamples: [Float] + ) async -> String? { + guard let vad = vadManager else { return nil } + + let ts = hypothesis.timestamps + guard ts.count >= 2 else { return nil } + + // Find blank gaps ≥ threshold + struct Gap { + let afterTokenIdx: Int + let afterFrame: Int + let beforeFrame: Int + } + var gaps: [Gap] = [] + for i in 0..<(ts.count - 1) { + let gf = ts[i + 1] - ts[i] + if gf >= Self.dictationRepairGapThreshold { + gaps.append(Gap(afterTokenIdx: i, afterFrame: ts[i], beforeFrame: ts[i + 1])) + } + } + guard !gaps.isEmpty else { return nil } + + // Run VAD on full audio once + guard let vadResults = try? await vad.process(audioSamples) else { return nil } + let chunkDuration = Double(VadManager.chunkSize) / 16_000.0 + + func hasSpeech(t1: Double, t2: Double) -> Bool { + for (ci, r) in vadResults.enumerated() { + let cs = Double(ci) * chunkDuration + let ce = cs + chunkDuration + guard cs <= t2 && ce >= t1 else { continue } + if r.isVoiceActive { return true } + } + return false + } + + // For each speech-confirmed gap, run window TDT and check for a dictation command + var repairs: [Int: String] = [:] + for gap in gaps { + let t1 = Double(gap.afterFrame + 1) * Self.frameSeconds + let t2 = Double(gap.beforeFrame) * Self.frameSeconds + guard hasSpeech(t1: t1, t2: t2) else { continue } + + let winStart = (gap.afterFrame + 3) * Self.samplesPerFrame + let winEnd = min((gap.beforeFrame + 2) * Self.samplesPerFrame, audioSamples.count) + guard winStart < audioSamples.count, winStart < winEnd else { continue } + + let winSamples = Array(audioSamples[winStart.." + } + + // ── Find blank gaps ─────────────────────────────────────── + let ts = hypothesis.timestamps + struct Gap { + let afterIdx: Int + let afterFrame: Int + let beforeFrame: Int + let gapFrames: Int + } + var gaps: [Gap] = [] + if ts.count >= 2 { + for gi in 0..<(ts.count - 1) { + let gf = ts[gi + 1] - ts[gi] + if gf >= Self.gapThresholdFrames { + gaps.append(Gap(afterIdx: gi, afterFrame: ts[gi], beforeFrame: ts[gi + 1], gapFrames: gf)) + } + } + } + + guard !gaps.isEmpty else { + results.append( + makeResult( + index: idx, category: category, input: inputPrompt, + expected: expected, rawSTT: rawSTT, recordedAt: recordedAt, + baseline: baseline, repaired: baseline, + repairApplied: false, error: nil)) + continue + } + + // ── Full-audio VAD ──────────────────────────────────────── + let fullVad = try await vad.process(samples) + func hasSpeechInRange(t1: Double, t2: Double) -> (Bool, Float) { + var maxP: Float = 0 + var active = false + for (ci, r) in fullVad.enumerated() { + let cs = Double(ci) * chunkDuration + let ce = cs + chunkDuration + guard cs <= t2 && ce >= t1 else { continue } + if r.probability > maxP { maxP = r.probability } + if r.isVoiceActive { active = true } + } + return (active, maxP) + } + + // ── Per-gap: VAD check + isolated TDT for every gap candidate ── + // We save full metadata so the viewer can apply different thresholds + // without re-running inference. + struct GapResult { + let afterTokenIdx: Int + let afterFrame: Int + let beforeFrame: Int + let gapFrames: Int + let vadMaxProb: Float + let hasSpeech: Bool + let recovered: String // "" if TDT found nothing + } + + var gapResults: [GapResult] = [] + for gap in gaps { + let gapT1 = Double(gap.afterFrame + 1) * Self.frameSeconds + let gapT2 = Double(gap.beforeFrame) * Self.frameSeconds + let (hasSpeech, vadProb) = hasSpeechInRange(t1: gapT1, t2: gapT2) + + var recovered = "" + if hasSpeech { + let winStart = (gap.afterFrame + 3) * Self.samplesPerFrame + let winEnd = min((gap.beforeFrame + 2) * Self.samplesPerFrame, samples.count) + if winStart < samples.count && winStart < winEnd { + let winSamples = Array(samples[winStart..= defaultThreshold && !g.recovered.isEmpty { + repairs[g.afterTokenIdx] = g.recovered + } + + var parts: [String] = [] + for (ti, tokenId) in hypothesis.ySequence.enumerated() { + let name = vocab[tokenId]?.replacingOccurrences(of: "▁", with: " ") ?? "<\(tokenId)>" + parts.append(name) + if let inserted = repairs[ti] { + parts.append(" \(inserted)") + } + } + let repaired = parts.joined() + .replacingOccurrences(of: " ", with: " ") + .trimmingCharacters(in: .whitespaces) + + let didRepair = !repairs.isEmpty + if didRepair { repairFired += 1 } + + let didChange = normalize(repaired) != normalize(baseline) + if didChange { changed += 1 } + + var r = makeResult( + index: idx, category: category, input: inputPrompt, + expected: expected, rawSTT: rawSTT, recordedAt: recordedAt, + baseline: baseline, repaired: repaired, + repairApplied: didRepair, error: nil) + r["gaps"] = gapData + r["tokenTexts"] = tokenTexts + r["tokenFrames"] = hypothesis.timestamps + results.append(r) + + } catch { + results.append( + makeResult( + index: idx, category: category, input: inputPrompt, + expected: expected, rawSTT: rawSTT, recordedAt: recordedAt, + baseline: nil, repaired: nil, + repairApplied: false, error: error.localizedDescription)) + errors += 1 + } + } + + // ── Write output ───────────────────────────────────────────────── + let summary: [String: Any] = [ + "total": allEntries.count, + "repairFired": repairFired, + "changed": changed, + "errors": errors, + ] + let output: [String: Any] = ["summary": summary, "results": results] + let data = try JSONSerialization.data(withJSONObject: output, options: [.sortedKeys]) + try data.write(to: URL(fileURLWithPath: Self.outputPath)) + + print("\n=== SUMMARY ===") + print("Total entries : \(allEntries.count)") + print("Repair fired : \(repairFired)") + print("Output changed : \(changed)") + print("Errors : \(errors)") + print("Written to : \(Self.outputPath) (\(data.count / 1024)KB)") + } + + private func normalize(_ s: String) -> String { + s.lowercased().trimmingCharacters(in: .whitespacesAndNewlines) + } + + private func makeResult( + index: Int, category: String, input: String, expected: String, + rawSTT: String, recordedAt: String, + baseline: String?, repaired: String?, + repairApplied: Bool, error: String? + ) -> [String: Any] { + var d: [String: Any] = [ + "index": index, + "category": category, + "input": input, + "expected": expected, + "rawSTT": rawSTT, + "recordedAt": recordedAt, + "baseline": baseline ?? "", + "repaired": repaired ?? "", + "repairApplied": repairApplied, + "changed": (baseline != nil && repaired != nil && baseline != repaired), + ] + if let e = error { d["error"] = e } + return d + } +} diff --git a/Tests/FluidAudioTests/ASR/TDT/HistoryAudioGapScanTests.swift b/Tests/FluidAudioTests/ASR/TDT/HistoryAudioGapScanTests.swift new file mode 100644 index 000000000..43fabd22c --- /dev/null +++ b/Tests/FluidAudioTests/ASR/TDT/HistoryAudioGapScanTests.swift @@ -0,0 +1,346 @@ +import Foundation +import XCTest + +@testable import FluidAudio + +/// Scans all history audio files from ~/Library/Application Support/Onit/transcription_history_audio/ +/// and reports what percentage would trigger the ≥12f blank-gap window inference. +/// +/// Run with: +/// cd FluidAudio && swift test --filter HistoryAudioGapScanTests +final class HistoryAudioGapScanTests: XCTestCase { + + private static let audioDir = + FileManager.default.homeDirectoryForCurrentUser + .appendingPathComponent("Library/Application Support/Onit/transcription_history_audio") + private static let outputPath = "/tmp/history_gap_scan.json" + private static let recoveryOutputPath = "/tmp/history_gap_recovery.json" + + private static let gapThresholdFrames = 12 // 960ms — the production threshold + private static let frameSeconds = 0.08 + private static let samplesPerFrame = 1280 // 80ms @ 16kHz + + func testScanHistoryAudioGaps() async throws { + guard FileManager.default.fileExists(atPath: Self.audioDir.path) else { + throw XCTSkip("History audio dir not found: \(Self.audioDir.path)") + } + + let allFiles = try FileManager.default.contentsOfDirectory( + at: Self.audioDir, includingPropertiesForKeys: nil + ).filter { $0.pathExtension == "wav" }.sorted { $0.lastPathComponent < $1.lastPathComponent } + + print("Found \(allFiles.count) history audio files") + + let tdtModels = try await AsrModels.loadFromCache() + let manager = AsrManager() + try await manager.initialize(models: tdtModels) + + let vad = try await VadManager() + let chunkDuration = Double(VadManager.chunkSize) / 16000.0 + + let converter = AudioConverter() + + var totalFiles = 0 + var errors = 0 + var filesWithGap = 0 // ≥12f gap exists at all + var filesWithSpeechGap = 0 // ≥12f gap + VAD confirms speech + var totalSpeechGaps = 0 // total qualifying gaps across all files + var gapFrameDistribution: [Int: Int] = [:] + var results: [[String: Any]] = [] + + for (i, url) in allFiles.enumerated() { + if (i + 1) % 100 == 0 || i == 0 { + print("[\(i+1)/\(allFiles.count)] processed=\(totalFiles) " + + "triggered=\(filesWithSpeechGap) errors=\(errors)") + } + + do { + let samples = try converter.resampleAudioFile(path: url.path) + totalFiles += 1 + + // ── TDT inference ───────────────────────────────────────── + let padded = manager.padAudioIfNeeded(samples, targetLength: ASRConstants.maxModelSamples) + var state = TdtDecoderState.make() + let (hypothesis, _) = try await manager.executeMLInferenceWithTimings( + padded, + originalLength: samples.count, + actualAudioFrames: nil, + decoderState: &state, + contextFrameAdjustment: 0, + isLastChunk: true + ) + + let ts = hypothesis.timestamps + guard ts.count >= 2 else { continue } + + // ── Find gaps ≥ threshold ───────────────────────────────── + struct Gap { + let afterFrame: Int + let beforeFrame: Int + let gapFrames: Int + } + var gaps: [Gap] = [] + for gi in 0..<(ts.count - 1) { + let gf = ts[gi + 1] - ts[gi] + if gf >= Self.gapThresholdFrames { + gaps.append(Gap(afterFrame: ts[gi], beforeFrame: ts[gi + 1], gapFrames: gf)) + } + } + + guard !gaps.isEmpty else { continue } + filesWithGap += 1 + + // ── VAD on full audio ───────────────────────────────────── + let fullVad = try await vad.process(samples) + func hasSpeechInRange(t1: Double, t2: Double) -> (Bool, Float) { + var maxP: Float = 0 + var active = false + for (ci, r) in fullVad.enumerated() { + let cs = Double(ci) * chunkDuration + let ce = cs + chunkDuration + guard cs <= t2 && ce >= t1 else { continue } + if r.probability > maxP { maxP = r.probability } + if r.isVoiceActive { active = true } + } + return (active, maxP) + } + + var speechGaps: [[String: Any]] = [] + for gap in gaps { + let t1 = Double(gap.afterFrame + 1) * Self.frameSeconds + let t2 = Double(gap.beforeFrame) * Self.frameSeconds + let (hasSpeech, vadProb) = hasSpeechInRange(t1: t1, t2: t2) + gapFrameDistribution[gap.gapFrames, default: 0] += 1 + if hasSpeech { + totalSpeechGaps += 1 + speechGaps.append([ + "gapFrames": gap.gapFrames, + "gapSeconds": Double(gap.gapFrames) * Self.frameSeconds, + "vadMaxProb": Double(vadProb), + ]) + } + } + + if !speechGaps.isEmpty { + filesWithSpeechGap += 1 + results.append([ + "file": url.lastPathComponent, + "durationSamples": samples.count, + "durationSeconds": Double(samples.count) / 16000.0, + "gaps": speechGaps, + ]) + } + + } catch { + errors += 1 + } + } + + // ── Summary ─────────────────────────────────────────────────────── + let pctTriggered = totalFiles > 0 ? 100.0 * Double(filesWithSpeechGap) / Double(totalFiles) : 0 + let summary: [String: Any] = [ + "totalFiles": totalFiles, + "errors": errors, + "filesWithGap_noVad": filesWithGap, + "filesWithSpeechGap": filesWithSpeechGap, + "pctTriggered": pctTriggered, + "totalSpeechGaps": totalSpeechGaps, + "gapThresholdFrames": Self.gapThresholdFrames, + ] + + let output: [String: Any] = ["summary": summary, "triggered": results] + let data = try JSONSerialization.data(withJSONObject: output, options: [.sortedKeys, .prettyPrinted]) + try data.write(to: URL(fileURLWithPath: Self.outputPath)) + + print("\n=== GAP SCAN SUMMARY (≥\(Self.gapThresholdFrames)f / \(Int(Double(Self.gapThresholdFrames) * Self.frameSeconds * 1000))ms) ===") + print("Total files processed : \(totalFiles)") + print("Errors : \(errors)") + print("Files with gap (raw) : \(filesWithGap)") + print("Files with speech gap : \(filesWithSpeechGap) (\(String(format: "%.1f", pctTriggered))%)") + print("Total speech gaps : \(totalSpeechGaps)") + print("Written to : \(Self.outputPath)") + + print("\nGap size distribution (≥\(Self.gapThresholdFrames)f):") + for size in gapFrameDistribution.keys.sorted() { + let count = gapFrameDistribution[size]! + print(" \(size)f (\(String(format: "%.2f", Double(size) * Self.frameSeconds))s): \(count)") + } + } + + /// Runs window TDT on only the files that triggered in testScanHistoryAudioGaps, + /// to see what gets recovered and whether any are dictation commands. + /// + /// Run with: + /// cd FluidAudio && swift test --filter HistoryAudioGapScanTests/testRecoverTriggeredGaps + func testRecoverTriggeredGaps() async throws { + guard FileManager.default.fileExists(atPath: Self.outputPath) else { + throw XCTSkip("Run testScanHistoryAudioGaps first to generate \(Self.outputPath)") + } + + let scanData = try Data(contentsOf: URL(fileURLWithPath: Self.outputPath)) + let scan = try JSONSerialization.jsonObject(with: scanData) as! [String: Any] + let triggered = scan["triggered"] as! [[String: Any]] + + print("Running window TDT on \(triggered.count) triggered files…") + + let tdtModels = try await AsrModels.loadFromCache() + let manager = AsrManager() + try await manager.initialize(models: tdtModels) + let converter = AudioConverter() + + let dictationCommands: Set = [ + "period", "comma", "semicolon", "colon", + "exclamation point", "question mark", + "quote", "unquote", "end quote", "open quote", "close quote", + "new line", "newline", "new paragraph", "dash", "hyphen", + ] + + func normalize(_ s: String) -> String { + s.lowercased() + .replacingOccurrences(of: "[^a-z0-9 ]", with: " ", options: .regularExpression) + .components(separatedBy: .whitespaces).filter { !$0.isEmpty }.joined(separator: " ") + } + + var results: [[String: Any]] = [] + var dictationHits = 0 + var windowInferenceMs: [Double] = [] + var errors = 0 + + for (i, entry) in triggered.enumerated() { + let fileName = entry["file"] as! String + let fileURL = Self.audioDir.appendingPathComponent(fileName) + let gaps = entry["gaps"] as! [[String: Any]] + + if (i + 1) % 20 == 0 || i == 0 { + print("[\(i+1)/\(triggered.count)] hits=\(dictationHits)") + } + + do { + let samples = try converter.resampleAudioFile(path: fileURL.path) + + // Full TDT to get token timestamps + let padded = manager.padAudioIfNeeded(samples, targetLength: ASRConstants.maxModelSamples) + var state = TdtDecoderState.make() + let (hypothesis, _) = try await manager.executeMLInferenceWithTimings( + padded, + originalLength: samples.count, + actualAudioFrames: nil, + decoderState: &state, + contextFrameAdjustment: 0, + isLastChunk: true + ) + let ts = hypothesis.timestamps + guard ts.count >= 2 else { continue } + + let baseline = manager.processTranscriptionResult( + tokenIds: hypothesis.ySequence, + timestamps: hypothesis.timestamps, + confidences: hypothesis.tokenConfidences, + encoderSequenceLength: 0, + audioSamples: samples, + processingTime: 0 + ).text + + // Re-find qualifying gaps + var gapRecoveries: [[String: Any]] = [] + for gi in 0..<(ts.count - 1) { + let gf = ts[gi + 1] - ts[gi] + guard gf >= Self.gapThresholdFrames else { continue } + + let afterFrame = ts[gi] + let beforeFrame = ts[gi + 1] + + // Only process gaps that VAD confirmed as speech in the scan + guard gaps.contains(where: { ($0["gapFrames"] as? Int) == gf }) else { continue } + + let winStart = (afterFrame + 3) * Self.samplesPerFrame + let winEnd = min((beforeFrame + 2) * Self.samplesPerFrame, samples.count) + guard winStart < samples.count && winStart < winEnd else { continue } + + let winSamples = Array(samples[winStart..