diff --git a/Documentation/ASR/NemotronMultilingual.md b/Documentation/ASR/NemotronMultilingual.md index 832a93796..075811135 100644 --- a/Documentation/ASR/NemotronMultilingual.md +++ b/Documentation/ASR/NemotronMultilingual.md @@ -83,6 +83,60 @@ 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 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`; 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 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 + ("Pheynix" → "Phoenix") needs an alias, not more weight. +- **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 + 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/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 new file mode 100644 index 000000000..5a9b3e2ee --- /dev/null +++ b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/NemotronVocabularyBias.swift @@ -0,0 +1,512 @@ +@preconcurrency 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. +/// 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. + 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 weights above this fall back to `defaultBoost` (see + /// `effectiveBoost`). + 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 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 { + guard let weight = term.weight, weight <= maxBoost else { return defaultBoost } + return weight + } + + /// 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; 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 + await rebuildVocabularyBias() + } + + /// Bind the stored terms to the loaded tokenizer. Called from the model + /// load path and from `setCustomVocabulary`. + internal func rebuildVocabularyBias() async { + guard !vocabularyTerms.isEmpty, let tokenizer else { + vocabularyBias = nil + return + } + let anyStepDecoder = + decoderJointNoEncProj != nil || decoderJointArgmax != nil || decoderJoint != nil + || (decoder != nil && joint != nil) + 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 — + // 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 legacyWeights = vocabularyTerms.filter { ($0.weight ?? 0) > NemotronVocabularyBias.maxBoost } + if !legacyWeights.isEmpty { + logger.warning( + "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. + 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..914eb1c7a 100644 --- a/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift +++ b/Sources/FluidAudio/ASR/Parakeet/Streaming/Nemotron/StreamingNemotronMultilingualAsrManager.swift @@ -89,6 +89,14 @@ 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? + // 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 // startTime is its absolute encoder-frame index * secondsPerEncoderFrame. @@ -319,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) @@ -491,6 +502,7 @@ public actor StreamingNemotronMultilingualAsrManager { vocabPath: tokenizerURL, langTagTokenIds: config.langTagTokenIds ) + await rebuildVocabularyBias() // Initialize states try resetStates() @@ -719,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) @@ -781,6 +794,7 @@ public actor StreamingNemotronMultilingualAsrManager { lastFinishTokenTimings.removeAll() audioBufferOffset = 0 firstDetectedLanguage = nil + vocabularyBias?.resetMatchState() do { try resetStates() } catch { @@ -1008,6 +1022,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..d5beaf402 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,13 @@ public class NemotronMultilingualTranscribe { } else { logger.info("Using default prompt id (auto)") } + + if let vocabPath = config.customVocabPath { + let vocabURL = URL(fileURLWithPath: vocabPath) + let vocab = try CustomVocabularyContext.loadVocabularyFile(at: 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/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronVocabBenchmark.swift b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronVocabBenchmark.swift new file mode 100644 index 000000000..4b3afb044 --- /dev/null +++ b/Sources/FluidAudioCLI/Commands/ASR/Parakeet/Streaming/NemotronVocabBenchmark.swift @@ -0,0 +1,292 @@ +#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 in (0, 6.0] + (default boost 4.5; values above 6.0 + fall back to the default) + --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) diff --git a/Tests/FluidAudioTests/ASR/Parakeet/Streaming/NemotronVocabularyBiasTests.swift b/Tests/FluidAudioTests/ASR/Parakeet/Streaming/NemotronVocabularyBiasTests.swift new file mode 100644 index 000000000..a4f348bf9 --- /dev/null +++ b/Tests/FluidAudioTests/ASR/Parakeet/Streaming/NemotronVocabularyBiasTests.swift @@ -0,0 +1,266 @@ +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 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, 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.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) + } + + 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) + } + } +}