feat: recover dropped dictation commands via blank-gap TDT repair

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,

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<Float>.stride
+        memcpy(copiedOutput.dataPointer, rawEncoderOutput.dataPointer, byteCount)
+
+        return (copiedOutput, seqLen, actualFrames)
+    }
+
     private func prepareEncoderInput(
         encoder: MLModel,
         preprocessorOutput: MLFeatureProvider,

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<String> = [
+        "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..<winEnd])
+            let paddedWin = padAudioIfNeeded(winSamples, targetLength: ASRConstants.maxModelSamples)
+            var winState = TdtDecoderState.make()
+            guard
+                let (winHyp, _) = try? await executeMLInferenceWithTimings(
+                    paddedWin,
+                    originalLength: winSamples.count,
+                    actualAudioFrames: nil,
+                    decoderState: &winState,
+                    contextFrameAdjustment: 0,
+                    isLastChunk: true
+                )
+            else { continue }
+
+            let recovered = processTranscriptionResult(
+                tokenIds: winHyp.ySequence,
+                timestamps: winHyp.timestamps,
+                confidences: winHyp.tokenConfidences,
+                tokenDurations: winHyp.tokenDurations,
+                encoderSequenceLength: 0,
+                audioSamples: winSamples,
+                processingTime: 0
+            ).text.trimmingCharacters(in: .whitespacesAndNewlines)
+
+            let normalized =
+                recovered
+                .lowercased()
+                .replacingOccurrences(of: "[^a-z0-9 ]", with: " ", options: .regularExpression)
+                .components(separatedBy: .whitespaces)
+                .filter { !$0.isEmpty }
+                .joined(separator: " ")
+
+            if Self.dictationCommands.contains(normalized) {
+                repairs[gap.afterTokenIdx] = recovered
+                logger.info(
+                    "Dictation repair: inserted \"\(recovered)\" after token \(gap.afterTokenIdx) (gap=\(ts[gap.afterTokenIdx + 1] - ts[gap.afterTokenIdx])f)"
+                )
+            }
+        }
+
+        guard !repairs.isEmpty else { return nil }
+
+        // Reconstruct text from vocabulary tokens with insertions at repair points
+        var parts: [String] = []
+        for (i, tokenId) in hypothesis.ySequence.enumerated() {
+            let tok = vocabulary[tokenId]?.replacingOccurrences(of: "▁", with: " ") ?? ""
+            parts.append(tok)
+            if let inserted = repairs[i] {
+                parts.append(" \(inserted)")
+            }
+        }
+
+        let repaired = parts.joined()
+            .replacingOccurrences(of: "  +", with: " ", options: .regularExpression)
+            .trimmingCharacters(in: .whitespaces)
+
+        return repaired
+    }
+}