From 8b8b419cdf644335035901f27f741c9a587fb6dc Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Wed, 19 Aug 2026 02:29:51 -0400 Subject: [PATCH 1/4] feat(asr/nemotron): decode-time custom vocabulary biasing (no CTC head required) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements issue #841: token-level shallow fusion inside the Nemotron multilingual streaming greedy RNN-T decode, based on the engine contributed by the Alma team, with the changes requested in review: - Scalability: fresh-start candidates (offset 0, decode-state independent) are precomputed once at init; per-step continuation matching walks tail suffixes through a trie over term piece-forms, so per-emission cost is bounded by tail length and live partial matches, not vocabulary size. - Weight safety: CustomVocabularyTerm.weight predates this engine and is a CTC-rescoring scale (the simple text-list loader assigns 10.0). Applied raw as a per-token logit bonus it over-biases, so overrides are clamped to maxBoost=6.0 with a one-time warning; default boost is 4.5 (the measured recall peak, matching the tuned CTC rescorer cbw). - Fused-argmax (B2) handling is explicit, not silent: with a vocabulary active, decode prefers a logits-producing step decoder over decoder_joint_argmax; if B2 is the only step decoder, the vocabulary is rejected with an error log (all-or-nothing, no half-applied biasing). - Every argmax site is biased and traced: legacy per-frame loop (B3/B1/ bare-pair), speculative scan, and all drain branches report boosted flips via FLUIDAUDIO_BIAS_LOG=1. - Match state resets on reset()/finish(); terms survive like the selected language and may be set before models load. CLI: nemotron-multilingual-transcribe gains --custom-vocab (JSON config or one-term-per-line text). Validated on real speech (LibriSpeech test-clean, cached 2240ms B1 assets): 7127-75947-0033 baseline misspells all three rare names (Levallier / Tenerchalte / Bragalone); with the vocabulary, Tonnay Charente and Bragelonne substantially recover, with blank-overtake flips visible in the trace. Neutrality: a 40-term invented-distractor vocabulary leaves a neutral utterance byte-identical with zero flips; on the rare-name clip, distractors sharing the true audio's opening letters can contest the uncertain region — documented as the false-fire profile, with guidance to keep vocabularies small and relevant. Unit tests cover word-start anchoring (#702 over-fire mode), CJK unanchored 2-cluster terms, Devanagari conjunct rejoin in scalar space, multi-offset overlap, alias/weight clamp semantics, special-token immunity, reset lifecycle, biased selection incl. blank overtake, and a 2000-term scaling smoke. --- Documentation/ASR/NemotronMultilingual.md | 45 ++ .../Nemotron/NemotronVocabularyBias.swift | 466 ++++++++++++++++++ ...emotronMultilingualAsrManager+Decode.swift | 38 +- ...otronMultilingualAsrManager+Pipeline.swift | 13 +- ...eamingNemotronMultilingualAsrManager.swift | 10 + .../NemotronMultilingualTranscribe.swift | 22 + .../NemotronVocabularyBiasTests.swift | 259 ++++++++++ 7 files changed, 844 insertions(+), 9 deletions(-) create mode 100644 Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/NemotronVocabularyBias.swift create mode 100644 Tests/FluidAudioTests/ASR/Parakeet/Streaming/NemotronVocabularyBiasTests.swift diff --git a/Documentation/ASR/NemotronMultilingual.md b/Documentation/ASR/NemotronMultilingual.md index 832a93796..087d570e1 100644 --- a/Documentation/ASR/NemotronMultilingual.md +++ b/Documentation/ASR/NemotronMultilingual.md @@ -83,6 +83,51 @@ let detected = await manager.detectedLanguage() // e.g. "fr-FR" await manager.reset() ``` +### Custom vocabulary (hotword biasing) + +The Nemotron exports ship no CTC head, so the sliding-window path's CTC +rescorer does not apply here. Instead the manager biases inside the greedy +RNN-T decode: tokens that extend a partially matched vocabulary term get a +flat log-prob bonus at every emission step (decode-time shallow fusion, issue +#841). + +```swift +await manager.setCustomVocabulary([ + CustomVocabularyTerm(text: "Torvane"), + CustomVocabularyTerm(text: "Quexal", aliases: ["Kwexal"]), +]) +// Pass [] to disable. Survives reset(); may be set before loadModels(). +``` + +Semantics and limits: + +- **Weight scale.** `weight` here is a *per-token logit bonus* (default 4.5, + the measured recall peak), not the CTC rescoring weight. Values above 6.0 + are clamped — the simple text-list loader's `weight: 10.0` would otherwise + over-bias (measured artifacts: word splits like "build today" → "build to + day"). Most terms should omit `weight`. +- **Assets.** Biasing needs a logits-producing step decoder + (`decoder_joint_noencproj` or `decoder_joint`). When a vocabulary is + active the fused-argmax `decoder_joint_argmax` asset is bypassed in favor + of a logits path; if it is the *only* step decoder, the vocabulary is + rejected with an error log rather than silently half-applied. +- **Greedy decode.** Once a boosted token wins the argmax it is committed — + there is no beam to undo an over-fire. Keep vocabularies to genuinely + rare terms; a term the model hears as a word it already spells + ("Pheynix" → "Phoenix") needs an alias, not more weight. +- **Chunk boundaries.** A term whose audio spans a chunk boundary is decoded + as two independent halves; single-word terms are the reliable target. +- Terms shorter than 3 letters are skipped (2 for CJK); word-start anchoring + keeps "ran" from matching into "CRAN". CJK terms match unanchored. +- **False-fire profile** (LibriSpeech test-clean spot check, 40 invented + distractor terms): confidently decoded speech is untouched, but an + acoustically uncertain rare-name region can be captured by a distractor + that shares the true audio's opening letters. Keep vocabularies small and + relevant; do not feed speculative or screen-harvested term lists. + +`FLUIDAUDIO_BIAS_LOG=1` traces every boosted flip on stderr for weight +tuning and over-fire attribution. + ## Benchmark Results Apple M2, FLEURS test set, int8 encoder, `MLComputeUnits.cpuAndNeuralEngine`. diff --git a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/NemotronVocabularyBias.swift b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/NemotronVocabularyBias.swift new file mode 100644 index 000000000..61df4be78 --- /dev/null +++ b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/NemotronVocabularyBias.swift @@ -0,0 +1,466 @@ +import CoreML +import Foundation + +/// Decode-time shallow-fusion hotword biasing for the Nemotron streaming +/// RNN-T greedy decode loop. +/// +/// The sliding-window (Parakeet TDT) path boosts custom vocabulary by +/// rescoring decoded text against CTC log-probabilities +/// (`SlidingWindow/CustomVocabulary/`). The Nemotron multilingual export +/// ships no CTC head, so this biases inside the greedy decode instead: at +/// every emission step the tokens that would *extend a partially matched +/// vocabulary term* receive a flat log-prob bonus, and the boosted best is +/// emitted when it beats the unbiased argmax. Adding a constant before the +/// argmax is equivalent to boosting the token's log-softmax probability by +/// the same amount (the normalizer is shared), so the bonus lives on the +/// same "context-biasing weight" scale as the CTC path's — see +/// `defaultBoost`. +/// +/// Matching is on *piece text*, not on one fixed token segmentation: +/// whichever segmentation of the term the decoder is drifting toward stays +/// boostable, because every vocab piece that continues the term from a +/// viable match offset is a candidate. Three details carry the +/// multilingual weight: +/// +/// - Comparison runs in NFC **Unicode-scalar space**, because SentencePiece +/// pieces split grapheme clusters (a Devanagari conjunct arrives as +/// `स्` + `ते`, an Arabic diacritic as a lone mark) and a cluster-level +/// tail could never re-join them. +/// - **Every viable match offset** is tracked, not just the longest — +/// a term that repeats its leading word ("ab ab c" hearing "ab ab ab c") +/// stays boostable through the shorter overlap. +/// - Space-delimited scripts anchor at word starts (the term's piece form +/// leads with `▁`, and `▁` appears only at word boundaries), which is +/// what keeps "ran" from matching into "CRAN" — the over-fire mode +/// issue #702 documents; the `minTermLength` guard applies here too. +/// CJK terms are indexed *unanchored* instead (the shipped vocab marks +/// almost no CJK piece with `▁`) and accept 2-cluster terms (a +/// two-character CJK word is a full word, not a short keyword). +/// +/// Scaling: the fresh-start candidate set (offset 0 for every term) does +/// not depend on decode state, so it is computed once at init. Per step, +/// only terms whose prefix matches a suffix of the emitted tail need work; +/// those are found through a trie over term piece-forms, so the per-step +/// cost is bounded by the tail length and the number of *live* partial +/// matches — not by the vocabulary size. +/// +/// Greedy-decoding caveat: this biases a greedy argmax, so once a boosted +/// token wins it is committed — there is no beam to recover an over-fire. +/// Streaming caveat (same as the CTC path documents): a term whose audio +/// spans a chunk boundary is decoded as two independent halves, so +/// single-word terms are the reliable target. +final class NemotronVocabularyBias { + + /// One boostable token continuation at the current match state. + struct Candidate { + let tokenId: Int + let boost: Float + } + + private struct Entry { + /// NFC scalars of the term's piece form: `"▁steve▁jobs"`, or the + /// bare characters for an unanchored (CJK-initial) term. + let pieceForm: [Unicode.Scalar] + let boost: Float + } + + /// Trie over entry piece-forms. Each node at depth `d` records every + /// entry whose form starts with the node's path, so one walk of a tail + /// suffix yields all entries continuable from offset `d`. + private final class TrieNode { + var children: [UInt32: TrieNode] = [:] + var entryIndices: [Int] = [] + } + + /// Terms shorter than this (letters, not counting whitespace) are + /// skipped — the CTC rescorer's own guard against short-keyword + /// over-firing. CJK surfaces accept two clusters (see `isUsable`). + static let minTermLength = 3 + + /// Per-token log-prob bonus for terms without their own weight. + /// Measured on the issue #841 rig (8 rare-term clips + 4 phonetic + /// neighbours + 2 neutrals): 3.0 recalled 2/8, 4.5 recalled 5/8, 6.0 + /// fell back to 4/8 with over-boost artifacts ("build today" splitting + /// into "build to day"); every weight left neighbours and neutrals + /// untouched. 4.5 is also the cbw the tuned CTC rescorer settled on + /// (`ContextBiasingConstants.rescorerConfig`). + static let defaultBoost: Float = 4.5 + + /// Ceiling for per-term `weight` overrides. `CustomVocabularyTerm.weight` + /// predates this engine and is set to 10.0 by the simple text-list + /// loader — a CTC-rescoring scale, not a per-token logit bonus. Applied + /// raw it would over-bias badly (degradation is measurable from 6.0), + /// so overrides are clamped here rather than reinterpreted. + static let maxBoost: Float = 6.0 + + private let entries: [Entry] + private let trieRoot: TrieNode + /// Offset-0 candidates for every term — decode-state independent, so + /// computed once. `candidates()` starts from a copy of this map. + private let freshStartCandidates: [Int: Float] + /// Lowercased NFC piece text → every token id whose piece folds to it + /// (`"▁the"` → ids of `"▁The"`, `"▁the"`, …). + private let idsByPiece: [String: [Int]] + /// Lowercased NFC piece per id, for tracking emissions. Special pieces + /// (``, lang tags) are absent — they never advance a match. + private let pieceById: [Int: String] + /// Lowercased NFC scalars of the most recent emissions, long enough to + /// hold any term's partial match. + private var tail: [Unicode.Scalar] = [] + private let tailCap: Int + private var cached: [Candidate]? + + /// Builds the biasing state, or `nil` when no usable term remains. + /// + /// - Parameters: + /// - terms: vocabulary; `weight` is a per-term flat log-prob bonus + /// override (clamped to `maxBoost`), `aliases` are additional + /// surface forms boosted (and emitted) as themselves. + /// - pieces: the model's id → piece table (`▁` markers intact, case + /// preserved). + /// - defaultBoost: bonus for terms without a `weight`. + init?( + terms: [CustomVocabularyTerm], + pieces: [Int: String], + defaultBoost: Float = NemotronVocabularyBias.defaultBoost + ) { + var entries: [Entry] = [] + var maxLen = 0 + for term in terms { + let boost = Self.effectiveBoost(of: term, defaultBoost: defaultBoost) + for surface in Self.usableSurfaces(of: term, defaultBoost: defaultBoost) { + let form = Self.pieceForm(surface) + entries.append(Entry(pieceForm: form, boost: boost)) + maxLen = max(maxLen, form.count) + } + } + guard !entries.isEmpty else { return nil } + self.entries = entries + self.tailCap = maxLen + + var idsByPiece: [String: [Int]] = [:] + var pieceById: [Int: String] = [:] + for (id, piece) in pieces { + // Special tokens (``, ``, …) are not text the + // decoder spells words with; letting one into the index would + // boost it, and letting one into the tail would corrupt the + // match position. + if piece.hasPrefix("<") && piece.hasSuffix(">") { continue } + let folded = piece.lowercased().precomposedStringWithCanonicalMapping + idsByPiece[folded, default: []].append(id) + pieceById[id] = folded + } + self.idsByPiece = idsByPiece + self.pieceById = pieceById + + let root = TrieNode() + for (index, entry) in entries.enumerated() { + var node = root + for scalar in entry.pieceForm { + let next = + node.children[scalar.value] + ?? { + let created = TrieNode() + node.children[scalar.value] = created + return created + }() + next.entryIndices.append(index) + node = next + } + } + self.trieRoot = root + + var fresh: [Int: Float] = [:] + for entry in entries { + Self.accumulate(from: entry, offset: 0, idsByPiece: idsByPiece, into: &fresh) + } + self.freshStartCandidates = fresh + } + + /// Whether a term would survive this engine's hygiene — the one copy of + /// the rule callers use to keep their "applied" accounting honest. + static func isUsable(_ term: CustomVocabularyTerm) -> Bool { + !usableSurfaces(of: term, defaultBoost: defaultBoost).isEmpty + } + + /// The per-token bonus a term will actually receive: its `weight` when + /// set (clamped into `(0, maxBoost]`), the default otherwise. + static func effectiveBoost(of term: CustomVocabularyTerm, defaultBoost: Float = defaultBoost) -> Float { + min(term.weight ?? defaultBoost, maxBoost) + } + + /// The surfaces of `term` (text + aliases) that pass the length and + /// weight guards. + private static func usableSurfaces( + of term: CustomVocabularyTerm, defaultBoost: Float + ) -> [String] { + guard effectiveBoost(of: term, defaultBoost: defaultBoost) > 0 else { return [] } + return ([term.text] + (term.aliases ?? [])).filter { surface in + let letters = surface.filter { !$0.isWhitespace } + let minimum = surface.unicodeScalars.contains(where: isCJK) ? 2 : minTermLength + return letters.count >= minimum + } + } + + /// The SentencePiece word-boundary marker. + private static let marker: Unicode.Scalar = "\u{2581}" + + /// Han, kana and Hangul — scripts the vocab writes with unmarked + /// pieces, where `▁`-anchoring would make every term inert. + private static func isCJK(_ scalar: Unicode.Scalar) -> Bool { + switch scalar.value { + case 0x3040...0x30FF, // hiragana + katakana + 0x3400...0x4DBF, 0x4E00...0x9FFF, // Han + 0xAC00...0xD7AF, 0xF900...0xFAFF: // Hangul syllables, compat Han + return true + default: + return false + } + } + + /// `"Steve Jobs"` → the NFC scalars of `"▁steve▁jobs"`; a CJK-initial + /// term stays unanchored (no markers at all), because the vocab spells + /// CJK with unmarked pieces. + private static func pieceForm(_ text: String) -> [Unicode.Scalar] { + let folded = text.lowercased().precomposedStringWithCanonicalMapping + let anchored = !(folded.unicodeScalars.first.map(isCJK) ?? false) + var form: [Unicode.Scalar] = [] + var atBoundary = true + for ch in folded { + if ch.isWhitespace { + atBoundary = true + continue + } + if atBoundary { + if anchored { form.append(marker) } + atBoundary = false + } + form.append(contentsOf: ch.unicodeScalars) + } + return form + } + + /// Record an emitted token so the match state follows the decode. + /// Call for every committed non-blank emission, in order. + func observe(_ tokenId: Int) { + guard let piece = pieceById[tokenId] else { return } + tail.append(contentsOf: piece.unicodeScalars) + if tail.count > tailCap { + tail.removeFirst(tail.count - tailCap) + } + cached = nil + } + + /// Forget the match state (stream reset / finish). The vocabulary + /// itself survives, like the selected language does. + func resetMatchState() { + tail.removeAll(keepingCapacity: true) + cached = nil + } + + /// The boostable continuations at the current match state, one entry + /// per token id with the strongest boost that reaches it. Cached until + /// the next `observe`/`resetMatchState`. + /// + /// For each term, every viable offset is live: the fresh start (a new + /// word can begin at any time, precomputed) plus every strict prefix of + /// the term the emitted tail currently ends with (found by walking each + /// tail suffix through the piece-form trie). The candidates are every + /// vocab piece equal to a prefix of the term's remaining text from any + /// of those offsets. + func candidates() -> [Candidate] { + if let cached { return cached } + var best = freshStartCandidates + if !tail.isEmpty { + let maxOffset = min(tail.count, tailCap) + for start in (tail.count - maxOffset).. 0) may still narrow an open match + // one letter at a time, and an unanchored CJK term's single + // characters stay (each is a selective full syllable). + if offset == 0, form.first == marker, letters < 2 { continue } + for id in ids where entry.boost > best[id, default: -.infinity] { + best[id] = entry.boost + } + } + } +} + +/// `FLUIDAUDIO_BIAS_LOG=1` traces every boosted flip on stderr — which +/// piece the plain argmax wanted, which the vocabulary promoted, and both +/// logits. Off by default; diagnostic rigs read it to attribute +/// split/insert artifacts to the exact flip that caused them. +let nemotronBiasLogEnabled: Bool = { + let value = ProcessInfo.processInfo.environment["FLUIDAUDIO_BIAS_LOG"] ?? "" + return !(value.isEmpty || value == "0" || value.lowercased() == "false") +}() + +extension StreamingNemotronMultilingualAsrManager { + + /// Configure decode-time hotword biasing (the Nemotron counterpart of + /// the sliding-window path's vocabulary boosting; no CTC models + /// needed). Pass an empty list to disable. Survives `reset()`; takes + /// effect from the next emission. May be called before models load — + /// the vocabulary is (re)bound whenever a tokenizer becomes available. + /// + /// `weight` is a per-token log-prob bonus clamped to + /// `NemotronVocabularyBias.maxBoost`; most terms should omit it. + /// Biasing requires a logits-producing step decoder: when only the + /// fused-argmax (B2) asset is loaded, the vocabulary is rejected with + /// an error log rather than silently half-applied. + public func setCustomVocabulary(_ terms: [CustomVocabularyTerm]) { + vocabularyTerms = terms + rebuildVocabularyBias() + } + + /// Bind the stored terms to the loaded tokenizer. Called from the model + /// load path and from `setCustomVocabulary`. + internal func rebuildVocabularyBias() { + guard !vocabularyTerms.isEmpty, let tokenizer else { + vocabularyBias = nil + return + } + let anyStepDecoder = + decoderJointNoEncProj != nil || decoderJointArgmax != nil || decoderJoint != nil + || (decoder != nil && joint != nil) + let logitsStepDecoder = + decoderJointNoEncProj != nil || decoderJoint != nil || (decoder != nil && joint != nil) + if anyStepDecoder && !logitsStepDecoder { + // B2-only asset set: the fused-argmax model never exposes + // logits, so no decode site can be biased. All-or-nothing — + // a silently unbiased decode reads exactly like a weak boost. + vocabularyBias = nil + logger.error( + "Custom vocabulary requires a logits-producing step decoder, but only the fused-argmax " + + "(decoder_joint_argmax) asset is loaded. Vocabulary biasing is DISABLED. Ship the " + + "decoder_joint_noencproj or decoder_joint asset to enable it.") + return + } + let clamped = vocabularyTerms.filter { ($0.weight ?? 0) > NemotronVocabularyBias.maxBoost } + if !clamped.isEmpty { + logger.warning( + "Custom vocabulary: \(clamped.count) term(s) with weight > " + + "\(NemotronVocabularyBias.maxBoost) clamped (weight is a per-token logit bonus " + + "here, not a CTC rescoring weight)") + } + // The multilingual wrapper hides the base vocabulary map, but ids + // are contiguous and `rawToken(for:)` is exact — rebuild the table. + var pieces: [Int: String] = [:] + pieces.reserveCapacity(config.vocabSize) + for id in 0.. Int { + let plain = findMaxIndex(logits) + guard let bias = vocabularyBias else { return plain } + let candidates = bias.candidates() + guard !candidates.isEmpty else { return plain } + let count = logits.count + let ptr = logits.dataPointer.bindMemory(to: Float.self, capacity: count) + let picked = Self.pickBiased(plain: plain, candidates: candidates, count: count) { ptr[$0] } + if picked != plain { + traceBiasFlip(plain: plain, picked: picked, plainLogit: ptr[plain], pickedLogit: ptr[picked]) + } + return picked + } + + /// The comparison itself, layout-agnostic: `logit` reads one vocab + /// index. Shared by `selectToken` and the speculative scan (whose + /// logits are strided and possibly Float16). + internal static func pickBiased( + plain: Int, + candidates: [NemotronVocabularyBias.Candidate], + count: Int, + logit: (Int) -> Float + ) -> Int { + var bestId = plain + var bestScore = logit(plain) + for candidate in candidates where candidate.tokenId < count { + let score = logit(candidate.tokenId) + candidate.boost + if score > bestScore { + bestScore = score + bestId = candidate.tokenId + } + } + return bestId + } + + /// Emit one `FLUIDAUDIO_BIAS_LOG` trace line for a boosted flip. Every + /// biased site reports through here — the speculative scan included, + /// which is where most flips happen (it is the one site that always + /// has logits); a drains-only trace undercounts and mis-attributes. + internal func traceBiasFlip(plain: Int, picked: Int, plainLogit: Float, pickedLogit: Float) { + guard nemotronBiasLogEnabled else { return } + let plainPiece = + plain == config.blankIdx ? "" : (tokenizer?.rawToken(for: plain) ?? "?") + let pickedPiece = tokenizer?.rawToken(for: picked) ?? "?" + FileHandle.standardError.write( + Data("bias-flip: '\(plainPiece)'(\(plainLogit)) -> '\(pickedPiece)'(\(pickedLogit))\n".utf8)) + } +} diff --git a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager+Decode.swift b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager+Decode.swift index a2d64330a..e5be9ada5 100644 --- a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager+Decode.swift +++ b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager+Decode.swift @@ -220,6 +220,29 @@ extension StreamingNemotronMultilingualAsrManager { } } } + // Hotword shallow fusion: a boosted vocabulary continuation + // may overtake the plain argmax — including blank, which is + // how a term the decoder was about to drop gets emitted. + if let bias = vocabularyBias { + let candidates = bias.candidates() + if !candidates.isEmpty { + let readLogit: (Int) -> Float = { v in + logitsIsF16 + ? nemotronHalfBitsToFloat( + logitsF16Ptr![frameBase + v * logitsStride3]) + : logitsF32Ptr![frameBase + v * logitsStride3] + } + let plain = bestIdx + bestIdx = Self.pickBiased( + plain: plain, candidates: candidates, count: vocabSize, + logit: readLogit) + if bestIdx != plain { + traceBiasFlip( + plain: plain, picked: bestIdx, + plainLogit: readLogit(plain), pickedLogit: readLogit(bestIdx)) + } + } + } if bestIdx != blankIdx { firstNonBlankAt = kk emittedToken = bestIdx @@ -334,10 +357,13 @@ extension StreamingNemotronMultilingualAsrManager { let djneH = djneOutput.featureValue(for: "h_out")?.multiArrayValue, let djneC = djneOutput.featureValue(for: "c_out")?.multiArrayValue else { throw ASRError.processingFailed("Drain B3+B1 failed") } - dBestIdx = findMaxIndex(djneLogits) + dBestIdx = selectToken(djneLogits) newH = djneH newC = djneC - } else if let dja = self.decoderJointArgmax { + } else if let dja = self.decoderJointArgmax, !vocabularyBiasPrefersLogits { + // No logits out of the fused-argmax model, so an active + // vocabulary routes this drain to a logits path instead + // (same as the inner loop's B2 branch). let djaInput = try MLDictionaryFeatureProvider(dictionary: [ "token": MLFeatureValue(multiArray: tokInput2), "token_length": MLFeatureValue(multiArray: tokLen2), @@ -376,7 +402,7 @@ extension StreamingNemotronMultilingualAsrManager { let djH = djOutput.featureValue(for: "h_out")?.multiArrayValue, let djC = djOutput.featureValue(for: "c_out")?.multiArrayValue else { throw ASRError.processingFailed("Drain B1 failed") } - dBestIdx = findMaxIndex(djLogits) + dBestIdx = selectToken(djLogits) newH = djH newC = djC } else { @@ -400,7 +426,7 @@ extension StreamingNemotronMultilingualAsrManager { let jOut = try await self.joint!.prediction(from: jIn) guard let jLogits = jOut.featureValue(for: "logits")?.multiArrayValue else { throw ASRError.processingFailed("Drain joint failed") } - dBestIdx = findMaxIndex(jLogits) + dBestIdx = selectToken(jLogits) newH = h2 newC = c2 } @@ -501,6 +527,10 @@ extension StreamingNemotronMultilingualAsrManager { internal func appendTokenTiming( _ tokenId: Int, frameInChunk: Int, tokenizer: NemotronMultilingualTokenizer ) { + // Every decode path commits an emission through here, which makes it + // the one place the hotword matcher has to watch. Lang tags carry no + // piece text (the bias skips special pieces itself). + vocabularyBias?.observe(tokenId) guard !config.langTagTokenIds.contains(tokenId) else { return } let startTime = Double(absoluteFrameBase + frameInChunk) * ASRConstants.secondsPerEncoderFrame diff --git a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager+Pipeline.swift b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager+Pipeline.swift index 132a98872..0c48e4e04 100644 --- a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager+Pipeline.swift +++ b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager+Pipeline.swift @@ -423,11 +423,14 @@ extension StreamingNemotronMultilingualAsrManager { else { throw ASRError.processingFailed("B3+B1 fused decoder_joint_noencproj failed") } - predToken = findMaxIndex(fl) + predToken = selectToken(fl) hOut = fh cOut = fc - } else if let dja = decoderJointArgmax { - // Triple-fused path: token + h + c + encoder → token_id (int32) + h + c + } else if let dja = decoderJointArgmax, !vocabularyBiasPrefersLogits { + // Triple-fused path: token + h + c + encoder → token_id (int32) + h + c. + // Hotword biasing cannot reach this branch — the argmax is + // fused inside the CoreML model and no logits come out — + // so an active vocabulary routes to a logits path instead. let tripleInput = try MLDictionaryFeatureProvider(dictionary: [ "token": MLFeatureValue(multiArray: tokenInput), "token_length": MLFeatureValue(multiArray: tokenLen), @@ -470,7 +473,7 @@ extension StreamingNemotronMultilingualAsrManager { else { throw ASRError.processingFailed("Fused decoder_joint failed") } - predToken = findMaxIndex(fl) + predToken = selectToken(fl) hOut = fh cOut = fc } else if let decoder = self.decoder, let joint = self.joint { @@ -506,7 +509,7 @@ extension StreamingNemotronMultilingualAsrManager { guard let jl = jointOutput.featureValue(for: "logits")?.multiArrayValue else { throw ASRError.processingFailed("Joint failed") } - predToken = findMaxIndex(jl) + predToken = selectToken(jl) hOut = dh cOut = dc } else { diff --git a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift index 54e3608ea..fda0a1483 100644 --- a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift +++ b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift @@ -89,6 +89,12 @@ public actor StreamingNemotronMultilingualAsrManager { // Accumulated token IDs (raw, including any lang-tag tokens) internal var accumulatedTokenIds: [Int] = [] + // Decode-time hotword biasing (see NemotronVocabularyBias.swift). The + // terms survive reset() like the selected language does; the bias is + // rebuilt whenever a tokenizer becomes available. + internal var vocabularyTerms: [CustomVocabularyTerm] = [] + internal var vocabularyBias: NemotronVocabularyBias? + // Per-token absolute timings captured during the RNNT decode loop, parallel // to the user-visible (lang-tag-stripped) token stream. Each token's // startTime is its absolute encoder-frame index * secondsPerEncoderFrame. @@ -491,6 +497,7 @@ public actor StreamingNemotronMultilingualAsrManager { vocabPath: tokenizerURL, langTagTokenIds: config.langTagTokenIds ) + rebuildVocabularyBias() // Initialize states try resetStates() @@ -781,6 +788,7 @@ public actor StreamingNemotronMultilingualAsrManager { lastFinishTokenTimings.removeAll() audioBufferOffset = 0 firstDetectedLanguage = nil + vocabularyBias?.resetMatchState() do { try resetStates() } catch { @@ -1008,6 +1016,8 @@ public actor StreamingNemotronMultilingualAsrManager { lastFinishTokenTimings = accumulatedTokenTimings accumulatedTokenIds.removeAll() accumulatedTokenTimings.removeAll() + // The emitted-token tail must follow the accumulated ids it mirrors. + vocabularyBias?.resetMatchState() if appendTerminalPunctuation { return Self.tidyTerminalPunctuation( diff --git a/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronMultilingualTranscribe.swift b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronMultilingualTranscribe.swift index ba340f26a..1ba52809a 100644 --- a/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronMultilingualTranscribe.swift +++ b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronMultilingualTranscribe.swift @@ -22,6 +22,9 @@ public class NemotronMultilingualTranscribe { var promptId: Int? /// Chunk-size tier in ms (560 / 1120 / 2240 / 4480) for auto-download. var chunkMs: Int = 2240 + /// Custom vocabulary file (JSON config or one-term-per-line text) for + /// decode-time hotword biasing. + var customVocabPath: String? public init() {} } @@ -70,6 +73,11 @@ public class NemotronMultilingualTranscribe { if i < arguments.count, let ms = Int(arguments[i]) { config.chunkMs = ms } + case "--custom-vocab": + i += 1 + if i < arguments.count { + config.customVocabPath = arguments[i] + } case "--help", "-h": printUsage() return @@ -125,6 +133,8 @@ public class NemotronMultilingualTranscribe { --chunk-ms Chunk-size tier for auto-download: 560 / 1120 / 2240 (default, recommended) / 4480. --prompt-id Raw prompt id (overrides --language) + --custom-vocab Vocabulary file for decode-time hotword biasing + (JSON config or one-term-per-line text) --help, -h Show this help Notes: @@ -189,6 +199,18 @@ public class NemotronMultilingualTranscribe { } else { logger.info("Using default prompt id (auto)") } + + if let vocabPath = config.customVocabPath { + let vocabURL = URL(fileURLWithPath: vocabPath) + let vocab: CustomVocabularyContext + do { + vocab = try CustomVocabularyContext.load(from: vocabURL) + } catch { + vocab = try CustomVocabularyContext.loadFromSimpleFormat(from: vocabURL) + } + await manager.setCustomVocabulary(vocab.terms) + logger.info("Custom vocabulary: \(vocab.terms.count) term(s) from \(vocabPath)") + } logger.info("") for (index, fileURL) in config.inputFiles.enumerated() { diff --git a/Tests/FluidAudioTests/ASR/Parakeet/Streaming/NemotronVocabularyBiasTests.swift b/Tests/FluidAudioTests/ASR/Parakeet/Streaming/NemotronVocabularyBiasTests.swift new file mode 100644 index 000000000..b9285f257 --- /dev/null +++ b/Tests/FluidAudioTests/ASR/Parakeet/Streaming/NemotronVocabularyBiasTests.swift @@ -0,0 +1,259 @@ +import XCTest + +@testable import FluidAudio + +/// Matcher-level tests for the Nemotron decode-time hotword biasing engine, +/// against a toy SentencePiece-style table — no CoreML needed. +final class NemotronVocabularyBiasTests: XCTestCase { + + /// Toy piece table. `▁` marks word starts, case variants fold together, + /// `<...>` pieces are specials the engine must ignore. + private let pieces: [Int: String] = [ + 0: "▁", + 1: "▁to", 2: "r", 3: "vane", 4: "▁tor", 5: "▁torvane", + 6: "▁The", 7: "▁the", + 8: "▁ran", 9: "▁t", 10: "▁cr", 11: "an", + 13: "东", 14: "京", + 15: "▁नम", 16: "स्", 17: "ते", + 18: "", 19: "", + 21: "▁ab", 22: "▁c", 23: "v", 24: "▁vor", + ] + + private func makeBias( + _ terms: [CustomVocabularyTerm], defaultBoost: Float = NemotronVocabularyBias.defaultBoost + ) -> NemotronVocabularyBias? { + NemotronVocabularyBias(terms: terms, pieces: pieces, defaultBoost: defaultBoost) + } + + private func candidateMap(_ bias: NemotronVocabularyBias) -> [Int: Float] { + Dictionary(uniqueKeysWithValues: bias.candidates().map { ($0.tokenId, $0.boost) }) + } + + // MARK: - Term hygiene + + func testInitReturnsNilWithNoUsableTerms() { + XCTAssertNil(makeBias([CustomVocabularyTerm(text: "ab")])) // too short + XCTAssertNil(makeBias([CustomVocabularyTerm(text: "torvane", weight: 0)])) + XCTAssertNil(makeBias([CustomVocabularyTerm(text: "torvane", weight: -1)])) + XCTAssertNil(makeBias([])) + } + + func testIsUsableMatchesEngineGuards() { + XCTAssertTrue(NemotronVocabularyBias.isUsable(CustomVocabularyTerm(text: "torvane"))) + XCTAssertFalse(NemotronVocabularyBias.isUsable(CustomVocabularyTerm(text: "ab"))) + XCTAssertFalse(NemotronVocabularyBias.isUsable(CustomVocabularyTerm(text: "torvane", weight: -1))) + // Two CJK clusters are a full word, not a short keyword. + XCTAssertTrue(NemotronVocabularyBias.isUsable(CustomVocabularyTerm(text: "东京"))) + // A term rescued by a usable alias is usable. + XCTAssertTrue( + NemotronVocabularyBias.isUsable(CustomVocabularyTerm(text: "ab", aliases: ["torvane"]))) + } + + // MARK: - Fresh-start candidates (word-start anchoring) + + func testFreshStartBoostsAnchoredPrefixPieces() throws { + let bias = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "Torvane")])) + let map = candidateMap(bias) + XCTAssertNotNil(map[1], "▁to is a 2-letter anchored prefix") + XCTAssertNotNil(map[4], "▁tor is an anchored prefix") + XCTAssertNotNil(map[5], "▁torvane is the whole term") + XCTAssertNil(map[3], "vane does not start the term at a word boundary") + XCTAssertNil(map[2], "bare r is not an anchored prefix") + } + + func testFreshStartExcludesSingleLetterAndBareMarkerPieces() throws { + let bias = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "torvane")])) + let map = candidateMap(bias) + XCTAssertNil(map[9], "single-letter word-initial piece ▁t must not carry a standing boost") + XCTAssertNil(map[0], "the bare ▁ piece narrows nothing and must never be boosted") + } + + func testCaseVariantsOfAPieceAllFold() throws { + let bias = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "the")])) + let map = candidateMap(bias) + XCTAssertNotNil(map[6], "▁The folds to the term") + XCTAssertNotNil(map[7], "▁the folds to the term") + } + + // MARK: - Continuation matching + + func testContinuationAfterPartialMatch() throws { + let bias = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "torvane")])) + bias.observe(4) // ▁tor + let map = candidateMap(bias) + XCTAssertNotNil(map[3], "vane completes the term from the matched offset") + XCTAssertNotNil(map[23], "a single-letter continuation (v) may narrow an open match") + XCTAssertNotNil(map[5], "fresh start stays live alongside the continuation") + } + + func testWordStartAnchoringBlocksMidWordMatch() throws { + // The #702 over-fire mode: "ran" must not match into "CRAN". + let bias = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "ran")])) + bias.observe(10) // ▁cr + let map = candidateMap(bias) + XCTAssertNil(map[11], "an must not be boosted mid-word after ▁cr") + XCTAssertNotNil(map[8], "the anchored fresh start ▁ran remains the only candidate") + XCTAssertEqual(map.count, 1) + } + + func testMultiOffsetOverlapSurvivesRepeatedLeadingWord() throws { + // "ab ab c" hearing "ab ab ab c": the shorter overlap keeps the + // term boostable. + let bias = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "ab ab c")])) + bias.observe(21) // ▁ab + bias.observe(21) // ▁ab + var map = candidateMap(bias) + XCTAssertNotNil(map[22], "▁c completes the full two-word match") + XCTAssertNotNil(map[21], "▁ab is also live through the one-word overlap") + bias.observe(21) // a third ▁ab — full-match offset shifts to the overlap + map = candidateMap(bias) + XCTAssertNotNil(map[22], "▁c stays live through the repeated leading word") + } + + // MARK: - Multilingual + + func testCJKTermIsUnanchoredAndAcceptsTwoClusters() throws { + let bias = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "东京")])) + var map = candidateMap(bias) + XCTAssertNotNil(map[13], "东 opens the unanchored CJK term") + XCTAssertNil(map[14], "京 does not start the term") + bias.observe(13) + map = candidateMap(bias) + XCTAssertNotNil(map[14], "京 completes the term after 东") + } + + func testDevanagariConjunctRejoinsInScalarSpace() throws { + // SentencePiece splits the नमस्ते conjunct across pieces (स् is a + // half-form); matching must run on scalars to re-join them. + let bias = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "नमस्ते")])) + var map = candidateMap(bias) + XCTAssertNotNil(map[15], "▁नम opens the term") + bias.observe(15) + map = candidateMap(bias) + XCTAssertNotNil(map[16], "स् continues the term mid-cluster") + bias.observe(16) + map = candidateMap(bias) + XCTAssertNotNil(map[17], "ते completes the term across the split conjunct") + } + + // MARK: - Aliases and weights + + func testAliasIsBoostedAsItself() throws { + let bias = try XCTUnwrap( + makeBias([CustomVocabularyTerm(text: "torvane", aliases: ["vortane"])])) + let map = candidateMap(bias) + XCTAssertNotNil(map[4], "the primary surface is live") + XCTAssertNotNil(map[24], "▁vor opens the alias surface") + } + + func testDefaultAndOverrideBoosts() throws { + let plain = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "torvane")])) + XCTAssertEqual(candidateMap(plain)[5], NemotronVocabularyBias.defaultBoost) + let weighted = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "torvane", weight: 5.5)])) + XCTAssertEqual(candidateMap(weighted)[5], 5.5) + } + + func testWeightsAboveCapAreClamped() throws { + // The simple text-list loader assigns weight 10.0 on a CTC + // rescoring scale; applied raw as a per-token logit bonus it would + // over-bias, so the engine clamps. + let bias = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "torvane", weight: 10.0)])) + XCTAssertEqual(candidateMap(bias)[5], NemotronVocabularyBias.maxBoost) + XCTAssertEqual( + NemotronVocabularyBias.effectiveBoost(of: CustomVocabularyTerm(text: "x", weight: 10)), + NemotronVocabularyBias.maxBoost) + } + + func testStrongestBoostWinsPerToken() throws { + let bias = try XCTUnwrap( + makeBias([ + CustomVocabularyTerm(text: "torvane", weight: 2.0), + CustomVocabularyTerm(text: "torment", weight: 5.0), + ])) + // ▁tor prefixes both terms; the stronger boost must win. + XCTAssertEqual(candidateMap(bias)[4], 5.0) + } + + // MARK: - Match-state lifecycle + + func testSpecialPiecesNeverDisturbMatchState() throws { + let bias = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "torvane")])) + bias.observe(4) // ▁tor + bias.observe(19) // lang tag — no piece text + bias.observe(18) // + XCTAssertNotNil(candidateMap(bias)[3], "vane stays live across special-token emissions") + } + + func testUnmatchedEmissionDropsContinuation() throws { + let bias = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "torvane")])) + bias.observe(4) // ▁tor + bias.observe(7) // ▁the — breaks the match + XCTAssertNil(candidateMap(bias)[3], "vane is no longer a continuation") + XCTAssertNotNil(candidateMap(bias)[4], "fresh start is still live") + } + + func testResetMatchStateClearsContinuationsKeepsVocabulary() throws { + let bias = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "torvane")])) + bias.observe(4) + XCTAssertNotNil(candidateMap(bias)[3]) + bias.resetMatchState() + let map = candidateMap(bias) + XCTAssertNil(map[3], "continuation gone after reset") + XCTAssertNotNil(map[5], "fresh-start candidates survive reset") + } + + // MARK: - Biased selection + + func testPickBiasedFlipsOnlyWhenBoostedLogitWins() { + let logits: [Float] = [1.0, 5.0, 3.0] + let read: (Int) -> Float = { logits[$0] } + // 3.0 + 4.5 beats 5.0 → flip. + XCTAssertEqual( + StreamingNemotronMultilingualAsrManager.pickBiased( + plain: 1, + candidates: [.init(tokenId: 2, boost: 4.5)], count: 3, logit: read), + 2) + // 1.0 + 3.0 loses to 5.0 → no flip. + XCTAssertEqual( + StreamingNemotronMultilingualAsrManager.pickBiased( + plain: 1, + candidates: [.init(tokenId: 0, boost: 3.0)], count: 3, logit: read), + 1) + // Out-of-range candidate ids are ignored, never read. + XCTAssertEqual( + StreamingNemotronMultilingualAsrManager.pickBiased( + plain: 1, + candidates: [.init(tokenId: 7, boost: 100)], count: 3, logit: read), + 1) + } + + func testPickBiasedCanOvertakeBlank() { + // Blank (id 2) is the plain argmax; a boosted term token must be + // able to overtake it — that is how a term the decoder was about to + // drop gets emitted. + let logits: [Float] = [2.0, 1.0, 4.0] + XCTAssertEqual( + StreamingNemotronMultilingualAsrManager.pickBiased( + plain: 2, + candidates: [.init(tokenId: 0, boost: 4.5)], count: 3, logit: { logits[$0] }), + 0) + } + + // MARK: - Scale + + func testLargeVocabularySmoke() throws { + // 2000 terms sharing no prefix with the emitted tail: per-step cost + // must stay bounded by the tail walk, not the term count. This is a + // functional smoke (timing belongs to a benchmark rig), but a + // regression to per-term scanning would show up as a timeout here. + var terms = (0..<2000).map { CustomVocabularyTerm(text: "zqterm\($0)vx") } + terms.append(CustomVocabularyTerm(text: "torvane")) + let bias = try XCTUnwrap(makeBias(terms)) + for _ in 0..<200 { + bias.observe(4) // ▁tor + XCTAssertNotNil(candidateMap(bias)[3]) + bias.observe(7) // ▁the + XCTAssertFalse(bias.candidates().isEmpty) + } + } +} From d959981baa373b4b1bd74f4383b07d87d9bcba42 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Wed, 19 Aug 2026 02:37:53 -0400 Subject: [PATCH 2/4] =?UTF-8?q?feat(cli):=20nemotron-vocab-benchmark=20?= =?UTF-8?q?=E2=80=94=20paired=20baseline-vs-biased=20custom=20vocabulary?= =?UTF-8?q?=20benchmark?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Runs each earnings22-kws chunk twice with the same manager (empty vocabulary vs the chunk's dictionary terms) and scores both with the presence-based TP/FP/FN scheme ctc-earnings-benchmark uses, so the Nemotron decode-time biasing path and the CTC rescoring path are comparable on the same rig. First 50 chunks, default boost 4.5, multilingual 2240ms B1 assets: baseline WER 21.58% vocab recall 29.8% (TP=28 FN=66) FP=0 RTFx 83.4x biased WER 20.07% vocab recall 48.9% (TP=46 FN=48) FP=0 RTFx 83.2x Recall +19.1 points with WER improving 1.5 points and no measurable decode overhead from the matcher. Two per-file regressions observed, both the greedy over-boost artifact class: an already-correct term perturbed by its own boost (Latam->Latan, Andres->andress). --- .../Streaming/NemotronVocabBenchmark.swift | 291 ++++++++++++++++++ Sources/FluidAudioCLI/FluidAudioCLI.swift | 3 + 2 files changed, 294 insertions(+) create mode 100644 Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronVocabBenchmark.swift diff --git a/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronVocabBenchmark.swift b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronVocabBenchmark.swift new file mode 100644 index 000000000..ff1d5d087 --- /dev/null +++ b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronVocabBenchmark.swift @@ -0,0 +1,291 @@ +#if os(macOS) +import AVFoundation +import FluidAudio +import Foundation + +/// Paired baseline-vs-biased benchmark for Nemotron decode-time custom +/// vocabulary (issue #841) on the earnings22-kws chunk dataset (the NeMo +/// CTC-WS keyword-spotting rig: per-chunk `.wav` + `.text.txt` reference + +/// `.dictionary.txt` keywords). +/// +/// Each file is transcribed twice by the same manager — once with an empty +/// vocabulary, once with the chunk's dictionary terms — and scored with the +/// same presence-based TP/FP/FN scheme `ctc-earnings-benchmark` uses, so the +/// two vocabulary paths are comparable. +public enum NemotronVocabBenchmark { + + private struct FileResult { + let fileId: String + let audioSeconds: Double + var wer: [String: Double] = [:] + var tp: [String: Int] = [:] + var fp: [String: Int] = [:] + var fn: [String: Int] = [:] + var seconds: [String: Double] = [:] + var hypothesis: [String: String] = [:] + } + + private static let conditions = ["baseline", "biased"] + + public static func runCLI(arguments: [String]) async { + var dataDir = + FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0] + .appendingPathComponent("FluidAudio/earnings22-kws/test-dataset").path + var modelDir: String? = nil + var maxFiles = 50 + var language = "en-US" + var weightOverride: Float? = nil + var outputFile: String? = nil + + var i = 0 + while i < arguments.count { + switch arguments[i] { + case "--data-dir": + i += 1 + if i < arguments.count { dataDir = arguments[i] } + case "--model-dir", "-m": + i += 1 + if i < arguments.count { modelDir = arguments[i] } + case "--max-files": + i += 1 + if i < arguments.count, let n = Int(arguments[i]) { maxFiles = n } + case "--language", "-l": + i += 1 + if i < arguments.count { language = arguments[i] } + case "--weight": + i += 1 + if i < arguments.count, let w = Float(arguments[i]) { weightOverride = w } + case "--output", "-o": + i += 1 + if i < arguments.count { outputFile = arguments[i] } + case "--help", "-h": + printUsage() + return + default: + print("Unknown argument: \(arguments[i])") + } + i += 1 + } + + let dataURL = URL(fileURLWithPath: dataDir) + guard FileManager.default.fileExists(atPath: dataURL.path) else { + print("Data directory not found: \(dataDir)") + print("Download with: fluidaudiocli download --dataset earnings22-kws") + return + } + + do { + let resolvedModelDir: URL + if let modelDir { + resolvedModelDir = URL(fileURLWithPath: modelDir) + } else { + resolvedModelDir = + try await StreamingNemotronMultilingualAsrManager + .downloadVariant(languageCode: language, chunkMs: 2240) + } + + let manager = StreamingNemotronMultilingualAsrManager() + try await manager.loadModels(from: resolvedModelDir) + await manager.setLanguage(language) + + // File ids = every .wav with a reference and a dictionary beside it. + let allFiles = try FileManager.default.contentsOfDirectory(atPath: dataURL.path) + let fileIds = allFiles.filter { $0.hasSuffix(".wav") } + .map { String($0.dropLast(4)) } + .filter { id in + allFiles.contains("\(id).text.txt") && allFiles.contains("\(id).dictionary.txt") + } + .sorted() + .prefix(maxFiles) + + print("Nemotron custom-vocabulary benchmark (issue #841)") + print("Model: \(resolvedModelDir.path)") + print("Files: \(fileIds.count) (of \(allFiles.filter { $0.hasSuffix(".wav") }.count) available)") + print("Boost: \(weightOverride.map { String($0) } ?? "default (4.5)")\n") + + var results: [FileResult] = [] + let converter = AudioConverter() + + for (index, fileId) in fileIds.enumerated() { + let wavURL = dataURL.appendingPathComponent("\(fileId).wav") + let reference = + (try? String( + contentsOf: dataURL.appendingPathComponent("\(fileId).text.txt"), + encoding: .utf8)) ?? "" + let dictWords = + ((try? String( + contentsOf: dataURL.appendingPathComponent("\(fileId).dictionary.txt"), + encoding: .utf8)) ?? "") + .components(separatedBy: .newlines) + .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } + .filter { !$0.isEmpty } + guard !reference.isEmpty, !dictWords.isEmpty else { continue } + + let audioFile = try AVAudioFile(forReading: wavURL) + guard + let buffer = AVAudioPCMBuffer( + pcmFormat: audioFile.processingFormat, + frameCapacity: AVAudioFrameCount(audioFile.length)) + else { continue } + try audioFile.read(into: buffer) + let samples = try converter.resampleBuffer(buffer) + let audioSeconds = Double(audioFile.length) / audioFile.processingFormat.sampleRate + + let terms = dictWords.map { CustomVocabularyTerm(text: $0, weight: weightOverride) } + var result = FileResult(fileId: fileId, audioSeconds: audioSeconds) + + for condition in conditions { + await manager.setCustomVocabulary(condition == "biased" ? terms : []) + let start = Date() + _ = try await manager.process(samples: samples) + let hypothesis = try await manager.finish() + result.seconds[condition] = Date().timeIntervalSince(start) + await manager.reset() + + let metrics = WERCalculator.calculateWERMetrics( + hypothesis: hypothesis, reference: reference) + result.wer[condition] = metrics.wer + result.hypothesis[condition] = hypothesis + + let refLower = TextNormalizer.normalize(reference).lowercased() + let hypLower = TextNormalizer.normalize(hypothesis).lowercased() + var tp = 0 + var fp = 0 + var fn = 0 + for word in dictWords { + let inRef = containsWholeWord(refLower, word) + let inHyp = containsWholeWord(hypLower, word) + if inRef && inHyp { + tp += 1 + } else if inHyp { + fp += 1 + } else if inRef { + fn += 1 + } + } + result.tp[condition] = tp + result.fp[condition] = fp + result.fn[condition] = fn + } + results.append(result) + + let bWer = (result.wer["baseline"] ?? 0) * 100 + let vWer = (result.wer["biased"] ?? 0) * 100 + let bRecall = recallString(result, "baseline") + let vRecall = recallString(result, "biased") + print( + String( + format: "[%3d/%d] %@ WER %5.1f%% -> %5.1f%% dict %@ -> %@", + index + 1, fileIds.count, + fileId.padding(toLength: 24, withPad: " ", startingAt: 0), + bWer, vWer, bRecall, vRecall)) + } + + printSummary(results, weightOverride: weightOverride) + if let outputFile { + try writeJSON(results, to: outputFile, weightOverride: weightOverride) + print("\nResults written to \(outputFile)") + } + } catch { + print("Benchmark failed: \(error)") + } + } + + private static func containsWholeWord(_ text: String, _ word: String) -> Bool { + let pattern = "\\b\(NSRegularExpression.escapedPattern(for: word.lowercased()))\\b" + guard let regex = try? NSRegularExpression(pattern: pattern) else { + return text.contains(word.lowercased()) + } + return regex.firstMatch(in: text, range: NSRange(text.startIndex..., in: text)) != nil + } + + private static func recallString(_ result: FileResult, _ condition: String) -> String { + let tp = result.tp[condition] ?? 0 + let fn = result.fn[condition] ?? 0 + return "\(tp)/\(tp + fn)" + } + + private static func printSummary(_ results: [FileResult], weightOverride: Float?) { + guard !results.isEmpty else { + print("No files processed") + return + } + print("\n" + String(repeating: "=", count: 72)) + print("NEMOTRON CUSTOM VOCABULARY BENCHMARK (earnings22-kws)") + print(String(repeating: "=", count: 72)) + print("Files: \(results.count) Boost: \(weightOverride.map { String($0) } ?? "default (4.5)")") + let audio = results.reduce(0.0) { $0 + $1.audioSeconds } + print(String(format: "Audio: %.1fs", audio)) + for condition in conditions { + let avgWer = + results.reduce(0.0) { $0 + ($1.wer[condition] ?? 0) } / Double(results.count) * 100 + let tp = results.reduce(0) { $0 + ($1.tp[condition] ?? 0) } + let fp = results.reduce(0) { $0 + ($1.fp[condition] ?? 0) } + let fn = results.reduce(0) { $0 + ($1.fn[condition] ?? 0) } + let recall = tp + fn > 0 ? Double(tp) / Double(tp + fn) * 100 : 0 + let precision = tp + fp > 0 ? Double(tp) / Double(tp + fp) * 100 : 0 + let time = results.reduce(0.0) { $0 + ($1.seconds[condition] ?? 0) } + let rtfx = time > 0 ? audio / time : 0 + print( + String( + format: "%@ WER %6.2f%% recall %5.1f%% (TP=%d FN=%d) precision %5.1f%% (FP=%d) RTFx %.1fx", + condition.padding(toLength: 9, withPad: " ", startingAt: 0), + avgWer, recall, tp, fn, precision, fp, rtfx)) + } + } + + private static func writeJSON( + _ results: [FileResult], to path: String, weightOverride: Float? + ) throws { + let files: [[String: Any]] = results.map { result in + var entry: [String: Any] = [ + "fileId": result.fileId, + "audioSeconds": result.audioSeconds, + ] + for condition in conditions { + entry[condition] = [ + "wer": result.wer[condition] ?? 0, + "tp": result.tp[condition] ?? 0, + "fp": result.fp[condition] ?? 0, + "fn": result.fn[condition] ?? 0, + "seconds": result.seconds[condition] ?? 0, + "hypothesis": result.hypothesis[condition] ?? "", + ] + } + return entry + } + let payload: [String: Any] = [ + "benchmark": "nemotron-vocab-benchmark", + "boost": weightOverride.map { Double($0) } ?? 4.5, + "files": files, + ] + let data = try JSONSerialization.data( + withJSONObject: payload, options: [.prettyPrinted, .sortedKeys]) + try data.write(to: URL(fileURLWithPath: path)) + } + + private static func printUsage() { + print( + """ + Nemotron decode-time custom vocabulary benchmark (issue #841) + + Runs each earnings22-kws chunk twice — with and without the chunk's + dictionary terms — and reports paired WER / vocab recall / precision. + + Usage: fluidaudio nemotron-vocab-benchmark [options] + + Options: + --data-dir earnings22-kws test-dataset directory + (default: app-support copy) + --model-dir Nemotron multilingual model directory + (default: auto-download 2240ms variant) + --max-files Number of chunks to run (default: 50) + --language Language hint (default: en-US) + --weight Per-term weight override (default boost 4.5; + clamped to 6.0 by the engine) + --output, -o Write per-file JSON results + """ + ) + } +} +#endif diff --git a/Sources/FluidAudioCLI/FluidAudioCLI.swift b/Sources/FluidAudioCLI/FluidAudioCLI.swift index 988b47a60..057ba1263 100644 --- a/Sources/FluidAudioCLI/FluidAudioCLI.swift +++ b/Sources/FluidAudioCLI/FluidAudioCLI.swift @@ -84,6 +84,8 @@ struct FluidAudioCLI { await NemotronMultilingualTranscribe.run(arguments: Array(arguments.dropFirst(2))) case "nemotron-multilingual-benchmark": await NemotronMultilingualFleursBenchmark.runCLI(arguments: Array(arguments.dropFirst(2))) + case "nemotron-vocab-benchmark": + await NemotronVocabBenchmark.runCLI(arguments: Array(arguments.dropFirst(2))) case "nemotron-multilingual-multi-stream-bench": await NemotronMultilingualMultiStreamBench.run(arguments: Array(arguments.dropFirst(2))) case "sensevoice-transcribe": @@ -149,6 +151,7 @@ struct FluidAudioCLI { nemotron-transcribe Transcribe custom audio files with Nemotron nemotron-multilingual-transcribe Transcribe audio with Nemotron multilingual (local model path) nemotron-multilingual-benchmark Run Nemotron multilingual benchmark on FLEURS / MCV-17 / MLS (local model path) + nemotron-vocab-benchmark Paired baseline-vs-biased custom vocabulary benchmark (earnings22-kws) nemotron-multilingual-multi-stream-bench Parallel multi-stream benchmark (N concurrent managers) ja-benchmark Run Japanese ASR benchmark on JSUT/Common Voice cohere-transcribe Transcribe using Cohere Transcribe (cache-external pipeline, 14 languages) From dc8d9992021953de2463309ca44de3792b9260c3 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Wed, 19 Aug 2026 02:43:22 -0400 Subject: [PATCH 3/4] docs(asr/nemotron): chunk boundaries are not a barrier for vocabulary biasing The 'decoded as two independent halves' caveat was inherited from the sliding-window CTC path, where windows really are re-decoded independently. In the Nemotron streaming manager both the decoder state and the bias match tail persist across chunk boundaries (only reset()/finish() clear them), so cross-boundary terms keep their bias. Verified empirically: stepping leading silence through a full 2240ms chunk period (0-1960ms, 280ms steps) on LibriSpeech 7127-75947-0033 sweeps the boundary through the 'Tonnay Charente' audio; biased recovery holds at every phase (Tonnay in 6/8, full Tonnay Charente at two phases) while the unbiased baseline never recovers it. Also records the full 772-chunk earnings22-kws paired benchmark: baseline WER 19.77% recall 39.3% (TP=494 FN=762) FP=1 RTFx 88.2x biased WER 18.68% recall 59.7% (TP=750 FN=506) FP=3 RTFx 88.0x --- Documentation/ASR/NemotronMultilingual.md | 9 +++++++-- .../Streaming/Nemotron/NemotronVocabularyBias.swift | 9 ++++++--- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/Documentation/ASR/NemotronMultilingual.md b/Documentation/ASR/NemotronMultilingual.md index 087d570e1..cee1e945e 100644 --- a/Documentation/ASR/NemotronMultilingual.md +++ b/Documentation/ASR/NemotronMultilingual.md @@ -115,8 +115,13 @@ Semantics and limits: there is no beam to undo an over-fire. Keep vocabularies to genuinely rare terms; a term the model hears as a word it already spells ("Pheynix" → "Phoenix") needs an alias, not more weight. -- **Chunk boundaries.** A term whose audio spans a chunk boundary is decoded - as two independent halves; single-word terms are the reliable target. +- **Chunk boundaries are not a barrier.** Unlike the sliding-window CTC + path, the decoder state and the term match state persist across chunks + (only `reset()`/`finish()` clear them), so a term whose audio spans a + boundary keeps its bias. Verified by a boundary-phase sweep: stepping + leading silence through a full 2240 ms chunk period, term recovery holds + at every phase. Multi-word terms are still harder — more greedy steps + must go right — but not because of chunking. - Terms shorter than 3 letters are skipped (2 for CJK); word-start anchoring keeps "ran" from matching into "CRAN". CJK terms match unanchored. - **False-fire profile** (LibriSpeech test-clean spot check, 40 invented diff --git a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/NemotronVocabularyBias.swift b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/NemotronVocabularyBias.swift index 61df4be78..de81c6024 100644 --- a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/NemotronVocabularyBias.swift +++ b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/NemotronVocabularyBias.swift @@ -46,9 +46,12 @@ import Foundation /// /// Greedy-decoding caveat: this biases a greedy argmax, so once a boosted /// token wins it is committed — there is no beam to recover an over-fire. -/// Streaming caveat (same as the CTC path documents): a term whose audio -/// spans a chunk boundary is decoded as two independent halves, so -/// single-word terms are the reliable target. +/// Chunk boundaries are NOT a barrier, unlike the sliding-window CTC path: +/// the decoder state and this match tail both persist across chunks (only +/// `reset()`/`finish()` clear them), so a term whose audio spans a +/// boundary keeps its bias. Verified by a boundary-phase sweep — leading +/// silence stepped through a full chunk period leaves term recovery intact +/// at every phase. final class NemotronVocabularyBias { /// One boostable token continuation at the current match state. From c44696f757be519533ce754c4af2559fdc6135ae Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Wed, 19 Aug 2026 04:11:03 -0400 Subject: [PATCH 4/4] fix(asr/nemotron): address vocabulary-biasing review findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. B2 no longer blocks the logits fallback it advertises. The load path's tier priority skips decoder_joint_noencproj/decoder_joint when decoder_joint_argmax exists, so a bundle shipping B2 alongside a logits decoder still had biasing rejected (or silently falling to the slow bare pair). rebuildVocabularyBias now lazily loads decoder_joint from the remembered model directory before rejecting. B1 specifically: it takes the plain encoder step every decode path has, whereas B3 also needs an encoder that emits encoder_proj — loading a model the decode loop can't feed would silently fall back to unbiased B2, the exact failure this avoids. setCustomVocabulary/rebuildVocabularyBias become async for the lazy load. 2. Legacy weights above maxBoost fall back to the 4.5 default instead of clamping to the aggressive 6.0 ceiling. The simple text-list loader blankets every term with CTC-scale weight 10.0, so plain text files were silently pinned to the hottest allowed boost with no way to request the measured default. Explicit weights in (0, 6.0] are still honored. Verified live: a text-list vocabulary now logs the fallback warning and decodes identically to the 4.5 JSON run. 3. Malformed JSON vocabulary files now surface their parse error instead of being reinterpreted as one-hotword-per-line text. CustomVocabularyContext.loadVocabularyFile (first-meaningful-byte format detection) is now public and the CLI uses it; the try/fallback in nemotron-multilingual-transcribe is gone. Verified live: truncated JSON aborts the run with the decode error. --- Documentation/ASR/NemotronMultilingual.md | 12 ++- .../CustomVocabularyContext.swift | 6 +- .../Nemotron/NemotronVocabularyBias.swift | 79 ++++++++++++++----- ...eamingNemotronMultilingualAsrManager.swift | 12 ++- .../NemotronMultilingualTranscribe.swift | 7 +- .../Streaming/NemotronVocabBenchmark.swift | 5 +- .../NemotronVocabularyBiasTests.swift | 13 ++- 7 files changed, 96 insertions(+), 38 deletions(-) diff --git a/Documentation/ASR/NemotronMultilingual.md b/Documentation/ASR/NemotronMultilingual.md index cee1e945e..075811135 100644 --- a/Documentation/ASR/NemotronMultilingual.md +++ b/Documentation/ASR/NemotronMultilingual.md @@ -103,14 +103,18 @@ Semantics and limits: - **Weight scale.** `weight` here is a *per-token logit bonus* (default 4.5, the measured recall peak), not the CTC rescoring weight. Values above 6.0 - are clamped — the simple text-list loader's `weight: 10.0` would otherwise + are treated as legacy CTC-scale weights and fall back to the 4.5 default — + the simple text-list loader's blanket `weight: 10.0` would otherwise over-bias (measured artifacts: word splits like "build today" → "build to - day"). Most terms should omit `weight`. + day"). Most terms should omit `weight`; explicit values in (0, 6.0] are + honored. - **Assets.** Biasing needs a logits-producing step decoder (`decoder_joint_noencproj` or `decoder_joint`). When a vocabulary is active the fused-argmax `decoder_joint_argmax` asset is bypassed in favor - of a logits path; if it is the *only* step decoder, the vocabulary is - rejected with an error log rather than silently half-applied. + of a logits path; if the load-time tier priority skipped `decoder_joint`, + it is loaded lazily from the model directory when the vocabulary is set. + Only if no logits decoder exists at all is the vocabulary rejected, with + an error log rather than silent half-application. - **Greedy decode.** Once a boosted token wins the argmax it is committed — there is no beam to undo an over-fire. Keep vocabularies to genuinely rare terms; a term the model hears as a word it already spells diff --git a/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/CustomVocabulary/CustomVocabularyContext.swift b/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/CustomVocabulary/CustomVocabularyContext.swift index 666dc36ad..3a70b160b 100644 --- a/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/CustomVocabulary/CustomVocabularyContext.swift +++ b/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/CustomVocabulary/CustomVocabularyContext.swift @@ -317,8 +317,10 @@ public struct CustomVocabularyContext: Sendable { /// Load a vocabulary file, auto-detecting the structured JSON config vs the /// simple one-term-per-line text format. JSON config files begin with `{` /// (after optional leading whitespace); anything else is treated as simple - /// text. - static func loadVocabularyFile(at url: URL) throws -> CustomVocabularyContext { + /// text. Detection is by first meaningful byte, not by try-and-fallback, + /// so a malformed JSON config surfaces its parse error instead of being + /// silently reinterpreted as a list of hotwords. + public static func loadVocabularyFile(at url: URL) throws -> CustomVocabularyContext { let data = try Data(contentsOf: url) let whitespace: Set = [0x20, 0x09, 0x0a, 0x0d] // space, tab, LF, CR let firstMeaningfulByte = data.first { !whitespace.contains($0) } diff --git a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/NemotronVocabularyBias.swift b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/NemotronVocabularyBias.swift index de81c6024..5a9b3e2ee 100644 --- a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/NemotronVocabularyBias.swift +++ b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/NemotronVocabularyBias.swift @@ -1,4 +1,4 @@ -import CoreML +@preconcurrency import CoreML import Foundation /// Decode-time shallow-fusion hotword biasing for the Nemotron streaming @@ -93,7 +93,8 @@ final class NemotronVocabularyBias { /// predates this engine and is set to 10.0 by the simple text-list /// loader — a CTC-rescoring scale, not a per-token logit bonus. Applied /// raw it would over-bias badly (degradation is measurable from 6.0), - /// so overrides are clamped here rather than reinterpreted. + /// so weights above this fall back to `defaultBoost` (see + /// `effectiveBoost`). static let maxBoost: Float = 6.0 private let entries: [Entry] @@ -187,9 +188,14 @@ final class NemotronVocabularyBias { } /// The per-token bonus a term will actually receive: its `weight` when - /// set (clamped into `(0, maxBoost]`), the default otherwise. + /// set and within `(0, maxBoost]`, the default otherwise. A weight above + /// `maxBoost` was tuned for the CTC rescoring scale (the simple text-list + /// loader assigns 10.0 to every term), not for this engine — treating it + /// as "untuned" and using the measured default beats pinning those terms + /// to the aggressive ceiling. static func effectiveBoost(of term: CustomVocabularyTerm, defaultBoost: Float = defaultBoost) -> Float { - min(term.weight ?? defaultBoost, maxBoost) + guard let weight = term.weight, weight <= maxBoost else { return defaultBoost } + return weight } /// The surfaces of `term` (text + aliases) that pass the length and @@ -350,19 +356,21 @@ extension StreamingNemotronMultilingualAsrManager { /// effect from the next emission. May be called before models load — /// the vocabulary is (re)bound whenever a tokenizer becomes available. /// - /// `weight` is a per-token log-prob bonus clamped to - /// `NemotronVocabularyBias.maxBoost`; most terms should omit it. - /// Biasing requires a logits-producing step decoder: when only the - /// fused-argmax (B2) asset is loaded, the vocabulary is rejected with - /// an error log rather than silently half-applied. - public func setCustomVocabulary(_ terms: [CustomVocabularyTerm]) { + /// `weight` is a per-token log-prob bonus; most terms should omit it + /// (default 4.5). Weights above `NemotronVocabularyBias.maxBoost` are + /// treated as legacy CTC-scale values and fall back to the default. + /// Biasing requires a logits-producing step decoder: one is loaded + /// lazily from the model directory when the tier priority skipped it, + /// and when none exists the vocabulary is rejected with an error log + /// rather than silently half-applied. + public func setCustomVocabulary(_ terms: [CustomVocabularyTerm]) async { vocabularyTerms = terms - rebuildVocabularyBias() + await rebuildVocabularyBias() } /// Bind the stored terms to the loaded tokenizer. Called from the model /// load path and from `setCustomVocabulary`. - internal func rebuildVocabularyBias() { + internal func rebuildVocabularyBias() async { guard !vocabularyTerms.isEmpty, let tokenizer else { vocabularyBias = nil return @@ -370,8 +378,16 @@ extension StreamingNemotronMultilingualAsrManager { let anyStepDecoder = decoderJointNoEncProj != nil || decoderJointArgmax != nil || decoderJoint != nil || (decoder != nil && joint != nil) - let logitsStepDecoder = + var logitsStepDecoder = decoderJointNoEncProj != nil || decoderJoint != nil || (decoder != nil && joint != nil) + if anyStepDecoder && !logitsStepDecoder { + // The normal load path skips the logits-producing fused decoders + // when B2 wins the tier priority, but the bundle may still ship + // one — bring up B1 (plain `encoder` input, works with every + // encoder) before rejecting. + await loadLogitsStepDecoderForVocabulary() + logitsStepDecoder = decoderJoint != nil + } if anyStepDecoder && !logitsStepDecoder { // B2-only asset set: the fused-argmax model never exposes // logits, so no decode site can be biased. All-or-nothing — @@ -383,12 +399,13 @@ extension StreamingNemotronMultilingualAsrManager { + "decoder_joint_noencproj or decoder_joint asset to enable it.") return } - let clamped = vocabularyTerms.filter { ($0.weight ?? 0) > NemotronVocabularyBias.maxBoost } - if !clamped.isEmpty { + let legacyWeights = vocabularyTerms.filter { ($0.weight ?? 0) > NemotronVocabularyBias.maxBoost } + if !legacyWeights.isEmpty { logger.warning( - "Custom vocabulary: \(clamped.count) term(s) with weight > " - + "\(NemotronVocabularyBias.maxBoost) clamped (weight is a per-token logit bonus " - + "here, not a CTC rescoring weight)") + "Custom vocabulary: \(legacyWeights.count) term(s) with weight > " + + "\(NemotronVocabularyBias.maxBoost) use the \(NemotronVocabularyBias.defaultBoost) " + + "default instead (a weight that large is a CTC rescoring value, not a per-token " + + "logit bonus)") } // The multilingual wrapper hides the base vocabulary map, but ids // are contiguous and `rawToken(for:)` is exact — rebuild the table. @@ -409,6 +426,32 @@ extension StreamingNemotronMultilingualAsrManager { } } + /// Load the B1 fused decoder (`decoder_joint`) from the remembered model + /// directory so an active vocabulary has a logits-producing step decoder + /// on a bundle where B2 won the tier priority. B1 is the safe choice: + /// it takes the plain `encoder` step every decode path already has, + /// whereas B3 (`decoder_joint_noencproj`) also needs the encoder to emit + /// `encoder_proj` — loading a model the decode loop can't feed would + /// silently fall back to unbiased B2, the exact failure this avoids. + private func loadLogitsStepDecoderForVocabulary() async { + guard decoderJoint == nil, let directory = modelDirectory else { return } + do { + guard + let fusedURL = try await locateOptionalModelBundle( + in: directory, compiled: "decoder_joint.mlmodelc", + uncompiled: "decoder_joint.mlpackage") + else { return } + decoderJoint = try await MLModel.load(contentsOf: fusedURL, configuration: mlConfiguration) + logger.info( + "Loaded decoder_joint for custom vocabulary — the fused-argmax (B2) path exposes no " + + "logits, so biased decoding uses B1") + } catch { + logger.warning( + "Custom vocabulary: failed to load decoder_joint from \(directory.path): " + + "\(error.localizedDescription)") + } + } + /// When hotword biasing is active the fused-argmax step decoder (B2) /// must yield to a logits-producing alternative — no logits, no bias. internal var vocabularyBiasPrefersLogits: Bool { diff --git a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift index fda0a1483..914eb1c7a 100644 --- a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift +++ b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift @@ -94,6 +94,8 @@ public actor StreamingNemotronMultilingualAsrManager { // rebuilt whenever a tokenizer becomes available. internal var vocabularyTerms: [CustomVocabularyTerm] = [] internal var vocabularyBias: NemotronVocabularyBias? + // Directory loadModels() read from, for lazy vocabulary-driven loads. + internal var modelDirectory: URL? // Per-token absolute timings captured during the RNNT decode loop, parallel // to the user-visible (lang-tag-stripped) token stream. Each token's @@ -325,6 +327,9 @@ public actor StreamingNemotronMultilingualAsrManager { } logger.info("Loading Nemotron multilingual CoreML models from \(directory.path)...") + // Remembered so vocabulary biasing can lazily load a logits-producing + // step decoder the tier priority skipped (see rebuildVocabularyBias). + self.modelDirectory = directory // Load config from metadata.json (required — the prompt dictionary lives here) let metadataPath = directory.appendingPathComponent(ModelNames.NemotronMultilingualStreaming.metadata) @@ -497,7 +502,7 @@ public actor StreamingNemotronMultilingualAsrManager { vocabPath: tokenizerURL, langTagTokenIds: config.langTagTokenIds ) - rebuildVocabularyBias() + await rebuildVocabularyBias() // Initialize states try resetStates() @@ -726,8 +731,9 @@ public actor StreamingNemotronMultilingualAsrManager { /// load site still gets the caching behavior (mlpackage → cached /// .mlmodelc next to source) instead of compiling to a temp dir per /// cold start. - private func locateOptionalModelBundle(in directory: URL, compiled: String, uncompiled: String) async throws -> URL? - { + internal func locateOptionalModelBundle( + in directory: URL, compiled: String, uncompiled: String + ) async throws -> URL? { let compiledURL = directory.appendingPathComponent(compiled) let uncompiledURL = directory.appendingPathComponent(uncompiled) if !FileManager.default.fileExists(atPath: compiledURL.path) diff --git a/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronMultilingualTranscribe.swift b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronMultilingualTranscribe.swift index 1ba52809a..d5beaf402 100644 --- a/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronMultilingualTranscribe.swift +++ b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronMultilingualTranscribe.swift @@ -202,12 +202,7 @@ public class NemotronMultilingualTranscribe { if let vocabPath = config.customVocabPath { let vocabURL = URL(fileURLWithPath: vocabPath) - let vocab: CustomVocabularyContext - do { - vocab = try CustomVocabularyContext.load(from: vocabURL) - } catch { - vocab = try CustomVocabularyContext.loadFromSimpleFormat(from: vocabURL) - } + let vocab = try CustomVocabularyContext.loadVocabularyFile(at: vocabURL) await manager.setCustomVocabulary(vocab.terms) logger.info("Custom vocabulary: \(vocab.terms.count) term(s) from \(vocabPath)") } diff --git a/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronVocabBenchmark.swift b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronVocabBenchmark.swift index ff1d5d087..4b3afb044 100644 --- a/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronVocabBenchmark.swift +++ b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronVocabBenchmark.swift @@ -281,8 +281,9 @@ public enum NemotronVocabBenchmark { (default: auto-download 2240ms variant) --max-files Number of chunks to run (default: 50) --language Language hint (default: en-US) - --weight Per-term weight override (default boost 4.5; - clamped to 6.0 by the engine) + --weight Per-term weight override in (0, 6.0] + (default boost 4.5; values above 6.0 + fall back to the default) --output, -o Write per-file JSON results """ ) diff --git a/Tests/FluidAudioTests/ASR/Parakeet/Streaming/NemotronVocabularyBiasTests.swift b/Tests/FluidAudioTests/ASR/Parakeet/Streaming/NemotronVocabularyBiasTests.swift index b9285f257..a4f348bf9 100644 --- a/Tests/FluidAudioTests/ASR/Parakeet/Streaming/NemotronVocabularyBiasTests.swift +++ b/Tests/FluidAudioTests/ASR/Parakeet/Streaming/NemotronVocabularyBiasTests.swift @@ -153,14 +153,21 @@ final class NemotronVocabularyBiasTests: XCTestCase { XCTAssertEqual(candidateMap(weighted)[5], 5.5) } - func testWeightsAboveCapAreClamped() throws { + func testLegacyWeightsFallBackToDefault() throws { // The simple text-list loader assigns weight 10.0 on a CTC // rescoring scale; applied raw as a per-token logit bonus it would - // over-bias, so the engine clamps. + // over-bias, and pinning it to the aggressive 6.0 ceiling is still + // hotter than the measured recall peak — so out-of-range weights + // use the default instead. let bias = try XCTUnwrap(makeBias([CustomVocabularyTerm(text: "torvane", weight: 10.0)])) - XCTAssertEqual(candidateMap(bias)[5], NemotronVocabularyBias.maxBoost) + XCTAssertEqual(candidateMap(bias)[5], NemotronVocabularyBias.defaultBoost) XCTAssertEqual( NemotronVocabularyBias.effectiveBoost(of: CustomVocabularyTerm(text: "x", weight: 10)), + NemotronVocabularyBias.defaultBoost) + // An explicit in-range weight is still honored, up to the ceiling. + XCTAssertEqual( + NemotronVocabularyBias.effectiveBoost( + of: CustomVocabularyTerm(text: "x", weight: NemotronVocabularyBias.maxBoost)), NemotronVocabularyBias.maxBoost) }