diff --git a/Sources/FluidAudio/ASR/Parakeet/AsrTypes.swift b/Sources/FluidAudio/ASR/Parakeet/AsrTypes.swift index 7b849b0fc..0030f6697 100644 --- a/Sources/FluidAudio/ASR/Parakeet/AsrTypes.swift +++ b/Sources/FluidAudio/ASR/Parakeet/AsrTypes.swift @@ -96,13 +96,19 @@ public struct ASRResult: Codable, Sendable { public let performanceMetrics: ASRPerformanceMetrics? public let ctcDetectedTerms: [String]? public let ctcAppliedTerms: [String]? + /// The rescoring decision behind each applied replacement, aligned 1:1 with + /// `ctcAppliedTerms`. Carries the decoded word each term displaced and the + /// scores the decision was made on, so callers can tell a recognizer error + /// from a bad vocabulary replacement. `nil` when boosting is not configured. + public let ctcReplacements: [VocabularyRescorer.RescoringResult]? public init( text: String, confidence: Float, duration: TimeInterval, processingTime: TimeInterval, tokenTimings: [TokenTiming]? = nil, performanceMetrics: ASRPerformanceMetrics? = nil, ctcDetectedTerms: [String]? = nil, - ctcAppliedTerms: [String]? = nil + ctcAppliedTerms: [String]? = nil, + ctcReplacements: [VocabularyRescorer.RescoringResult]? = nil ) { self.text = text self.confidence = confidence @@ -112,6 +118,7 @@ public struct ASRResult: Codable, Sendable { self.performanceMetrics = performanceMetrics self.ctcDetectedTerms = ctcDetectedTerms self.ctcAppliedTerms = ctcAppliedTerms + self.ctcReplacements = ctcReplacements } /// Real-time factor (RTFx) - how many times faster than real-time @@ -125,8 +132,12 @@ public struct ASRResult: Codable, Sendable { /// - text: The rescored transcript text /// - detected: Vocabulary terms detected by CTC (candidates considered for replacement) /// - applied: Vocabulary terms actually applied as replacements + /// - replacements: The rescoring decision behind each applied term /// - Returns: A new ASRResult with updated text and CTC metadata - public func withRescoring(text: String, detected: [String]?, applied: [String]?) -> ASRResult { + public func withRescoring( + text: String, detected: [String]?, applied: [String]?, + replacements: [VocabularyRescorer.RescoringResult]? = nil + ) -> ASRResult { ASRResult( text: text, confidence: confidence, @@ -135,7 +146,8 @@ public struct ASRResult: Codable, Sendable { tokenTimings: tokenTimings, performanceMetrics: performanceMetrics, ctcDetectedTerms: detected, - ctcAppliedTerms: applied + ctcAppliedTerms: applied, + ctcReplacements: replacements ) } } diff --git a/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/CustomVocabulary/Rescorer/VocabularyRescorer.swift b/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/CustomVocabulary/Rescorer/VocabularyRescorer.swift index b24e2d69f..465e20278 100644 --- a/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/CustomVocabulary/Rescorer/VocabularyRescorer.swift +++ b/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/CustomVocabulary/Rescorer/VocabularyRescorer.swift @@ -197,7 +197,7 @@ public struct VocabularyRescorer: Sendable { // MARK: - Result Types /// Result of rescoring a word - public struct RescoringResult: Sendable { + public struct RescoringResult: Codable, Sendable { public let originalWord: String public let originalScore: Float public let replacementWord: String? diff --git a/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/SlidingWindowAsrManager.swift b/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/SlidingWindowAsrManager.swift index ead8d833c..7acfe0215 100644 --- a/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/SlidingWindowAsrManager.swift +++ b/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/SlidingWindowAsrManager.swift @@ -564,7 +564,8 @@ public actor SlidingWindowAsrManager { displayResult = interim.withRescoring( text: rescored?.text ?? interim.text, detected: rescored?.detectedTerms ?? [], - applied: applied.isEmpty ? nil : applied + applied: applied.isEmpty ? nil : applied, + replacements: appliedReplacements.isEmpty ? nil : appliedReplacements ) } @@ -580,7 +581,8 @@ public actor SlidingWindowAsrManager { tokenIds: tokens, tokenTimings: displayResult.tokenTimings ?? [], ctcDetectedTerms: displayResult.ctcDetectedTerms, - ctcAppliedTerms: displayResult.ctcAppliedTerms + ctcAppliedTerms: displayResult.ctcAppliedTerms, + ctcReplacements: displayResult.ctcReplacements ) updateContinuation?.yield(update) @@ -959,6 +961,11 @@ public struct SlidingWindowTranscriptionUpdate: Sendable { public let ctcDetectedTerms: [String]? /// Vocabulary terms applied as replacements in this window's text. public let ctcAppliedTerms: [String]? + /// The rescoring decision behind each applied replacement, aligned 1:1 with + /// `ctcAppliedTerms`. Carries the decoded word each term displaced and the + /// scores the decision was made on, so callers can tell a recognizer error + /// from a bad vocabulary replacement. `nil` when nothing was replaced. + public let ctcReplacements: [VocabularyRescorer.RescoringResult]? public init( text: String, @@ -968,7 +975,8 @@ public struct SlidingWindowTranscriptionUpdate: Sendable { tokenIds: [Int] = [], tokenTimings: [TokenTiming] = [], ctcDetectedTerms: [String]? = nil, - ctcAppliedTerms: [String]? = nil + ctcAppliedTerms: [String]? = nil, + ctcReplacements: [VocabularyRescorer.RescoringResult]? = nil ) { self.text = text self.isConfirmed = isConfirmed @@ -978,5 +986,6 @@ public struct SlidingWindowTranscriptionUpdate: Sendable { self.tokenTimings = tokenTimings self.ctcDetectedTerms = ctcDetectedTerms self.ctcAppliedTerms = ctcAppliedTerms + self.ctcReplacements = ctcReplacements } } diff --git a/Sources/FluidAudio/ASR/Parakeet/Unified/StreamingUnifiedAsrManager.swift b/Sources/FluidAudio/ASR/Parakeet/Unified/StreamingUnifiedAsrManager.swift index c29ad4404..1612f0da6 100644 --- a/Sources/FluidAudio/ASR/Parakeet/Unified/StreamingUnifiedAsrManager.swift +++ b/Sources/FluidAudio/ASR/Parakeet/Unified/StreamingUnifiedAsrManager.swift @@ -84,6 +84,8 @@ public actor StreamingUnifiedAsrManager { /// Audio kept before the first pending token when a segment is released, /// so the next CTC pass sees the word's onset. private let vocabPreRollSeconds: Double = 1.0 + // Replacement decisions from the segments rescored since the last drain. + private var pendingVocabReplacements: [VocabularyRescorer.RescoringResult] = [] public private(set) var mlConfiguration: MLModelConfiguration @@ -288,6 +290,17 @@ public actor StreamingUnifiedAsrManager { buildWordTimings(from: consumeTokenTimings()) } + /// Returns the vocabulary replacements applied since the previous call and + /// clears them, draining the same way `consumeTokenTimings()` does so the + /// buffer stays bounded over long streams. Each decision carries the + /// decoded word the term displaced and the scores behind the replacement, + /// which `finish()`'s transcript cannot express. Empty when boosting is not + /// configured, or when no segment has been rescored since the last call. + public func consumeVocabularyReplacements() -> [VocabularyRescorer.RescoringResult] { + defer { pendingVocabReplacements.removeAll(keepingCapacity: true) } + return pendingVocabReplacements + } + public func reset() async throws { samples.removeAll() samplesGlobalStart = 0 @@ -300,6 +313,7 @@ public actor StreamingUnifiedAsrManager { vocabTimings.removeAll() vocabAudio.removeAll() vocabAudioGlobalStart = 0 + pendingVocabReplacements.removeAll() try rnntDecoder?.reset() } @@ -478,6 +492,7 @@ public actor StreamingUnifiedAsrManager { // dropping the segment's leading separator — restore it so the // released text still butts cleanly against the rescored prefix. releasedText = (segmentText.hasPrefix(" ") ? " " : "") + rescored.text + pendingVocabReplacements.append(contentsOf: rescored.replacements.filter { $0.shouldReplace }) } rescoredTranscript += releasedText vocabTimings.removeFirst(cut) diff --git a/Sources/FluidAudio/ASR/Parakeet/Unified/UnifiedAsrManager.swift b/Sources/FluidAudio/ASR/Parakeet/Unified/UnifiedAsrManager.swift index bf7613ffb..6e0fc1031 100644 --- a/Sources/FluidAudio/ASR/Parakeet/Unified/UnifiedAsrManager.swift +++ b/Sources/FluidAudio/ASR/Parakeet/Unified/UnifiedAsrManager.swift @@ -233,7 +233,58 @@ public actor UnifiedAsrManager { guard let tokenizer = tokenizer else { throw ASRError.notInitialized } let merged = try await decodedTokens(samples, tokenizer: tokenizer) let text = tokenizer.decode(ids: merged.map(\.token)) - return await rescoreIfConfigured(text: text, merged: merged, samples: samples) + return await rescoreIfConfigured(text: text, merged: merged, samples: samples)?.text ?? text + } + + /// Transcribe as `transcribe(_:)` does, returning the full ``ASRResult`` + /// instead of only its text. + /// + /// This is the only batch entry point that reports vocabulary boosting: + /// `transcribe(_:)` returns a `String`, so the CTC metadata behind a + /// rescored transcript — which terms the spotter detected, which were + /// applied, and the decision behind each one — has nowhere to go. Callers + /// that need to explain a replacement to a user, or to tell a recognizer + /// error from a bad replacement, use this instead. + public func transcribeDetailed(_ samples: [Float]) async throws -> ASRResult { + guard let tokenizer = tokenizer else { throw ASRError.notInitialized } + let startTime = Date() + let merged = try await decodedTokens(samples, tokenizer: tokenizer) + let text = tokenizer.decode(ids: merged.map(\.token)) + let duration = Double(samples.count) / Double(config.sampleRate) + // Rescored text can replace words, so token timings no longer decode + // to the text verbatim; they remain the raw emissions, as in + // `transcribeWithTimings(_:)`. + let timings = Self.tokenTimings( + from: merged, + secondsPerFrame: Double(config.frameSamples) / Double(config.sampleRate), + vocabulary: tokenizer.vocabulary, + clipDuration: duration + ) + let rescored = await rescoreIfConfigured(text: text, tokenTimings: timings, samples: samples) + let result = ASRResult( + text: text, + confidence: Self.meanConfidence( + of: merged, isEmpty: text.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty), + duration: duration, + processingTime: Date().timeIntervalSince(startTime), + tokenTimings: timings + ) + guard let rescored else { return result } + return Self.applying(rescored, to: result) + } + + /// Fold a rescorer output into a decoded result, reporting the applied + /// terms alongside the decision behind each one. Pure, so the mapping is + /// testable without loading a 600M parameter model. + static func applying(_ rescored: VocabularyRescorer.RescoreOutput, to result: ASRResult) -> ASRResult { + let applied = rescored.replacements.filter { $0.shouldReplace } + let appliedTerms = applied.compactMap { $0.replacementWord } + return result.withRescoring( + text: rescored.text, + detected: rescored.detectedTerms, + applied: appliedTerms.isEmpty ? nil : appliedTerms, + replacements: applied.isEmpty ? nil : applied + ) } /// Transcribe as `transcribe(_:)` does, additionally reporting the encoder @@ -278,19 +329,42 @@ public actor UnifiedAsrManager { } /// Apply vocabulary rescoring to a finished transcript when boosting is - /// configured; otherwise return the transcript unchanged. + /// configured; otherwise return `nil`. + /// + /// Returns the rescorer's whole output, not just its text: the replacement + /// decisions behind it are what `transcribeDetailed(_:)` reports. Building + /// the token timings stays behind the boosting check, so callers that never + /// configured boosting pay nothing for them. private func rescoreIfConfigured( text: String, merged: [ChunkProcessor.TokenWindow], samples: [Float] - ) async -> String { - guard let boosting = vocabularyBoosting, let tokenizer = tokenizer else { return text } + ) async -> VocabularyRescorer.RescoreOutput? { + guard vocabularyBoosting != nil, let tokenizer = tokenizer else { return nil } let timings = Self.tokenTimings( from: merged, secondsPerFrame: Double(config.frameSamples) / Double(config.sampleRate), vocabulary: tokenizer.vocabulary, clipDuration: Double(samples.count) / Double(config.sampleRate) ) - let rescored = await boosting.rescore(text: text, tokenTimings: timings, audioSamples: samples) - return rescored?.text ?? text + return await rescoreIfConfigured(text: text, tokenTimings: timings, samples: samples) + } + + /// As above, for callers that already built the token timings. + private func rescoreIfConfigured( + text: String, tokenTimings: [TokenTiming], samples: [Float] + ) async -> VocabularyRescorer.RescoreOutput? { + guard let boosting = vocabularyBoosting else { return nil } + return await boosting.rescore(text: text, tokenTimings: tokenTimings, audioSamples: samples) + } + + /// Mean token confidence for a decoded window, matching `AsrManager`'s + /// rule: an empty transcript scores 0.1, otherwise the mean softmax + /// probability clamped to [0.1, 1.0]. Pure, so it is testable without + /// loading a 600M parameter model. + static func meanConfidence(of emissions: [ChunkProcessor.TokenWindow], isEmpty: Bool) -> Float { + if isEmpty { return 0.1 } + guard !emissions.isEmpty else { return 0.5 } + let mean = emissions.reduce(Float(0)) { $0 + $1.confidence } / Float(emissions.count) + return max(0.1, min(1.0, mean)) } /// Emission frames → seconds. Pure, so the back-fill rule can be tested diff --git a/Sources/FluidAudio/ITN/TextNormalizer.swift b/Sources/FluidAudio/ITN/TextNormalizer.swift index 50868b19e..08b733133 100644 --- a/Sources/FluidAudio/ITN/TextNormalizer.swift +++ b/Sources/FluidAudio/ITN/TextNormalizer.swift @@ -165,7 +165,8 @@ public final class TextNormalizer: Sendable { processingTime: result.processingTime, tokenTimings: result.tokenTimings, ctcDetectedTerms: result.ctcDetectedTerms, - ctcAppliedTerms: result.ctcAppliedTerms + ctcAppliedTerms: result.ctcAppliedTerms, + ctcReplacements: result.ctcReplacements ) } diff --git a/Tests/FluidAudioTests/ASR/Parakeet/VocabularyRescoringSurfacingTests.swift b/Tests/FluidAudioTests/ASR/Parakeet/VocabularyRescoringSurfacingTests.swift new file mode 100644 index 000000000..3e550a3d3 --- /dev/null +++ b/Tests/FluidAudioTests/ASR/Parakeet/VocabularyRescoringSurfacingTests.swift @@ -0,0 +1,157 @@ +import XCTest + +@testable import FluidAudio + +/// `ctcAppliedTerms` says which vocabulary terms went in; `ctcReplacements` +/// says which decoded word each one displaced and on what scores. These tests +/// pin that the decisions survive to the public result on the batch +/// (`ASRResult`, `UnifiedAsrManager.transcribeDetailed`) and streaming +/// (`SlidingWindowTranscriptionUpdate`, +/// `StreamingUnifiedAsrManager.consumeVocabularyReplacements`) paths. +final class VocabularyRescoringSurfacingTests: XCTestCase { + + private typealias Decision = VocabularyRescorer.RescoringResult + + private let applied = Decision( + originalWord: "codecs", originalScore: 0.21, replacementWord: "Codex", replacementScore: 0.88, + shouldReplace: true, reason: "ctc-score") + + private let declined = Decision( + originalWord: "favor", originalScore: 0.61, replacementWord: "flavor", replacementScore: 0.30, + shouldReplace: false, reason: "below-floor") + + private func baseResult(text: String) -> ASRResult { + ASRResult(text: text, confidence: 0.9, duration: 2.0, processingTime: 0.1) + } + + // MARK: - Batch path + + func testWithRescoringCarriesReplacementDecisions() { + let rescored = baseResult(text: "validate with codecs").withRescoring( + text: "validate with Codex", detected: ["Codex"], applied: ["Codex"], replacements: [applied]) + + XCTAssertEqual(rescored.text, "validate with Codex") + XCTAssertEqual(rescored.ctcDetectedTerms, ["Codex"]) + XCTAssertEqual(rescored.ctcAppliedTerms, ["Codex"]) + XCTAssertEqual(rescored.ctcReplacements?.count, 1) + XCTAssertEqual(rescored.ctcReplacements?.first?.originalWord, "codecs") + XCTAssertEqual(rescored.ctcReplacements?.first?.replacementWord, "Codex") + XCTAssertEqual(rescored.ctcReplacements?.first?.originalScore ?? 0, 0.21, accuracy: 1e-6) + XCTAssertEqual(rescored.ctcReplacements?.first?.replacementScore ?? 0, 0.88, accuracy: 1e-6) + XCTAssertTrue(rescored.ctcReplacements?.first?.shouldReplace ?? false) + XCTAssertEqual(rescored.ctcReplacements?.first?.reason, "ctc-score") + } + + /// The new parameter is defaulted, so the pre-existing three-argument call + /// still compiles and leaves the decisions unset. + func testWithRescoringWithoutReplacementsLeavesDecisionsNil() { + let rescored = baseResult(text: "raw").withRescoring(text: "raw", detected: [], applied: nil) + XCTAssertNil(rescored.ctcReplacements) + XCTAssertEqual(rescored.ctcDetectedTerms, []) + } + + /// `ASRResult` is `Codable`, so the decisions have to survive a round trip. + func testReplacementDecisionsSurviveCodableRoundTrip() throws { + let rescored = baseResult(text: "validate with Codex").withRescoring( + text: "validate with Codex", detected: ["Codex"], applied: ["Codex"], replacements: [applied]) + + let decoded = try JSONDecoder().decode(ASRResult.self, from: JSONEncoder().encode(rescored)) + + XCTAssertEqual(decoded.ctcAppliedTerms, ["Codex"]) + XCTAssertEqual(decoded.ctcReplacements?.first?.originalWord, "codecs") + XCTAssertEqual(decoded.ctcReplacements?.first?.replacementWord, "Codex") + XCTAssertEqual(decoded.ctcReplacements?.first?.reason, "ctc-score") + } + + // MARK: - Streaming path + + func testStreamingUpdateCarriesReplacementDecisions() { + let rescored = baseResult(text: "validate with codecs").withRescoring( + text: "validate with Codex", detected: ["Codex"], applied: ["Codex"], replacements: [applied]) + + let update = SlidingWindowTranscriptionUpdate( + text: rescored.text, + isConfirmed: true, + confidence: rescored.confidence, + timestamp: Date(), + ctcDetectedTerms: rescored.ctcDetectedTerms, + ctcAppliedTerms: rescored.ctcAppliedTerms, + ctcReplacements: rescored.ctcReplacements + ) + + XCTAssertEqual(update.ctcAppliedTerms, ["Codex"]) + XCTAssertEqual(update.ctcReplacements?.count, 1) + XCTAssertEqual(update.ctcReplacements?.first?.originalWord, "codecs") + XCTAssertEqual(update.ctcReplacements?.first?.replacementScore ?? 0, 0.88, accuracy: 1e-6) + } + + func testStreamingUpdateWithoutBoostingLeavesDecisionsNil() { + let update = SlidingWindowTranscriptionUpdate( + text: "validate with codecs", isConfirmed: false, confidence: 0.5, timestamp: Date()) + XCTAssertNil(update.ctcReplacements) + XCTAssertNil(update.ctcAppliedTerms) + } + + // MARK: - UnifiedAsrManager.transcribeDetailed mapping + + /// `transcribeDetailed(_:)` folds the rescorer output in through + /// `UnifiedAsrManager.applying(_:to:)`; only accepted decisions are + /// reported, and they stay aligned with `ctcAppliedTerms`. + func testUnifiedApplyingReportsAcceptedDecisionsOnly() { + let output = VocabularyRescorer.RescoreOutput( + text: "validate with Codex, in your favor", + replacements: [applied, declined], + wasModified: true, + detectedTerms: ["Codex", "flavor"]) + + let result = UnifiedAsrManager.applying(output, to: baseResult(text: "validate with codecs, in your favor")) + + XCTAssertEqual(result.text, "validate with Codex, in your favor") + XCTAssertEqual(result.ctcDetectedTerms, ["Codex", "flavor"]) + XCTAssertEqual(result.ctcAppliedTerms, ["Codex"]) + XCTAssertEqual(result.ctcReplacements?.count, 1) + XCTAssertEqual(result.ctcReplacements?.first?.originalWord, "codecs") + XCTAssertEqual(result.ctcReplacements?.first?.reason, "ctc-score") + } + + /// Rescoring ran but replaced nothing: the detections are still reported, + /// and the decision list is `nil` rather than empty, as `ctcAppliedTerms` is. + func testUnifiedApplyingWithNoAcceptedDecisions() { + let output = VocabularyRescorer.RescoreOutput( + text: "in your favor", replacements: [declined], wasModified: false, detectedTerms: ["flavor"]) + + let result = UnifiedAsrManager.applying(output, to: baseResult(text: "in your favor")) + + XCTAssertEqual(result.ctcDetectedTerms, ["flavor"]) + XCTAssertNil(result.ctcAppliedTerms) + XCTAssertNil(result.ctcReplacements) + } + + /// The result `transcribeDetailed(_:)` builds carries a confidence, so the + /// mean follows `AsrManager`'s rule rather than a fresh one. + func testMeanConfidenceFollowsAsrManagerRule() { + let emissions: [ChunkProcessor.TokenWindow] = [ + (token: 1, timestamp: 0, confidence: 0.8, duration: 1), + (token: 2, timestamp: 2, confidence: 0.6, duration: 1), + ] + XCTAssertEqual(UnifiedAsrManager.meanConfidence(of: emissions, isEmpty: false), 0.7, accuracy: 1e-6) + XCTAssertEqual(UnifiedAsrManager.meanConfidence(of: emissions, isEmpty: true), 0.1, accuracy: 1e-6) + XCTAssertEqual(UnifiedAsrManager.meanConfidence(of: [], isEmpty: false), 0.5, accuracy: 1e-6) + XCTAssertEqual( + UnifiedAsrManager.meanConfidence( + of: [(token: 1, timestamp: 0, confidence: 0.01, duration: 1)], isEmpty: false), + 0.1, accuracy: 1e-6) + } + + // MARK: - StreamingUnifiedAsrManager drain + + /// Without boosting configured nothing is ever accumulated, so the drain is + /// empty and stays empty. + func testStreamingUnifiedDrainIsEmptyWithoutBoosting() async { + let manager = StreamingUnifiedAsrManager() + let first = await manager.consumeVocabularyReplacements() + let second = await manager.consumeVocabularyReplacements() + XCTAssertTrue(first.isEmpty) + XCTAssertTrue(second.isEmpty) + } +}