From 69c3d3ed569699d5d25896122719bed83f0ea0cd Mon Sep 17 00:00:00 2001 From: Wassim Gharbi Date: Sun, 2 Aug 2026 11:05:15 -1000 Subject: [PATCH 1/7] forced word alignment --- lib/align.ts | 358 +++++++++++++++++++++++++++++--- lib/forcedAlign.ts | 345 ++++++++++++++++++++++++++++++ tests/align-test.ts | 235 +++++++++++++++++---- tests/forced-align-test.ts | 152 ++++++++++++++ workers/transcription.worker.ts | 155 +++++++++++++- 5 files changed, 1163 insertions(+), 82 deletions(-) create mode 100644 lib/forcedAlign.ts create mode 100644 tests/forced-align-test.ts diff --git a/lib/align.ts b/lib/align.ts index 394674c..d5b08b0 100644 --- a/lib/align.ts +++ b/lib/align.ts @@ -19,8 +19,7 @@ * 2. `alignWordsToSpeech` interpolates the correction linearly between those * anchors, holding it constant outside them, which tracks the decay instead of * assuming it away. Guards keep the map strictly increasing. - * 3. `snapWordsToSpeech` then nudges remaining starts onto nearby onsets, and a - * small fixed `ALIGN_LEAD_S` is applied so words land a hair early. + * 3. `snapWordsToSpeech` then nudges remaining starts onto nearby onsets. * * Only *starts* are snapped, and only onsets are used as anchors. Silero holds a * speech flag for ~150 ms after speech actually stops, so anchoring word ends to @@ -52,14 +51,169 @@ export interface AlignOptions { landmarkTolS?: number; /** Words are never shortened below this. Default 0.02. */ minWordS?: number; - /** - * Extra seconds subtracted from every timestamp after alignment, so words land - * a hair early rather than a hair late. Default `ALIGN_LEAD_S`; set 0 to - * measure raw alignment accuracy. - */ - leadS?: number; /** Media duration; times are clamped to it when > 0. */ duration?: number; + /** + * Mono PCM at `sampleRate`. When supplied, the global shift is estimated from + * loudness rises rather than the VAD mask, which is the only thing that works on + * speech with no pauses. Strongly recommended. + */ + audio?: Float32Array; + /** + * Speech onsets to anchor and snap to, already refined against the audio. When + * omitted they are derived from `speechFrames` unrefined. `alignWordsToSpeech` + * sets this so the envelope is built once rather than per helper. + */ + onsets?: number[]; +} + +/** Hop between envelope frames, in seconds. 5 ms is well under the error we chase. */ +export const ENVELOPE_HOP_S = 0.005; + +/** High-pass corner and order used to isolate fricative energy. */ +const HIGH_BAND_HZ = 2000; +const HIGH_BAND_STAGES = 3; + +/** + * Loudness of the audio, sampled every `ENVELOPE_HOP_S`. + * + * `level` is the max of a broadband and a high-band envelope, each normalised + * against its own 95th percentile. The high band matters: an unvoiced fricative + * (/ʃ/, /s/, /f/) carries almost all its energy above 3 kHz and is far quieter + * than a vowel, so a broadband envelope puts the onset of "shot" on the vowel + * rather than on the /ʃ/ — about 0.2 s late. Taking the max of two separately + * normalised bands lets either one mark an onset. + */ +export interface SpeechEnvelope { + hopS: number; + /** Normalised loudness, roughly 0..1. */ + level: Float32Array; + /** Half-wave-rectified rise of `level`; word onsets land on peaks here. */ + rise: Float32Array; + /** Noise floor, taken as the 10th percentile of `level`. */ + floor: number; +} + +/** Value at percentile `p` (0..1) of a copy of `values`. */ +function percentile(values: Float32Array, p: number): number { + if (values.length === 0) return 0; + const sorted = values.slice(); + sorted.sort(); + return sorted[Math.min(sorted.length - 1, Math.floor(p * sorted.length))]; +} + +/** + * Build a speech envelope from mono PCM in a single pass. The high-band filter + * runs sample-by-sample against a small state array rather than materialising a + * filtered copy, which would double peak memory on long files. + */ +export function speechEnvelope( + audio: Float32Array, + sampleRate = VAD_SAMPLE_RATE, + hopS = ENVELOPE_HOP_S +): SpeechEnvelope { + const hop = Math.max(1, Math.round(hopS * sampleRate)); + const n = Math.floor(audio.length / hop); + const broad = new Float32Array(n); + const high = new Float32Array(n); + + const dt = 1 / sampleRate; + const rc = 1 / (2 * Math.PI * HIGH_BAND_HZ); + const a = rc / (rc + dt); + const prevIn = new Float64Array(HIGH_BAND_STAGES); + const prevOut = new Float64Array(HIGH_BAND_STAGES); + + for (let f = 0; f < n; f++) { + const start = f * hop; + const end = Math.min(audio.length, start + hop); + let sumBroad = 0; + let sumHigh = 0; + for (let i = start; i < end; i++) { + const x = audio[i]; + sumBroad += x * x; + let v: number = x; + for (let s = 0; s < HIGH_BAND_STAGES; s++) { + const y = a * (prevOut[s] + v - prevIn[s]); + prevIn[s] = v; + prevOut[s] = y; + v = y; + } + sumHigh += v * v; + } + const count = Math.max(1, end - start); + broad[f] = Math.sqrt(sumBroad / count); + high[f] = Math.sqrt(sumHigh / count); + } + + // Normalise each band against its own loud level, so a quiet fricative can + // still register as strongly as a vowel. + const normBroad = percentile(broad, 0.95) || 1; + const normHigh = percentile(high, 0.95) || 1; + const combined = new Float32Array(n); + for (let f = 0; f < n; f++) { + combined[f] = Math.max(broad[f] / normBroad, high[f] / normHigh); + } + + // ~25 ms smoothing, so single-frame noise does not read as an onset. + const level = new Float32Array(n); + for (let f = 0; f < n; f++) { + let sum = 0; + let count = 0; + for (let k = Math.max(0, f - 2); k <= Math.min(n - 1, f + 2); k++) { + sum += combined[k]; + count++; + } + level[f] = sum / count; + } + + const rise = new Float32Array(n); + let peak = 0; + for (let f = 1; f < n; f++) { + const d = level[f] - level[f - 1]; + rise[f] = d > 0 ? d : 0; + if (rise[f] > peak) peak = rise[f]; + } + if (peak > 0) for (let f = 0; f < n; f++) rise[f] /= peak; + + return { hopS, level, rise, floor: percentile(level, 0.1) }; +} + +/** + * The global shift (seconds) that best lands word starts on loudness rises. + * + * Unlike the VAD-mask score this needs no pauses at all: it votes with every word + * in the transcript, so continuous speech — where Whisper emits words back to back + * and the mask score goes flat — still yields a sharp estimate. Positive means the + * transcript is late. Ties resolve toward the smallest correction. + */ +export function estimateLagFromEnvelope( + words: Word[], + env: SpeechEnvelope, + { maxLagS = 0.6, lagStepS = 0.01 }: AlignOptions = {} +): number { + const { hopS, rise } = env; + if (words.length === 0 || rise.length === 0) return 0; + // Tolerate a frame or two of jitter between a word boundary and its onset. + const slack = Math.max(1, Math.round(0.01 / hopS)); + const scoreAt = (t: number) => { + const c = Math.round(t / hopS); + let best = 0; + for (let i = c - slack; i <= c + slack; i++) { + if (i >= 0 && i < rise.length && rise[i] > best) best = rise[i]; + } + return best; + }; + let bestLag = 0; + let bestScore = -1; + for (const lag of lagCandidates(maxLagS, lagStepS)) { + let score = 0; + for (const w of words) score += scoreAt(w.start - lag); + if (score > bestScore) { + bestScore = score; + bestLag = lag; + } + } + return bestScore > 0 ? bestLag : 0; } /** Times (seconds) where speech runs begin and end in the VAD flags. */ @@ -239,6 +393,81 @@ export function estimateSpeechLag( ).bestLag; } +/** How far back an onset may be dragged to find the true start of the sound. */ +const MAX_ONSET_REFINE_S = 0.35; + +/** + * Refined onsets closer together than this describe one sound, not two. + * + * Silero often splits a fricative-initial word into two runs — on the test clip + * "shot" gives one run for the /ʃ/ and another for the vowel 0.08 s later. Both + * refine to real onsets, and anchor matching then picks the *later* one, putting + * the word on its vowel again. Keeping only the earliest of a cluster fixes that. + * Well below the ~0.2 s spacing of genuinely distinct syllables. + */ +const MIN_ONSET_SEPARATION_S = 0.15; + +/** + * Pull each VAD onset back to where the sound actually starts. + * + * Silero fires on voicing, so a word beginning with an unvoiced fricative gets an + * onset on the vowel: on the test clip "shot" was flagged at 8.128 s when its /ʃ/ + * begins at 7.92 s. Walking back through the envelope recovers the missing 0.2 s. + * + * The walk stops at the first real gap in the sound. If it reaches the look-back + * limit without finding one there is no evidence of a distinct onset, so the VAD + * time is kept — without that check two onsets on the test clip ran the full + * look-back and annexed the previous word. + * + * An onset landing within `MIN_ONSET_SEPARATION_S` of the previous one describes + * the same sound and is dropped. That covers Silero splitting a fricative-initial + * word into two runs, and doubles as the monotonicity guard. The returned list is + * therefore usually shorter than `onsets`. + */ +export function refineOnsets( + onsets: number[], + env: SpeechEnvelope, + maxLookBackS = MAX_ONSET_REFINE_S +): number[] { + const { hopS, level, floor } = env; + if (level.length === 0) return onsets.slice(); + const maxGapFrames = Math.round(0.03 / hopS); + const out: number[] = []; + let prev = -Infinity; + + for (const onset of onsets) { + const at = Math.round(onset / hopS); + const lo = Math.max(0, at - Math.round(maxLookBackS / hopS)); + // Loudness just after the onset tells us what "present" means here. + let loud = 0; + for (let i = at; i < Math.min(level.length, at + Math.round(0.08 / hopS)); i++) { + if (level[i] > loud) loud = level[i]; + } + const threshold = Math.max(floor * 2.5, floor + 0.08 * (loud - floor)); + + let k = at; + let gap = 0; + let foundGap = false; + while (k > lo) { + if (level[k - 1] >= threshold) { + gap = 0; + } else if (++gap > maxGapFrames) { + foundGap = true; + break; + } + k--; + } + while (k < at && level[k] < threshold) k++; + + const refined = foundGap ? k * hopS : onset; + // Landing on the previous anchor means this is the same sound, not a new one. + if (refined <= prev + MIN_ONSET_SEPARATION_S) continue; + out.push(refined); + prev = refined; + } + return out; +} + /** * One matched landmark: the decoded time of a word start, and the VAD onset it * belongs to. `from - to` is the correction to apply there. @@ -273,7 +502,7 @@ export function buildSpeechAnchors( options: AlignOptions = {} ): SpeechAnchor[] { const { minGapS = 0.06, landmarkTolS = 0.3 } = options; - const { onsets } = speechEdgesFromFrames(speechFrames, options); + const onsets = options.onsets ?? speechEdgesFromFrames(speechFrames, options).onsets; if (onsets.length === 0) return []; const anchors: SpeechAnchor[] = []; @@ -368,7 +597,7 @@ export function snapWordsToSpeech( const out = words.map((w) => ({ ...w })); if (out.length === 0) return out; - const { onsets } = speechEdgesFromFrames(speechFrames, options); + const onsets = options.onsets ?? speechEdgesFromFrames(speechFrames, options).onsets; for (let i = 0; i < out.length; i++) { const w = out[i]; const prevEnd = i > 0 ? out[i - 1].end : 0; @@ -396,25 +625,83 @@ function normalizeWords(words: Word[], minWordS: number, duration: number): Word } /** - * Fixed extra lead applied after alignment, in seconds. + * A multi-syllable word cannot really last this long or less; if Whisper says so, + * its DTW collapsed the word and gave the time to a neighbour. + */ +const SLIVER_MAX_S = 0.1; + +/** Rough syllable count: vowel groups, at least one. */ +function syllableCount(text: string): number { + const groups = text.toLowerCase().replace(/[^a-z]/g, "").match(/[aeiouy]+/g); + return Math.max(1, groups ? groups.length : 1); +} + +/** + * Re-split words that Whisper's DTW collapsed to a sliver. + * + * Distinct from everything else here: this is not a timing offset, it is time + * mis-distributed *between* adjacent words. On the test clip "evening" came back + * spanning 0.06 s while the "and" before it took 0.40 s and 0.34 s went to no word + * at all, even though both ends of that region were correctly placed. Shifting + * cannot fix that; only re-splitting can. * - * Alignment can only ever be as good as its reference, and the reference is a - * hair late in both directions: Silero raises its speech flag ~40 ms after speech - * actually starts, and a highlight that lands a touch early reads as in sync - * whereas one that lands late reads as lagging. This is a deliberate perceptual - * nudge, not a measured correction — against acoustic ground truth it slightly - * increases absolute error, so keep it small and change it here. + * The region runs from the over-long neighbour's start to the next word's start, + * trimmed back to where the audio actually goes quiet, and is divided by syllable + * count. That is a guess at relative duration — a real forced aligner would know — + * so it only fires on spans too short to be anything but broken. */ -export const ALIGN_LEAD_S = 0.08; +export function repairCollapsedWords( + words: Word[], + env: SpeechEnvelope | null, + options: AlignOptions = {} +): Word[] { + const { minWordS = 0.02, duration = 0 } = options; + const out = words.map((w) => ({ ...w })); + + for (let i = 0; i < out.length; i++) { + const word = out[i]; + if (word.end - word.start > SLIVER_MAX_S) continue; + if (syllableCount(word.text) < 2) continue; // "a", "I", "the" really are brief + + const next = out[i + 1]; + let regionEnd = next ? next.start : word.end; + // Do not reclaim time the audio says is silent. + if (env && env.level.length > 0) { + const threshold = env.floor * 2.5; + let k = Math.min(env.level.length - 1, Math.round(regionEnd / env.hopS)); + const floorFrame = Math.round(word.end / env.hopS); + while (k > floorFrame && env.level[k] < threshold) k--; + regionEnd = Math.max(word.end, k * env.hopS); + } + + // Pull in the previous word only when it is touching and looks over-long. + const prev = i > 0 ? out[i - 1] : undefined; + const joined = prev && word.start - prev.end < 1e-6; + const group = joined ? [prev, word] : [word]; + const regionStart = group[0].start; + const room = regionEnd - regionStart; + const held = group.reduce((s, w) => s + (w.end - w.start), 0); + if (room <= held + minWordS) continue; // nothing unassigned to reclaim + + const weights = group.map((w) => syllableCount(w.text)); + const total = weights.reduce((s, w) => s + w, 0); + let cursor = regionStart; + group.forEach((w, k) => { + w.start = cursor; + cursor = k === group.length - 1 ? regionEnd : cursor + (room * weights[k]) / total; + w.end = cursor; + }); + } + return normalizeWords(out, minWordS, duration); +} /** - * Correct Whisper's late word timestamps against the VAD flags. + * Correct Whisper's late word timestamps against the audio. * * Interpolates the correction between matched pause anchors so the decay across * the clip is tracked rather than averaged away, then snaps remaining starts onto - * nearby onsets, then applies `leadS`. Falls back to a single global shift when - * there are too few anchors to describe a curve. Returns new Word objects; the - * input is untouched. + * nearby onsets. Falls back to a single global shift when there are too few + * anchors to describe a curve. Returns new Word objects; the input is untouched. */ export function alignWordsToSpeech( words: Word[], @@ -422,21 +709,26 @@ export function alignWordsToSpeech( options: AlignOptions = {} ): Word[] { if (words.length === 0) return []; - const { leadS = ALIGN_LEAD_S, minWordS = 0.02, duration = 0 } = options; - const lag = estimateSpeechLag(words, speechFrames, options); - const anchors = buildSpeechAnchors(words, speechFrames, lag, options); + const { audio, sampleRate = VAD_SAMPLE_RATE } = options; + const edges = speechEdgesFromFrames(speechFrames, options); + // With audio we can do two things the VAD flags alone cannot: estimate the + // shift from loudness rises (the VAD-mask score goes flat on speech that never + // pauses, and silently returns ~0), and pull each onset back off the vowel onto + // the fricative that actually starts the word. + const env = audio ? speechEnvelope(audio, sampleRate) : null; + const resolved: AlignOptions = env + ? { ...options, onsets: refineOnsets(edges.onsets, env) } + : options; + const lag = env + ? estimateLagFromEnvelope(words, env, options) + : estimateSpeechLag(words, speechFrames, options); + const anchors = buildSpeechAnchors(words, speechFrames, lag, resolved); const usable = anchors.length >= MIN_ANCHORS ? anchors : []; const warped = words.map((w) => ({ ...w, start: w.start - correctionAt(usable, lag, w.start), end: w.end - correctionAt(usable, lag, w.end), })); - const snapped = snapWordsToSpeech(warped, speechFrames, options); - if (leadS === 0) return snapped; - // After snapping, so the snap cannot pull the lead back onto the VAD onset. - for (const w of snapped) { - w.start -= leadS; - w.end -= leadS; - } - return normalizeWords(snapped, minWordS, duration); + const snapped = snapWordsToSpeech(warped, speechFrames, resolved); + return repairCollapsedWords(snapped, env, options); } diff --git a/lib/forcedAlign.ts b/lib/forcedAlign.ts new file mode 100644 index 0000000..c238bbf --- /dev/null +++ b/lib/forcedAlign.ts @@ -0,0 +1,345 @@ +/** + * CTC forced alignment. + * + * Whisper's word timestamps come from DTW over the decoder's cross-attention, + * which is a by-product of decoding rather than a measurement: it runs late, the + * error drifts across a clip, it puts fricative-initial words on their vowel, and + * it sometimes collapses a word to a few milliseconds and hands its time to a + * neighbour. `lib/align.ts` corrects each of those after the fact. + * + * This module removes the need for that. Given the transcript text and a CTC + * acoustic model's per-frame character probabilities, Viterbi finds the single + * most likely assignment of frames to characters. Word boundaries then fall out + * of the path directly — measured against the audio rather than inferred from + * attention, and correct by construction. + * + * Everything here is pure. The caller supplies the emission matrix, so the + * algorithm is testable without loading a model. + */ +import type { Word } from "./types"; +import type { SpeechEnvelope } from "./align"; + +/** Per-frame log-probabilities from a CTC model, row-major `[frames][vocab]`. */ +export interface CtcEmission { + logProbs: Float32Array; + frames: number; + vocab: number; +} + +/** The pieces of a CTC tokenizer this module needs. */ +export interface CtcVocab { + /** Id of the CTC blank symbol (`` for wav2vec2). */ + blankId: number; + /** Encode already-normalised text to token ids. */ + encode(text: string): number[]; + /** Id of the between-words symbol (`|`), if the vocab has one. */ + delimiterId?: number; +} + +/** + * Keep only what a character-level CTC vocabulary can represent. Digits and + * punctuation have no acoustic spelling here, so a word made only of those + * cannot be aligned and is interpolated from its neighbours instead. + */ +export function normalizeForCtc(text: string): string { + return text.toUpperCase().replace(/[^A-Z']/g, ""); +} + +/** Which slice of the flattened token sequence belongs to which word. */ +interface TokenSpan { + /** Index into the caller's word array. */ + word: number; + from: number; + to: number; +} + +interface TokenSequence { + tokens: number[]; + spans: TokenSpan[]; +} + +/** Flatten words into one token sequence, separated by the delimiter symbol. */ +function buildTokens(words: Word[], vocab: CtcVocab): TokenSequence { + const tokens: number[] = []; + const spans: TokenSpan[] = []; + words.forEach((word, index) => { + const ids = vocab.encode(normalizeForCtc(word.text)); + if (ids.length === 0) return; + if (tokens.length > 0 && vocab.delimiterId !== undefined) { + tokens.push(vocab.delimiterId); + } + const from = tokens.length; + tokens.push(...ids); + spans.push({ word: index, from, to: tokens.length }); + }); + return { tokens, spans }; +} + +/** + * Viterbi over the standard CTC state lattice. + * + * The token sequence is expanded to `blank t0 blank t1 … blank` so a blank may + * sit between any two tokens and *must* sit between two identical ones — which is + * what makes doubled letters ("HELLO") align correctly. From each state a frame + * may stay, advance one, or skip a blank to the next token when that token + * differs from the previous one. + * + * Returns the state index chosen at each frame, or null when the audio is too + * short to host every token. + */ +export function ctcViterbi( + emission: CtcEmission, + tokens: number[], + blankId: number +): Int32Array | null { + const { logProbs, frames, vocab } = emission; + const n = tokens.length; + if (n === 0 || frames === 0) return null; + const states = 2 * n + 1; + // CTC needs a frame per token, plus a blank between any repeated pair. + let minimum = n; + for (let i = 1; i < n; i++) if (tokens[i] === tokens[i - 1]) minimum++; + if (frames < minimum) return null; + + const labelAt = (s: number) => (s % 2 === 1 ? tokens[(s - 1) / 2] : blankId); + const NEG = -1e30; + const score = new Float32Array(states).fill(NEG); + const next = new Float32Array(states); + // One backpointer per state per frame: 0 stay, 1 advance, 2 skip a blank. + const back = new Uint8Array(frames * states); + + const lp = (t: number, v: number) => logProbs[t * vocab + v]; + score[0] = lp(0, blankId); + if (states > 1) score[1] = lp(0, tokens[0]); + + for (let t = 1; t < frames; t++) { + const row = t * states; + for (let s = 0; s < states; s++) { + let best = score[s]; + let choice = 0; + if (s >= 1 && score[s - 1] > best) { + best = score[s - 1]; + choice = 1; + } + // Skipping the blank at s-1 is only legal between two different tokens. + if (s >= 2 && s % 2 === 1 && tokens[(s - 1) / 2] !== tokens[(s - 3) / 2] && score[s - 2] > best) { + best = score[s - 2]; + choice = 2; + } + next[s] = best <= NEG ? NEG : best + lp(t, labelAt(s)); + back[row + s] = choice; + } + score.set(next); + } + + // A valid path ends on the last token or the trailing blank. + let state = score[states - 1] >= score[states - 2] ? states - 1 : states - 2; + if (score[state] <= NEG) return null; + + const path = new Int32Array(frames); + for (let t = frames - 1; t >= 0; t--) { + path[t] = state; + state -= back[t * states + state]; + } + return path; +} + +/** + * Word timings from a Viterbi path. + * + * A word owns the frames whose state carries one of its tokens. Words with no + * alignable characters keep a slot so the caller can interpolate them. + */ +function spansFromPath( + path: Int32Array, + spans: TokenSpan[], + frameS: number, + offsetS: number +): Map { + const first = new Int32Array(spans.length).fill(-1); + const last = new Int32Array(spans.length).fill(-1); + // token index -> span index + const owner = new Map(); + spans.forEach((span, i) => { + for (let k = span.from; k < span.to; k++) owner.set(k, i); + }); + + for (let t = 0; t < path.length; t++) { + const state = path[t]; + if (state % 2 === 0) continue; // blank + const tokenIndex = (state - 1) / 2; + const i = owner.get(tokenIndex); + if (i === undefined) continue; + if (first[i] < 0) first[i] = t; + last[i] = t; + } + + const out = new Map(); + spans.forEach((span, i) => { + if (first[i] < 0) return; + out.set(span.word, { + start: offsetS + first[i] * frameS, + end: offsetS + (last[i] + 1) * frameS, + }); + }); + return out; +} + +/** + * Align one batch of words against one slice of audio. + * + * `frameS` is derived from the emission length rather than assumed, so the + * model's exact stride and any convolution edge effects are absorbed. Words the + * vocabulary cannot spell are spread evenly through the gap their neighbours + * leave. Returns null when the path cannot be found, so the caller can keep the + * timings it already had. + */ +export function alignBatch( + words: Word[], + emission: CtcEmission, + sliceStartS: number, + sliceDurationS: number, + vocab: CtcVocab +): Word[] | null { + if (words.length === 0) return []; + const { tokens, spans } = buildTokens(words, vocab); + if (spans.length === 0) return null; + + const path = ctcViterbi(emission, tokens, vocab.blankId); + if (!path) return null; + + const frameS = sliceDurationS / emission.frames; + const timings = spansFromPath(path, spans, frameS, sliceStartS); + if (timings.size === 0) return null; + + const out = words.map((w) => ({ ...w })); + for (const [index, span] of timings) { + out[index].start = span.start; + out[index].end = span.end; + } + + // Unalignable words (digits, bare punctuation) sit between aligned ones. + for (let i = 0; i < out.length; i++) { + if (timings.has(i)) continue; + let j = i; + while (j < out.length && !timings.has(j)) j++; + const from = i > 0 ? out[i - 1].end : sliceStartS; + const to = j < out.length ? out[j].start : sliceStartS + sliceDurationS; + const step = (to - from) / (j - i); + for (let k = i; k < j; k++) { + out[k].start = from + step * (k - i); + out[k].end = from + step * (k - i + 1); + } + i = j - 1; + } + return out; +} + +/** + * Widen CTC word spans out to the sound they cover. + * + * CTC emissions are peaky: the model fires a spike for each character and stays + * blank around it, so the Viterbi path hands the frames at a word's edges to + * blank and the span comes back too narrow — on the test clip "shot" started + * 0.13 s and "evening" 0.16 s late. Growing each boundary through contiguous + * sound recovers them. + * + * A boundary only ever grows into audio the envelope calls present, and never + * past its neighbour, so a real pause between words is preserved and no word can + * annexe another's audio. CTC has already decided *which* word owns a region; + * this only sharpens where the edge falls. + */ +export function expandToAcoustics( + words: Word[], + env: SpeechEnvelope, + maxGrowS = 0.3 +): Word[] { + const { hopS, level, floor } = env; + const out = words.map((w) => ({ ...w })); + if (level.length === 0) return out; + const lastFrame = level.length - 1; + const clampFrame = (f: number) => Math.min(lastFrame, Math.max(0, f)); + + // What counts as present is set by how loud each word actually is. + const thresholds = out.map((word) => { + const at = clampFrame(Math.round(word.start / hopS)); + const endAt = clampFrame(Math.round(word.end / hopS)); + let loud = 0; + for (let f = at; f <= endAt; f++) if (level[f] > loud) loud = level[f]; + return Math.max(floor * 2.5, floor + 0.08 * (loud - floor)); + }); + + // Each side grows independently, against the neighbour's *decoded* boundary + // rather than its grown one. Doing it sequentially would let whichever word is + // processed first swallow the whole transition into the next. + const decodedStart = out.map((w) => w.start); + const decodedEnd = out.map((w) => w.end); + const starts = out.map((word, i) => { + const at = clampFrame(Math.round(word.start / hopS)); + const limit = clampFrame( + Math.round(Math.max(i > 0 ? decodedEnd[i - 1] : 0, word.start - maxGrowS) / hopS) + ); + let start = at; + while (start > limit && level[start - 1] >= thresholds[i]) start--; + return start * hopS; + }); + const ends = out.map((word, i) => { + const at = clampFrame(Math.round(word.end / hopS)); + const limit = clampFrame( + Math.round( + Math.min( + i + 1 < out.length ? decodedStart[i + 1] : Infinity, + word.end + maxGrowS + ) / hopS + ) + ); + let end = at; + while (end < limit && level[end] >= thresholds[i]) end++; + return end * hopS; + }); + + // Where both neighbours reached into the same sound there is no acoustic + // evidence for the boundary — split the difference instead of picking a side. + for (let i = 0; i + 1 < out.length; i++) { + if (ends[i] > starts[i + 1]) { + const middle = (ends[i] + starts[i + 1]) / 2; + ends[i] = middle; + starts[i + 1] = middle; + } + } + + for (let i = 0; i < out.length; i++) { + out[i].start = starts[i]; + out[i].end = Math.max(ends[i], starts[i] + hopS); + } + return out; +} + +/** + * Split words into batches small enough to align in one pass. + * + * The Viterbi lattice is frames × tokens, so an unbroken hour would need + * gigabytes. Batches are cut at the widest pause available, falling back to a + * hard cut when speech runs longer than `maxSpanS` without one. + */ +export function groupWordsForAlignment( + words: Word[], + maxSpanS = 20, + minGapS = 0.25 +): Word[][] { + const batches: Word[][] = []; + let current: Word[] = []; + for (const word of words) { + if (current.length > 0) { + const span = word.end - current[0].start; + const gap = word.start - current[current.length - 1].end; + if (span > maxSpanS && (gap >= minGapS || span > maxSpanS * 1.5)) { + batches.push(current); + current = []; + } + } + current.push(word); + } + if (current.length > 0) batches.push(current); + return batches; +} diff --git a/tests/align-test.ts b/tests/align-test.ts index f525d45..518aa77 100644 --- a/tests/align-test.ts +++ b/tests/align-test.ts @@ -1,11 +1,14 @@ import { - ALIGN_LEAD_S, alignWordsToSpeech, buildSpeechAnchors, correctionAt, + estimateLagFromEnvelope, estimateSpeechLag, + speechEnvelope, snapWordsToSpeech, speechEdgesFromFrames, + refineOnsets, + repairCollapsedWords, } from "../lib/align"; import { VAD_FRAME_SIZE, VAD_SAMPLE_RATE } from "../lib/vad"; import type { Word } from "../lib/types"; @@ -80,7 +83,7 @@ function word(id: number, start: number, end: number): Word { `expected lag ~${LATE}, got ${lag}` ); - const fixed = alignWordsToSpeech(late, frames, { duration: 22, leadS: 0 }); + const fixed = alignWordsToSpeech(late, frames, { duration: 22 }); const errBefore = late.reduce((s, w, i) => s + Math.abs(w.start - truth[i]!.start), 0) / late.length; const errAfter = @@ -109,7 +112,7 @@ function word(id: number, start: number, end: number): Word { ["all silence", framesFor([], 4)], ] as const) { assert(estimateSpeechLag(words, frames) === 0, `${label} frames should give lag 0`); - const out = alignWordsToSpeech(words, frames, { duration: 4, leadS: 0 }); + const out = alignWordsToSpeech(words, frames, { duration: 4 }); assert( out.every((w, i) => w.start === words[i]!.start && w.end === words[i]!.end), `${label} frames should leave timings untouched` @@ -152,7 +155,7 @@ function word(id: number, start: number, end: number): Word { const late = correctionAt(anchors, 0, 20.5); assert(early > late + 0.1, `correction should decay: ${early} -> ${late}`); - const fixed = alignWordsToSpeech(drifted, frames, { duration: 23, leadS: 0 }); + const fixed = alignWordsToSpeech(drifted, frames, { duration: 23 }); const mae = (ws: Word[]) => ws.reduce((s, w, i) => s + Math.abs(w.start - truth[i]!.start), 0) / ws.length; const flat = snapWordsToSpeech( @@ -175,46 +178,6 @@ function word(id: number, start: number, end: number): Word { ); } -{ - // The perceptual lead: a flat extra shift on top of alignment. It must survive - // snapping (which would otherwise pull starts back onto the VAD onset) and must - // not break the ordering / length / range invariants. - const frames = framesFor( - [ - [1, 2], - [2.5, 4], - ], - 5 - ); - const words = [word(0, 1.2, 1.7), word(1, 1.7, 2.1), word(2, 2.7, 3.9)]; - const plain = alignWordsToSpeech(words, frames, { duration: 5, leadS: 0 }); - const led = alignWordsToSpeech(words, frames, { duration: 5, leadS: ALIGN_LEAD_S }); - for (let i = 0; i < plain.length; i++) { - assert( - Math.abs(plain[i]!.start - ALIGN_LEAD_S - led[i]!.start) < 1e-9, - `word ${i} should lead by exactly ${ALIGN_LEAD_S}: ${plain[i]!.start} vs ${led[i]!.start}` - ); - assert( - Math.abs((led[i]!.end - led[i]!.start) - (plain[i]!.end - plain[i]!.start)) < 1e-9, - `word ${i} duration should be unchanged by the lead` - ); - } - for (let i = 0; i < led.length; i++) { - assert(led[i]!.end > led[i]!.start, `word ${i} positive length`); - assert(led[i]!.start >= 0, `word ${i} start within range`); - if (i > 0) assert(led[i]!.start >= led[i - 1]!.start, `order preserved at ${i}`); - } - assert(ALIGN_LEAD_S > 0 && ALIGN_LEAD_S < 0.2, "lead should stay small"); - - // A word already at t=0 cannot lead further, and must not go negative. - const atZero = alignWordsToSpeech([word(0, 0.02, 0.5)], framesFor([[0, 1]], 2), { - duration: 2, - }); - assert(atZero[0]!.start === 0, `start clamped to 0, got ${atZero[0]!.start}`); - assert(atZero[0]!.end > 0, "end still positive at the clip start"); - console.log(`lead of ${ALIGN_LEAD_S}s applied after snapping: ok`); -} - { // Snapping must not reorder words or let one swallow its neighbour's audio. const frames = framesFor( @@ -248,4 +211,188 @@ function word(id: number, start: number, end: number): Word { console.log("clamping: ok", out.map((w) => [w.start, w.end])); } +{ + // Continuous speech: Whisper emits words back to back, so there are no pause + // landmarks and the VAD-mask score goes flat (a solid block of words inside a + // solid block of speech trades one edge mismatch for the other). The loudness + // envelope still votes with every word, so it stays sharp. + const SR = VAD_SAMPLE_RATE; + const TOTAL = 30; + const truth: Word[] = []; + for (let t = 1; t + 0.35 <= 29; t += 0.35) { + truth.push(word(truth.length, t, t + 0.35)); + } + // A decaying burst per word, so every onset sits at a known time. + const audio = new Float32Array(Math.round(TOTAL * SR)); + for (const w of truth) { + const s = Math.round(w.start * SR); + const e = Math.min(audio.length, Math.round(w.end * SR)); + for (let i = s; i < e; i++) { + audio[i] += Math.sin(i / 12) * 0.35 * Math.exp((-(i - s) / SR) * 6); + } + } + const frames = framesFor([[1, 29]], TOTAL); + + for (const LAG of [0.1, 0.25, 0.34]) { + const late = truth.map((w) => ({ ...w, start: w.start + LAG, end: w.end + LAG })); + const fromEnv = estimateLagFromEnvelope(late, speechEnvelope(audio, SR)); + assert( + Math.abs(fromEnv - LAG) <= 0.02, + `envelope should recover lag ${LAG}, got ${fromEnv}` + ); + // The VAD-only path is expected to miss this badly — that is why audio is passed. + const vadOnly = alignWordsToSpeech(late, frames, { duration: TOTAL }); + const withAudio = alignWordsToSpeech(late, frames, { + duration: TOTAL, + audio, + sampleRate: SR, + }); + const mae = (ws: Word[]) => + ws.reduce((s, w, i) => s + Math.abs(w.start - truth[i]!.start), 0) / ws.length; + assert( + mae(withAudio) < 0.03, + `audio-anchored alignment should be tight, got ${mae(withAudio)} for lag ${LAG}` + ); + assert( + mae(withAudio) < mae(vadOnly), + `audio should beat VAD-only on pauseless speech (${mae(withAudio)} vs ${mae(vadOnly)})` + ); + for (let i = 1; i < withAudio.length; i++) { + assert(withAudio[i]!.start >= withAudio[i - 1]!.start, `order broken at ${i}`); + assert(withAudio[i]!.end > withAudio[i]!.start, `zero-length word at ${i}`); + } + } + // Silence in, nothing out: no rises means no estimate rather than a wrong one. + assert( + estimateLagFromEnvelope(truth, speechEnvelope(new Float32Array(SR * 5), SR)) === 0, + "silent audio should yield no lag estimate" + ); + console.log("pauseless speech corrected via loudness envelope: ok"); +} + +{ + // Unvoiced fricatives: Silero fires on voicing, so a word like "shot" gets its + // VAD onset on the vowel and the /S/ in front of it is missed. Broadband RMS + // makes the same mistake — the fricative is much quieter than the vowel — so the + // envelope splits into a high band that can mark it. + const SR = VAD_SAMPLE_RATE; + const TOTAL = 3; + const audio = new Float32Array(Math.round(TOTAL * SR)); + const FRIC_START = 1.0, VOWEL_START = 1.22, VOWEL_END = 1.6; + // /S/: quiet, high frequency only. + for (let i = Math.round(FRIC_START * SR); i < Math.round(VOWEL_START * SR); i++) { + audio[i] = (Math.sin(i * 1.9) + Math.sin(i * 2.3)) * 0.03; + } + // vowel: loud, low frequency. + for (let i = Math.round(VOWEL_START * SR); i < Math.round(VOWEL_END * SR); i++) { + audio[i] = Math.sin(i / 40) * 0.3; + } + // VAD sees only the voiced part, exactly as Silero does. + const frames = framesFor([[VOWEL_START, VOWEL_END]], TOTAL); + const edges = speechEdgesFromFrames(frames); + const env = speechEnvelope(audio, SR); + const refined = refineOnsets(edges.onsets, env); + assert(refined.length >= 1, "refinement keeps at least one onset"); + assert( + Math.abs(refined[0]! - FRIC_START) < 0.05, + `onset should move back to the fricative at ${FRIC_START}, got ${refined[0]} (VAD said ${edges.onsets[0]})` + ); + + + // Silero often splits a fricative-initial word into two runs (one for the /S/, + // one for the vowel). Both refine to real onsets, and anchor matching would then + // pick the later one — putting the word back on its vowel. They must collapse. + const split = framesFor( + [ + [FRIC_START + 0.02, FRIC_START + 0.1], + [VOWEL_START, VOWEL_END], + ], + TOTAL + ); + const splitEdges = speechEdgesFromFrames(split); + assert(splitEdges.onsets.length === 2, "fixture should give two VAD runs"); + const merged = refineOnsets(splitEdges.onsets, env); + assert( + merged.length === 1, + `two runs of one sound should collapse to one onset, got ${merged.length}: ${merged}` + ); + assert( + Math.abs(merged[0]! - FRIC_START) < 0.05, + `the surviving onset should be the earliest, got ${merged[0]}` + ); + + // No gap inside the look-back means no evidence of a distinct onset, so the VAD + // time must be kept rather than dragged an arbitrary distance backwards. + const solid = new Float32Array(Math.round(TOTAL * SR)); + for (let i = 0; i < solid.length; i++) solid[i] = Math.sin(i / 40) * 0.3; + const solidEnv = speechEnvelope(solid, SR); + const mid = 1.5; + const kept = refineOnsets([mid], solidEnv); + assert( + kept.length === 1 && Math.abs(kept[0]! - mid) < 1e-9, + `continuous sound should keep its VAD onset, got ${kept}` + ); + + // Results stay ordered whatever the input. + const twoRuns = framesFor([[0.3, 0.7], [VOWEL_START, VOWEL_END]], TOTAL); + const twoEdges = speechEdgesFromFrames(twoRuns); + const bounded = refineOnsets(twoEdges.onsets, env); + for (let i = 1; i < bounded.length; i++) { + assert(bounded[i]! > bounded[i - 1]!, `refined onsets must stay ordered at ${i}`); + } + console.log( + `fricative onset recovered: ok (VAD ${edges.onsets[0]!.toFixed(3)} -> ${refined[0]!.toFixed(3)}, true ${FRIC_START})` + ); +} + +{ + // Time mis-distributed between neighbours, not shifted: Whisper gave "evening" + // 0.06 s while "and" took 0.40 s and 0.34 s went to no word at all, with both + // ends of the region correctly placed. Only re-splitting can fix that. + const words = [ + word(0, 20.0, 20.4), // "and" - over-long + word(1, 20.4, 20.46), // "evening" - collapsed to a sliver + word(2, 20.8, 21.2), // "bell" - correctly placed, must not move + ]; + words[0]!.text = "and"; + words[1]!.text = "evening"; + words[2]!.text = "bell"; + const fixed = repairCollapsedWords(words, null, { duration: 24 }); + assert( + fixed[1]!.end - fixed[1]!.start > 0.4, + `collapsed word should regain real duration, got ${(fixed[1]!.end - fixed[1]!.start).toFixed(3)}s` + ); + assert(fixed[2]!.start === 20.8, `the next word must not move, got ${fixed[2]!.start}`); + assert(fixed[0]!.start === 20.0, `the region start must not move, got ${fixed[0]!.start}`); + assert( + fixed[0]!.end < words[0]!.end, + "the over-long neighbour should give time back" + ); + for (let i = 1; i < fixed.length; i++) { + assert(fixed[i]!.start >= fixed[i - 1]!.start, `order broken at ${i}`); + assert(fixed[i]!.end > fixed[i]!.start, `zero-length word at ${i}`); + } + + // Single-syllable function words really can be brief - leave them alone. + const brief = [word(0, 1.0, 1.4), word(1, 1.4, 1.44), word(2, 2.0, 2.4)]; + brief[1]!.text = "a"; + const untouched = repairCollapsedWords(brief, null, { duration: 5 }); + assert( + untouched[1]!.start === 1.4 && untouched[1]!.end === 1.44, + "a one-syllable word must not be re-split" + ); + + // Nothing unassigned to reclaim -> no change. + const packed = [word(0, 1.0, 1.4), word(1, 1.4, 1.46), word(2, 1.46, 1.9)]; + packed[1]!.text = "evening"; + const same = repairCollapsedWords(packed, null, { duration: 5 }); + assert( + same[1]!.end - same[1]!.start < 0.1, + "with no spare room the word must be left as decoded" + ); + console.log( + `collapsed word re-split: ok (0.060s -> ${(fixed[1]!.end - fixed[1]!.start).toFixed(3)}s)` + ); +} + console.log("ALL ALIGN TESTS PASSED"); diff --git a/tests/forced-align-test.ts b/tests/forced-align-test.ts new file mode 100644 index 0000000..51774d4 --- /dev/null +++ b/tests/forced-align-test.ts @@ -0,0 +1,152 @@ +import { + alignBatch, + ctcViterbi, + groupWordsForAlignment, + normalizeForCtc, + type CtcEmission, + type CtcVocab, +} from "../lib/forcedAlign"; +import type { Word } from "../lib/types"; + +function assert(cond: boolean, msg: string) { + if (!cond) throw new Error(msg); +} + +// A toy character vocabulary: 0 = blank, 1 = "|", 2.. = A, B, C, ... +const LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ'"; +const VOCAB: CtcVocab = { + blankId: 0, + delimiterId: 1, + encode: (text) => + text + .split("") + .map((c) => LETTERS.indexOf(c)) + .filter((i) => i >= 0) + .map((i) => i + 2), +}; +const VOCAB_SIZE = LETTERS.length + 2; + +/** + * Emission where each frame is confidently one symbol. `plan` lists the symbol + * occupying each frame, so the correct alignment is known exactly. + */ +function emissionFor(plan: number[]): CtcEmission { + const frames = plan.length; + const logProbs = new Float32Array(frames * VOCAB_SIZE).fill(Math.log(0.001)); + for (let t = 0; t < frames; t++) logProbs[t * VOCAB_SIZE + plan[t]] = Math.log(0.9); + return { logProbs, frames, vocab: VOCAB_SIZE }; +} +const sym = (c: string) => (c === "|" ? 1 : c === "_" ? 0 : LETTERS.indexOf(c) + 2); +/** "_AB_|_C" -> frame plan, one character per frame. */ +const plan = (s: string) => s.split("").map(sym); + +function word(id: number, text: string, start: number, end: number): Word { + return { id, text, start, end, speaker: 0, deleted: false }; +} + +{ + assert(normalizeForCtc("Shot.") === "SHOT", "punctuation is stripped"); + assert(normalizeForCtc("don't") === "DON'T", "apostrophes survive"); + assert(normalizeForCtc("1985") === "", "digits have no spelling"); + console.log("normalisation: ok"); +} + +{ + // The path must follow the frames the symbols actually occupy. + const frames = plan("__CC_AA_TT__"); + const tokens = VOCAB.encode("CAT"); + const path = ctcViterbi(emissionFor(frames), tokens, 0); + assert(path !== null, "CAT should align"); + // state 2j+1 carries token j + const framesOf = (j: number) => { + const out: number[] = []; + path!.forEach((s, t) => { + if (s === 2 * j + 1) out.push(t); + }); + return out; + }; + assert(framesOf(0)[0] === 2, `C should start at frame 2, got ${framesOf(0)[0]}`); + assert(framesOf(1)[0] === 5, `A should start at frame 5, got ${framesOf(1)[0]}`); + assert(framesOf(2)[0] === 8, `T should start at frame 8, got ${framesOf(2)[0]}`); + console.log("viterbi finds the right frames: ok"); +} + +{ + // Doubled letters need a blank between them; without one the path is wrong. + const tokens = VOCAB.encode("BELL"); + const path = ctcViterbi(emissionFor(plan("_BB_EE_LL_L_")), tokens, 0); + assert(path !== null, "BELL should align"); + const states = new Set(Array.from(path!)); + // Both L tokens must be used (states 5 and 7), not collapsed into one. + assert(states.has(5) && states.has(7), `both L tokens should appear, got ${[...states]}`); + console.log("repeated letters keep their blank: ok"); +} + +{ + // Too few frames to host every token -> refuse rather than guess. + const short = ctcViterbi(emissionFor(plan("_A_")), VOCAB.encode("ALIGNMENT"), 0); + assert(short === null, "an impossible alignment must return null"); + assert(ctcViterbi(emissionFor(plan("___")), [], 0) === null, "no tokens -> null"); + console.log("impossible alignments refused: ok"); +} + +{ + // End to end: two words, deliberately wrong incoming times. + const frames = plan("__CC_AA_TT___|__SS_AA_TT__"); + const emission = emissionFor(frames); + const FRAME_S = 0.02; + const words = [word(0, "Cat", 5.0, 5.1), word(1, "sat.", 5.1, 5.16)]; + const aligned = alignBatch(words, emission, 10, frames.length * FRAME_S, VOCAB); + assert(aligned !== null, "batch should align"); + const [cat, sat] = aligned!; + assert( + Math.abs(cat.start - (10 + 2 * FRAME_S)) < 1e-6, + `cat should start at frame 2, got ${cat.start}` + ); + assert( + Math.abs(cat.end - (10 + 10 * FRAME_S)) < 1e-6, + `cat should end after frame 9, got ${cat.end}` + ); + assert( + Math.abs(sat.start - (10 + 16 * FRAME_S)) < 1e-6, + `sat should start at frame 16, got ${sat.start}` + ); + assert(sat.end > sat.start && cat.end <= sat.start, "words must not overlap"); + assert(aligned![0].text === "Cat" && aligned![1].text === "sat.", "text is preserved"); + console.log( + `end to end: ok (cat ${cat.start.toFixed(2)}-${cat.end.toFixed(2)}, sat ${sat.start.toFixed(2)}-${sat.end.toFixed(2)})` + ); +} + +{ + // A word the vocabulary cannot spell still gets a sensible slot. + const frames = plan("__CC_AA_TT___|__SS_AA_TT__"); + const words = [word(0, "Cat", 0, 0.1), word(1, "1985", 0.1, 0.2), word(2, "sat", 0.2, 0.3)]; + const aligned = alignBatch(words, emissionFor(frames), 0, frames.length * 0.02, VOCAB); + assert(aligned !== null, "batch with an unspellable word should still align"); + const [a, b, c] = aligned!; + assert(b.start >= a.end - 1e-9 && b.end <= c.start + 1e-9, "digits sit in the gap"); + assert(b.end > b.start, "digits keep a positive duration"); + console.log(`unspellable word interpolated: ok (${b.start.toFixed(2)}-${b.end.toFixed(2)})`); +} + +{ + const many: Word[] = []; + for (let i = 0; i < 60; i++) many.push(word(i, "w", i * 0.5, i * 0.5 + 0.4)); + const batches = groupWordsForAlignment(many, 10, 0.05); + assert(batches.length > 1, "long input should be split"); + assert( + batches.reduce((n, b) => n + b.length, 0) === many.length, + "batching must not drop words" + ); + for (const b of batches) { + assert(b.length > 0, "no empty batches"); + assert(b[b.length - 1]!.end - b[0]!.start <= 16, "batches stay near the limit"); + } + const flat = batches.flat().map((w) => w.id); + assert(flat.every((id, i) => id === i), "order is preserved across batches"); + assert(groupWordsForAlignment([], 10).length === 0, "no words -> no batches"); + console.log(`batching: ok (${many.length} words -> ${batches.length} batches)`); +} + +console.log("ALL FORCED ALIGN TESTS PASSED"); diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index ff76be0..1e686c9 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -15,6 +15,8 @@ import { AutoProcessor, AutoModel, AutoModelForAudioFrameClassification, + AutoModelForCTC, + AutoTokenizer, WhisperTextStreamer, Tensor, env, @@ -23,7 +25,14 @@ import { import type { Word, WorkerRequest, WorkerResponse } from "@/lib/types"; import { MODELS, type WhisperModel } from "@/lib/models"; import { cleanTranscript } from "@/lib/hallucinations"; -import { alignWordsToSpeech } from "@/lib/align"; +import { alignWordsToSpeech, speechEnvelope } from "@/lib/align"; +import { + alignBatch, + expandToAcoustics, + groupWordsForAlignment, + type CtcEmission, + type CtcVocab, +} from "@/lib/forcedAlign"; import { VAD_FRAME_SIZE, VAD_SAMPLE_RATE, @@ -41,6 +50,18 @@ if (env.backends?.onnx?.wasm) { const DIARIZATION_MODEL = "onnx-community/pyannote-segmentation-3.0"; const VAD_MODEL = "onnx-community/silero-vad"; +/** + * Character-level CTC model used to place word boundaries against the audio. + * Whisper's own timestamps are a by-product of decoding and only roughly right; + * this measures them. + */ +const ALIGN_MODEL = "Xenova/wav2vec2-base-960h"; +/** `|` in the wav2vec2 vocabulary — the between-words symbol. */ +const ALIGN_DELIMITER_ID = 4; +/** Viterbi needs a frames x tokens lattice, so alignment runs in bounded batches. */ +const ALIGN_BATCH_MAX_S = 20; +/** Context either side of a batch, so edge words are not clipped. */ +const ALIGN_BATCH_PAD_S = 0.2; /** Gaps longer than this split speech into separate Whisper jobs. */ const SPEECH_MAX_GAP_S = 1.5; /** Pad each speech region so phoneme edges are not clipped. */ @@ -241,6 +262,12 @@ interface DiarizationSegment { confidence: number; } +type Aligner = { + processor: Awaited>; + model: Awaited>; + vocab: CtcVocab; +}; + type Diarizer = { processor: Awaited>; model: Awaited>; @@ -421,6 +448,107 @@ function getDiarizer(): Promise { return diarizerPromise; } +/** + * Load the CTC acoustic model used to place word boundaries. Warmed in the + * background alongside the diarizer, so it is cached by the time Whisper is done. + * + * q4 is 86 MB against fp32's 360 MB and scored identically on the test clip + * (0.033 s mean boundary error). q8 is no smaller and markedly worse — it moved + * one word 0.22 s — and fp16 fails to load on onnxruntime-web. + */ +let alignerPromise: Promise | null = null; +function getAligner(): Promise { + if (!alignerPromise) { + alignerPromise = (async () => { + const [processor, model, tokenizer] = await Promise.all([ + AutoProcessor.from_pretrained(ALIGN_MODEL, {}), + AutoModelForCTC.from_pretrained(ALIGN_MODEL, { dtype: "q4" }), + AutoTokenizer.from_pretrained(ALIGN_MODEL), + ]); + /* eslint-disable-next-line @typescript-eslint/no-explicit-any */ + const tok = tokenizer as any; + const vocab: CtcVocab = { + blankId: tok.pad_token_id ?? 0, + delimiterId: ALIGN_DELIMITER_ID, + encode: (text) => + text ? (tok.encode(text, { add_special_tokens: false }) as number[]) : [], + }; + return { processor, model, vocab }; + })(); + alignerPromise.catch(() => { + alignerPromise = null; + }); + } + return alignerPromise; +} + +/** Per-frame log-probabilities for one slice of audio. */ +async function ctcEmission(aligner: Aligner, slice: Float32Array): Promise { + const inputs = await aligner.processor(slice); + /* eslint-disable-next-line @typescript-eslint/no-explicit-any */ + const { logits } = await (aligner.model as any)(inputs); + const [, frames, vocab] = logits.dims as number[]; + const data = logits.data as Float32Array; + // The model emits raw scores; Viterbi needs log-probabilities. + const logProbs = new Float32Array(frames * vocab); + for (let t = 0; t < frames; t++) { + const row = t * vocab; + let max = -Infinity; + for (let v = 0; v < vocab; v++) if (data[row + v] > max) max = data[row + v]; + let sum = 0; + for (let v = 0; v < vocab; v++) sum += Math.exp(data[row + v] - max); + const logSumExp = max + Math.log(sum); + for (let v = 0; v < vocab; v++) logProbs[row + v] = data[row + v] - logSumExp; + } + return { logProbs, frames, vocab }; +} + +/** + * Replace Whisper's word timings with boundaries measured against the audio. + * + * Whisper's own timestamps come from DTW over cross-attention and are only + * roughly right; they are used here just to decide which words go in which + * batch. Any batch that fails to align keeps the timings it came in with, so a + * bad slice costs accuracy on those words rather than losing them. + */ +async function forceAlign( + words: Word[], + audio: Float32Array, + duration: number +): Promise { + const aligner = await getAligner(); + const batches = groupWordsForAlignment(words, ALIGN_BATCH_MAX_S); + const out: Word[] = []; + let done = 0; + + for (const batch of batches) { + const from = Math.max(0, batch[0].start - ALIGN_BATCH_PAD_S); + const to = Math.min(duration, batch[batch.length - 1].end + ALIGN_BATCH_PAD_S); + const startSample = Math.floor(from * VAD_SAMPLE_RATE); + const endSample = Math.min(audio.length, Math.ceil(to * VAD_SAMPLE_RATE)); + const slice = audio.slice(startSample, endSample); + let aligned: Word[] | null = null; + if (slice.length > VAD_SAMPLE_RATE / 10) { + try { + const emission = await ctcEmission(aligner, slice); + aligned = alignBatch( + batch, + emission, + startSample / VAD_SAMPLE_RATE, + slice.length / VAD_SAMPLE_RATE, + aligner.vocab + ); + } catch (err) { + console.warn("Forced alignment failed for one batch; keeping decoded times.", err); + } + } + out.push(...(aligned ?? batch)); + done++; + post({ type: "progress", message: "Aligning words…", value: done / batches.length }); + } + return out; +} + async function diarize(audio: Float32Array): Promise { const { processor, model } = await getDiarizer(); const inputs = await processor(audio); @@ -474,6 +602,7 @@ self.onmessage = async (event: MessageEvent) => { // Overlap Whisper + Silero downloads; diarizer warms in the background. getDiarizer().catch(() => {}); + getAligner().catch(() => {}); const [asr, vad] = await Promise.all([getAsr(choice), getVad()]); let transcriber = asr; @@ -640,10 +769,26 @@ self.onmessage = async (event: MessageEvent) => { // phrases ("I'm sorry", "thanks for watching", …) that slip past decoding. const cleaned = cleanTranscript(rawWords); - // Whisper's DTW word timestamps run consistently late (~0.2 s on the test - // clips). Realign them against the VAD flags before diarization, so speakers - // are assigned from corrected times too. - const words = alignWordsToSpeech(cleaned, speechFrames, { duration }); + // Whisper's DTW word timestamps run late by an amount that drifts across the + // clip. Correct them against the VAD as a first pass: forced alignment only + // needs them accurate enough to batch by, and this is what we fall back to if + // the acoustic model is unavailable. + let words = alignWordsToSpeech(cleaned, speechFrames, { + duration, + audio, + sampleRate: VAD_SAMPLE_RATE, + }); + + // Then measure the boundaries properly. CTC forced alignment decides which + // audio each word occupies; the envelope sharpens the edges, which CTC leaves + // narrow because its emissions spike on each character and go blank around it. + try { + post({ type: "progress", message: "Aligning words…", value: 0 }); + const measured = await forceAlign(words, audio, duration); + words = expandToAcoustics(measured, speechEnvelope(audio, VAD_SAMPLE_RATE)); + } catch (err) { + console.warn("Forced alignment unavailable; using VAD-corrected times.", err); + } // Best-effort speaker diarization; a failure should not lose the transcript. try { From de89417e26d200b3c1f302db264ca143e8dd2bca Mon Sep 17 00:00:00 2001 From: Wassim Gharbi Date: Sun, 2 Aug 2026 14:08:05 -1000 Subject: [PATCH 2/7] Manage the alignment model with weightlift MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CTC aligner loaded itself with bare from_pretrained calls, so its 86 MB download happened behind a static "Aligning words…" with no bar — indistinguishable from a hang on a cold cache. Register it in the existing ModelManager instead, so it shares the download progress, cache labeling and unload-on-GPU-loss path with the ASR models. It needs the processor, model and tokenizer separately rather than a pipeline(), so it is a plain ModelDefinition wired to transformersProgress rather than transformersModel(). The aligner warms in the background while Whisper runs, so its bytes must not fight the "Transcribing…" line: speech models always own the progress line, and the aligner only claims it once transcription is done and we are actually waiting on it. Co-Authored-By: Claude Opus 5 (1M context) --- workers/transcription.worker.ts | 96 ++++++++++++++++++++++++--------- 1 file changed, 71 insertions(+), 25 deletions(-) diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index b8ee787..91b9b16 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -7,10 +7,11 @@ * original timeline. * 3. Pyannote segmentation 3.0 assigns a speaker to each word. * - * Both ASR backends are registered in a weightlift ModelManager so download - * progress, cache labeling, and WebGPU→WASM fallback share one path. Weights - * land in Cache Storage (Whisper) or IndexedDB (Parakeet); later runs are - * offline. ORT WASM is served same-origin from /vendor/ort* . + * Both ASR backends and the CTC aligner are registered in a weightlift + * ModelManager so download progress, cache labeling, and WebGPU→WASM fallback + * share one path. Weights land in Cache Storage (Whisper, aligner) or IndexedDB + * (Parakeet); later runs are offline. ORT WASM is served same-origin from + * /vendor/ort* . */ import { pipeline, @@ -27,7 +28,9 @@ import { import { ModelManager, type ModelDefinition } from "weightlift"; import { fallbackDevicePolicy, + isTransformersModelCached, transformersModel, + transformersProgress, } from "weightlift/transformers"; import type { Word, WorkerRequest, WorkerResponse } from "@/lib/types"; import { @@ -253,17 +256,41 @@ const models = new ModelManager({ }) ), }); +models.define(ALIGN_MODEL, alignerModel()); + +/** + * Whether the aligner currently owns the progress line. + * + * It warms in the background while Whisper transcribes, and its bytes must not + * fight the "Transcribing…" updates. Once transcription is done and we are + * actually waiting on it, its download is the thing worth showing — that wait + * used to sit behind a bare "Aligning words…" with no bar. + */ +let awaitingAligner = false; + +const MODEL_LABELS: Record = { + [ALIGN_MODEL]: { + cached: "Loading alignment model from cache…", + download: "Downloading alignment model…", + }, +}; +const ASR_LABELS = { + cached: "Loading speech model from cache…", + download: "Downloading speech model…", +}; + models.subscribe((snap) => { - const id = snap.loading[0]; + // Speech models always win the line; the aligner only gets it while awaited. + const id = + snap.loading.find((each) => each !== ALIGN_MODEL) ?? + (awaitingAligner ? snap.loading.find((each) => each === ALIGN_MODEL) : undefined); if (!id) return; const rec = snap.models[id]; if (!rec) return; + const labels = MODEL_LABELS[id] ?? ASR_LABELS; post({ type: "progress", - message: - rec.fromCache === true - ? "Loading speech model from cache…" - : "Downloading speech model…", + message: rec.fromCache === true ? labels.cached : labels.download, value: rec.indeterminate ? null : rec.percent, }); }); @@ -525,21 +552,32 @@ function getDiarizer(): Promise { } /** - * Load the CTC acoustic model used to place word boundaries. Warmed in the - * background alongside the diarizer, so it is cached by the time Whisper is done. + * The CTC acoustic model that places word boundaries, as a weightlift + * definition so its bytes are tracked like the ASR models rather than + * downloading silently behind an "Aligning words…" label. + * + * Not `transformersModel()`: that builds a `pipeline()`, and this needs the + * processor, model and tokenizer separately. The progress wiring is the same. * * q4 is 86 MB against fp32's 360 MB and scored identically on the test clip * (0.033 s mean boundary error). q8 is no smaller and markedly worse — it moved * one word 0.22 s — and fp16 fails to load on onnxruntime-web. */ -let alignerPromise: Promise | null = null; -function getAligner(): Promise { - if (!alignerPromise) { - alignerPromise = (async () => { +function alignerModel(): ModelDefinition { + return { + isCached: () => + isTransformersModelCached(ALIGN_MODEL, { + cacheKey: env.cacheKey ?? "transformers-cache", + }), + load: async ({ progress }) => { + const progress_callback = transformersProgress(progress); const [processor, model, tokenizer] = await Promise.all([ - AutoProcessor.from_pretrained(ALIGN_MODEL, {}), - AutoModelForCTC.from_pretrained(ALIGN_MODEL, { dtype: "q4" }), - AutoTokenizer.from_pretrained(ALIGN_MODEL), + AutoProcessor.from_pretrained(ALIGN_MODEL, { progress_callback }), + AutoModelForCTC.from_pretrained(ALIGN_MODEL, { + dtype: "q4", + progress_callback, + }), + AutoTokenizer.from_pretrained(ALIGN_MODEL, { progress_callback }), ]); /* eslint-disable-next-line @typescript-eslint/no-explicit-any */ const tok = tokenizer as any; @@ -550,12 +588,12 @@ function getAligner(): Promise { text ? (tok.encode(text, { add_special_tokens: false }) as number[]) : [], }; return { processor, model, vocab }; - })(); - alignerPromise.catch(() => { - alignerPromise = null; - }); - } - return alignerPromise; + }, + }; +} + +function getAligner(): Promise { + return models.load(ALIGN_MODEL); } /** Per-frame log-probabilities for one slice of audio. */ @@ -592,7 +630,15 @@ async function forceAlign( audio: Float32Array, duration: number ): Promise { - const aligner = await getAligner(); + awaitingAligner = true; + let aligner: Aligner; + try { + aligner = await getAligner(); + } finally { + awaitingAligner = false; + } + // Switch off the download label as soon as the weights are in hand. + post({ type: "progress", message: "Aligning words…", value: 0 }); const batches = groupWordsForAlignment(words, ALIGN_BATCH_MAX_S); const out: Word[] = []; let done = 0; From e82fafb5a85fc192f903febc2af81e6c0ffd0f0a Mon Sep 17 00:00:00 2001 From: Wassim Gharbi Date: Sun, 2 Aug 2026 16:25:30 -1000 Subject: [PATCH 3/7] Give the aligner its own ModelManager MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit wav2vec2 is not in the same family as the ASR models — it is an acoustic aligner, not a transcriber, and not something the user picks. But the deciding factor is lifecycle, not taxonomy. The ASR registry exists under a WebGPU→WASM fallback policy whose recovery step is unloadAll(). The aligner never asks for a device, so transformers.js runs it on WASM (DEFAULT_DEVICE), where a lost GPU cannot touch it. Sharing the registry meant a GPU loss threw away 86 MB of perfectly good weights and forced a needless session rebuild. Splitting also removes the id special-casing the shared registry forced into the progress subscription. The ASR handler goes back to "whatever is loading", and the aligner's own subscription is scoped to the await in forceAlign — the subscription's lifetime is the gate, so there is no module-level flag left to fall out of step if that path is ever re-entered. Co-Authored-By: Claude Opus 5 (1M context) --- workers/transcription.worker.ts | 71 ++++++++++++++++++--------------- 1 file changed, 38 insertions(+), 33 deletions(-) diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index 012cdc2..34739e2 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -257,45 +257,33 @@ const models = new ModelManager({ }) ), }); -models.define(ALIGN_MODEL, alignerModel()); - -/** - * Whether the aligner currently owns the progress line. - * - * It warms in the background while Whisper transcribes, and its bytes must not - * fight the "Transcribing…" updates. Once transcription is done and we are - * actually waiting on it, its download is the thing worth showing — that wait - * used to sit behind a bare "Aligning words…" with no bar. - */ -let awaitingAligner = false; - -const MODEL_LABELS: Record = { - [ALIGN_MODEL]: { - cached: "Loading alignment model from cache…", - download: "Downloading alignment model…", - }, -}; -const ASR_LABELS = { - cached: "Loading speech model from cache…", - download: "Downloading speech model…", -}; - models.subscribe((snap) => { - // Speech models always win the line; the aligner only gets it while awaited. - const id = - snap.loading.find((each) => each !== ALIGN_MODEL) ?? - (awaitingAligner ? snap.loading.find((each) => each === ALIGN_MODEL) : undefined); + const id = snap.loading[0]; if (!id) return; const rec = snap.models[id]; if (!rec) return; - const labels = MODEL_LABELS[id] ?? ASR_LABELS; post({ type: "progress", - message: rec.fromCache === true ? labels.cached : labels.download, + message: + rec.fromCache === true + ? "Loading speech model from cache…" + : "Downloading speech model…", value: rec.indeterminate ? null : rec.percent, }); }); +/** + * The aligner gets its own registry rather than joining `models`. + * + * It is a different family — an acoustic aligner, not a transcriber, and not + * something the user picks — but the deciding factor is lifecycle. The ASR + * registry exists under a WebGPU→WASM fallback policy whose recovery step is + * `unloadAll()`. The aligner never asks for a device, so transformers.js runs it + * on WASM (`DEFAULT_DEVICE`), where a lost GPU cannot touch it. Sharing the + * registry meant a GPU loss threw away 86 MB that was still perfectly good. + */ +const aligners = new ModelManager({ models: { [ALIGN_MODEL]: alignerModel() } }); + async function getAsr(choice: WhisperModel) { return models.load(MODELS[choice].id); } @@ -307,7 +295,8 @@ async function getParakeet() { /** * Drop dead WebGPU pipelines and reload on WASM. * A lost GPU device invalidates every WebGPU session, so clear the whole - * ASR cache — not just the model that was running. + * ASR cache — not just the model that was running. The aligner is a separate + * registry and runs on WASM, so it is deliberately untouched. */ async function fallbackAsrToWasm() { fallbackDevicePolicy.preferWasm(); @@ -594,7 +583,7 @@ function alignerModel(): ModelDefinition { } function getAligner(): Promise { - return models.load(ALIGN_MODEL); + return aligners.load(ALIGN_MODEL); } /** Per-frame log-probabilities for one slice of audio. */ @@ -631,12 +620,28 @@ async function forceAlign( audio: Float32Array, duration: number ): Promise { - awaitingAligner = true; + // Subscribe only while we are actually waiting. The aligner also warms in the + // background during transcription, and those bytes must not fight the + // "Transcribing…" line — but a wait here is worth a bar, since a cold cache + // otherwise looks like a hang. Scoping the subscription to the await is the + // whole gate: no shared flag to get out of step if this is ever re-entered. let aligner: Aligner; + const unsubscribe = aligners.subscribe((snap) => { + const rec = snap.models[ALIGN_MODEL]; + if (!rec || rec.status !== "loading") return; + post({ + type: "progress", + message: + rec.fromCache === true + ? "Loading alignment model from cache…" + : "Downloading alignment model…", + value: rec.indeterminate ? null : rec.percent, + }); + }); try { aligner = await getAligner(); } finally { - awaitingAligner = false; + unsubscribe(); } // Switch off the download label as soon as the weights are in hand. post({ type: "progress", message: "Aligning words…", value: 0 }); From b33dc3a85fb1ad7f6a7539245c263ff0a99e393b Mon Sep 17 00:00:00 2001 From: Wassim Gharbi Date: Sun, 2 Aug 2026 16:29:52 -1000 Subject: [PATCH 4/7] Prime the aligner progress line from current state subscribe() only fires on the next change, so a load that had started but not yet received its first byte left the UI on the previous message. Co-Authored-By: Claude Opus 5 (1M context) --- workers/transcription.worker.ts | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index 34739e2..2d136f5 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -626,9 +626,8 @@ async function forceAlign( // otherwise looks like a hang. Scoping the subscription to the await is the // whole gate: no shared flag to get out of step if this is ever re-entered. let aligner: Aligner; - const unsubscribe = aligners.subscribe((snap) => { - const rec = snap.models[ALIGN_MODEL]; - if (!rec || rec.status !== "loading") return; + const report = (rec: ReturnType) => { + if (rec.status !== "loading") return; post({ type: "progress", message: @@ -637,7 +636,12 @@ async function forceAlign( : "Downloading alignment model…", value: rec.indeterminate ? null : rec.percent, }); - }); + }; + const unsubscribe = aligners.subscribe((snap) => report(snap.models[ALIGN_MODEL])); + // subscribe() only fires on the next change, so prime from current state: a + // load that has started but not yet received its first byte would otherwise + // leave the UI sitting on "Transcribing…". + report(aligners.status(ALIGN_MODEL)); try { aligner = await getAligner(); } finally { From f7b6f26523e86e74837aea48bb54a68d46d46738 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 3 Aug 2026 03:18:14 +0000 Subject: [PATCH 5/7] Add Parakeet alignment and multilingual CTC model map MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Share the VAD → CTC → expand post-pass between Whisper and Parakeet, and pick a language-specific aligner (English wav2vec2, MMS for es/fr/de, XLS-R for zh) so non-English runs no longer warm an English-only model. Co-authored-by: Wassim Gharbi --- lib/alignModels.ts | 56 +++++++++ lib/forcedAlign.ts | 38 +++++- lib/types.ts | 5 +- tests/align-models-test.ts | 29 +++++ tests/forced-align-test.ts | 20 +++ workers/transcription.worker.ts | 214 ++++++++++++++++++++------------ 6 files changed, 281 insertions(+), 81 deletions(-) create mode 100644 lib/alignModels.ts create mode 100644 tests/align-models-test.ts diff --git a/lib/alignModels.ts b/lib/alignModels.ts new file mode 100644 index 0000000..b94e808 --- /dev/null +++ b/lib/alignModels.ts @@ -0,0 +1,56 @@ +import type { TranscriptLanguage } from "./languages"; +import type { CtcNormalizeMode } from "./forcedAlign"; + +/** CTC acoustic model used to force-align a transcript to audio. */ +export interface AlignModelInfo { + /** Hugging Face / transformers.js model id. */ + id: string; + /** How transcript text is folded into this model's character vocabulary. */ + normalize: CtcNormalizeMode; +} + +/** + * Language → CTC aligner map (WhisperX-style). + * + * English keeps the small Librispeech wav2vec2. European languages share the + * multilingual MMS forced-aligner (Latin a–z after accent folding). Chinese + * uses an XLS-R CTC model whose vocab includes Han characters. Languages with + * no entry skip CTC and keep the VAD/envelope heuristic only. + */ +export const ALIGN_MODELS: Partial< + Record +> = { + en: { + id: "Xenova/wav2vec2-base-960h", + normalize: "latin-upper", + }, + es: { + id: "onnx-community/mms-300m-1130-forced-aligner-ONNX", + normalize: "latin-lower", + }, + fr: { + id: "onnx-community/mms-300m-1130-forced-aligner-ONNX", + normalize: "latin-lower", + }, + de: { + id: "onnx-community/mms-300m-1130-forced-aligner-ONNX", + normalize: "latin-lower", + }, + zh: { + id: "onnx-community/wav2vec2-large-xlsr-53-chinese-zh-cn-ONNX", + normalize: "cjk", + }, +}; + +/** Align model for a transcript language, or null when CTC is unavailable. */ +export function alignModelFor( + language: TranscriptLanguage | string | undefined +): AlignModelInfo | null { + if (!language) return null; + return ALIGN_MODELS[language as TranscriptLanguage] ?? null; +} + +/** Unique Hugging Face ids referenced by {@link ALIGN_MODELS}. */ +export function uniqueAlignModelIds(): string[] { + return [...new Set(Object.values(ALIGN_MODELS).map((m) => m.id))]; +} diff --git a/lib/forcedAlign.ts b/lib/forcedAlign.ts index c238bbf..bab1a9a 100644 --- a/lib/forcedAlign.ts +++ b/lib/forcedAlign.ts @@ -26,23 +26,52 @@ export interface CtcEmission { vocab: number; } +/** + * How transcript text is folded into a CTC model's character vocabulary. + * + * - `latin-upper` — English wav2vec2 (A–Z, `'`) + * - `latin-lower` — MMS forced-aligner (a–z, `'`); accents are stripped + * - `cjk` — Chinese XLS-R (Han + A–Z, `'`) + */ +export type CtcNormalizeMode = "latin-upper" | "latin-lower" | "cjk"; + /** The pieces of a CTC tokenizer this module needs. */ export interface CtcVocab { - /** Id of the CTC blank symbol (`` for wav2vec2). */ + /** Id of the CTC blank symbol (`` / `` for wav2vec2). */ blankId: number; /** Encode already-normalised text to token ids. */ encode(text: string): number[]; /** Id of the between-words symbol (`|`), if the vocab has one. */ delimiterId?: number; + /** Text folding applied before `encode`. Default `latin-upper`. */ + normalizeMode?: CtcNormalizeMode; } /** * Keep only what a character-level CTC vocabulary can represent. Digits and * punctuation have no acoustic spelling here, so a word made only of those * cannot be aligned and is interpolated from its neighbours instead. + * + * Latin modes fold accents (`café` → `CAFE` / `cafe`) so European text can use + * ASCII-only models like MMS. CJK mode keeps Han characters. */ -export function normalizeForCtc(text: string): string { - return text.toUpperCase().replace(/[^A-Z']/g, ""); +export function normalizeForCtc( + text: string, + mode: CtcNormalizeMode = "latin-upper" +): string { + if (mode === "cjk") { + return Array.from(text) + .filter((ch) => { + const cp = ch.codePointAt(0) ?? 0; + const isHan = cp >= 0x4e00 && cp <= 0x9fff; + return isHan || /[A-Za-z']/.test(ch); + }) + .join("") + .toUpperCase(); + } + const folded = text.normalize("NFD").replace(/\p{M}/gu, ""); + const letters = folded.replace(/[^A-Za-z']/g, ""); + return mode === "latin-lower" ? letters.toLowerCase() : letters.toUpperCase(); } /** Which slice of the flattened token sequence belongs to which word. */ @@ -62,8 +91,9 @@ interface TokenSequence { function buildTokens(words: Word[], vocab: CtcVocab): TokenSequence { const tokens: number[] = []; const spans: TokenSpan[] = []; + const mode = vocab.normalizeMode ?? "latin-upper"; words.forEach((word, index) => { - const ids = vocab.encode(normalizeForCtc(word.text)); + const ids = vocab.encode(normalizeForCtc(word.text, mode)); if (ids.length === 0) return; if (tokens.length > 0 && vocab.delimiterId !== undefined) { tokens.push(vocab.delimiterId); diff --git a/lib/types.ts b/lib/types.ts index cb88c89..51d82f7 100644 --- a/lib/types.ts +++ b/lib/types.ts @@ -95,6 +95,9 @@ export interface WorkerRequest { duration: number; /** Which local speech model to use (see lib/models.ts). */ model: import("./models").ModelId; - /** Whisper language code for the transcript (ignored by Parakeet auto-detect). */ + /** + * Transcript language: conditions Whisper decoding, selects the CTC aligner, + * and is passed through for Parakeet alignment (Parakeet ASR still auto-detects). + */ language: import("./languages").TranscriptLanguage; } diff --git a/tests/align-models-test.ts b/tests/align-models-test.ts new file mode 100644 index 0000000..8ee87b1 --- /dev/null +++ b/tests/align-models-test.ts @@ -0,0 +1,29 @@ +import { + ALIGN_MODELS, + alignModelFor, + uniqueAlignModelIds, +} from "../lib/alignModels"; + +function assert(cond: boolean, msg: string) { + if (!cond) throw new Error(msg); +} + +{ + assert(alignModelFor("en")?.id === ALIGN_MODELS.en!.id, "en has an aligner"); + assert(alignModelFor("es")?.id === ALIGN_MODELS.es!.id, "es has an aligner"); + assert(alignModelFor("fr")?.normalize === "latin-lower", "fr uses MMS folding"); + assert(alignModelFor("de")?.id === alignModelFor("es")?.id, "eu langs share MMS"); + assert(alignModelFor("zh")?.normalize === "cjk", "zh uses CJK normalize"); + assert(alignModelFor("en")?.normalize === "latin-upper", "en stays uppercase"); + assert(alignModelFor(undefined) === null, "missing language → no CTC"); + assert(alignModelFor("xx") === null, "unknown language → no CTC"); + + const ids = uniqueAlignModelIds(); + assert(ids.length === 3, `expected 3 unique model ids, got ${ids.length}: ${ids}`); + assert(ids.includes(ALIGN_MODELS.en!.id), "english model listed"); + assert(ids.includes(ALIGN_MODELS.es!.id), "mms model listed"); + assert(ids.includes(ALIGN_MODELS.zh!.id), "chinese model listed"); + console.log("align model map: ok", ids); +} + +console.log("ALL ALIGN MODEL TESTS PASSED"); diff --git a/tests/forced-align-test.ts b/tests/forced-align-test.ts index 51774d4..09e883e 100644 --- a/tests/forced-align-test.ts +++ b/tests/forced-align-test.ts @@ -48,6 +48,26 @@ function word(id: number, text: string, start: number, end: number): Word { assert(normalizeForCtc("Shot.") === "SHOT", "punctuation is stripped"); assert(normalizeForCtc("don't") === "DON'T", "apostrophes survive"); assert(normalizeForCtc("1985") === "", "digits have no spelling"); + assert( + normalizeForCtc("café", "latin-upper") === "CAFE", + "accents fold for Latin-upper" + ); + assert( + normalizeForCtc("niño", "latin-lower") === "nino", + "accents fold for Latin-lower / MMS" + ); + assert( + normalizeForCtc("über", "latin-lower") === "uber", + "German umlaut folds" + ); + assert( + normalizeForCtc("你好世界!", "cjk") === "你好世界", + "CJK keeps Han and drops punctuation" + ); + assert( + normalizeForCtc("OK了", "cjk") === "OK了", + "CJK keeps Latin letters uppercased" + ); console.log("normalisation: ok"); } diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index 2d136f5..37cfb1d 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -5,13 +5,15 @@ * 2. ASR (Whisper via transformers.js, or Parakeet TDT v3 via parakeet.js) * transcribes each segment with per-word timestamps, remapped onto the * original timeline. - * 3. Pyannote segmentation 3.0 assigns a speaker to each word. + * 3. CTC forced alignment (language-specific wav2vec2 / MMS) measures word + * boundaries against the audio; a VAD/envelope heuristic is the fallback. + * 4. Pyannote segmentation 3.0 assigns a speaker to each word. * - * Both ASR backends and the CTC aligner are registered in a weightlift - * ModelManager so download progress, cache labeling, and WebGPU→WASM fallback - * share one path. Weights land in Cache Storage (Whisper, aligner) or IndexedDB - * (Parakeet); later runs are offline. ORT WASM is served same-origin from - * /vendor/ort* . + * ASR backends live in one weightlift ModelManager (download progress, cache + * labeling, WebGPU→WASM fallback). The CTC aligner has its own registry so a + * GPU loss does not unload WASM aligner weights. Weights land in Cache Storage + * (Whisper, aligner) or IndexedDB (Parakeet); later runs are offline. ORT WASM + * is served same-origin from /vendor/ort* . */ import { pipeline, @@ -43,6 +45,11 @@ import { } from "@/lib/models"; import { cleanTranscript } from "@/lib/hallucinations"; import { alignWordsToSpeech, speechEnvelope } from "@/lib/align"; +import { + ALIGN_MODELS, + alignModelFor, + type AlignModelInfo, +} from "@/lib/alignModels"; import { insertDisfluencyPlaceholders } from "@/lib/disfluencies"; import { alignBatch, @@ -51,6 +58,7 @@ import { type CtcEmission, type CtcVocab, } from "@/lib/forcedAlign"; +import type { TranscriptLanguage } from "@/lib/languages"; import { VAD_FRAME_SIZE, VAD_SAMPLE_RATE, @@ -71,14 +79,6 @@ if (env.backends?.onnx?.wasm) { const DIARIZATION_MODEL = "onnx-community/pyannote-segmentation-3.0"; const VAD_MODEL = "onnx-community/silero-vad"; -/** - * Character-level CTC model used to place word boundaries against the audio. - * Whisper's own timestamps are a by-product of decoding and only roughly right; - * this measures them. - */ -const ALIGN_MODEL = "Xenova/wav2vec2-base-960h"; -/** `|` in the wav2vec2 vocabulary — the between-words symbol. */ -const ALIGN_DELIMITER_ID = 4; /** Viterbi needs a frames x tokens lattice, so alignment runs in bounded batches. */ const ALIGN_BATCH_MAX_S = 20; /** Context either side of a batch, so edge words are not clipped. */ @@ -272,6 +272,12 @@ models.subscribe((snap) => { }); }); +type Aligner = { + processor: Awaited>; + model: Awaited>; + vocab: CtcVocab; +}; + /** * The aligner gets its own registry rather than joining `models`. * @@ -280,9 +286,22 @@ models.subscribe((snap) => { * registry exists under a WebGPU→WASM fallback policy whose recovery step is * `unloadAll()`. The aligner never asks for a device, so transformers.js runs it * on WASM (`DEFAULT_DEVICE`), where a lost GPU cannot touch it. Sharing the - * registry meant a GPU loss threw away 86 MB that was still perfectly good. + * registry meant a GPU loss threw away perfectly good WASM weights. + * + * Multiple language-specific CTC models are registered; only the one for the + * active transcript language is loaded. */ -const aligners = new ModelManager({ models: { [ALIGN_MODEL]: alignerModel() } }); +function buildAlignerRegistry(): Record> { + const byId = new Map(); + for (const info of Object.values(ALIGN_MODELS)) { + if (!byId.has(info.id)) byId.set(info.id, info); + } + return Object.fromEntries( + [...byId.entries()].map(([id, info]) => [id, alignerModel(info)]) + ); +} + +const aligners = new ModelManager({ models: buildAlignerRegistry() }); async function getAsr(choice: WhisperModel) { return models.load(MODELS[choice].id); @@ -355,12 +374,6 @@ interface DiarizationSegment { confidence: number; } -type Aligner = { - processor: Awaited>; - model: Awaited>; - vocab: CtcVocab; -}; - type Diarizer = { processor: Awaited>; model: Awaited>; @@ -549,31 +562,31 @@ function getDiarizer(): Promise { * Not `transformersModel()`: that builds a `pipeline()`, and this needs the * processor, model and tokenizer separately. The progress wiring is the same. * - * q4 is 86 MB against fp32's 360 MB and scored identically on the test clip - * (0.033 s mean boundary error). q8 is no smaller and markedly worse — it moved - * one word 0.22 s — and fp16 fails to load on onnxruntime-web. + * q4 is used throughout: English wav2vec2 is ~86 MB; MMS / Chinese XLS-R are + * ~240 MB. fp16 fails to load on onnxruntime-web. */ -function alignerModel(): ModelDefinition { +function alignerModel(info: AlignModelInfo): ModelDefinition { return { isCached: () => - isTransformersModelCached(ALIGN_MODEL, { + isTransformersModelCached(info.id, { cacheKey: env.cacheKey ?? "transformers-cache", }), load: async ({ progress }) => { const progress_callback = transformersProgress(progress); const [processor, model, tokenizer] = await Promise.all([ - AutoProcessor.from_pretrained(ALIGN_MODEL, { progress_callback }), - AutoModelForCTC.from_pretrained(ALIGN_MODEL, { + AutoProcessor.from_pretrained(info.id, { progress_callback }), + AutoModelForCTC.from_pretrained(info.id, { dtype: "q4", progress_callback, }), - AutoTokenizer.from_pretrained(ALIGN_MODEL, { progress_callback }), + AutoTokenizer.from_pretrained(info.id, { progress_callback }), ]); /* eslint-disable-next-line @typescript-eslint/no-explicit-any */ const tok = tokenizer as any; const vocab: CtcVocab = { blankId: tok.pad_token_id ?? 0, - delimiterId: ALIGN_DELIMITER_ID, + delimiterId: delimiterIdFromTokenizer(tok), + normalizeMode: info.normalize, encode: (text) => text ? (tok.encode(text, { add_special_tokens: false }) as number[]) : [], }; @@ -582,8 +595,27 @@ function alignerModel(): ModelDefinition { }; } -function getAligner(): Promise { - return aligners.load(ALIGN_MODEL); +/** `|` word delimiter when the tokenizer has one (wav2vec2); MMS has none. */ +function delimiterIdFromTokenizer(tok: { + encode?: (text: string, opts?: { add_special_tokens?: boolean }) => number[]; +}): number | undefined { + try { + const ids = tok.encode?.("|", { add_special_tokens: false }); + if (Array.isArray(ids) && ids.length === 1 && typeof ids[0] === "number") { + return ids[0]; + } + } catch { + // tokenizer rejected "|" + } + return undefined; +} + +function getAligner(language: TranscriptLanguage): Promise { + const info = alignModelFor(language); + if (!info) { + return Promise.reject(new Error(`No CTC aligner for language: ${language}`)); + } + return aligners.load(info.id); } /** Per-frame log-probabilities for one slice of audio. */ @@ -608,18 +640,22 @@ async function ctcEmission(aligner: Aligner, slice: Float32Array): Promise { + const info = alignModelFor(language); + if (!info) return words; + // Subscribe only while we are actually waiting. The aligner also warms in the // background during transcription, and those bytes must not fight the // "Transcribing…" line — but a wait here is worth a bar, since a cold cache @@ -637,13 +673,13 @@ async function forceAlign( value: rec.indeterminate ? null : rec.percent, }); }; - const unsubscribe = aligners.subscribe((snap) => report(snap.models[ALIGN_MODEL])); + const unsubscribe = aligners.subscribe((snap) => report(snap.models[info.id])); // subscribe() only fires on the next change, so prime from current state: a // load that has started but not yet received its first byte would otherwise // leave the UI sitting on "Transcribing…". - report(aligners.status(ALIGN_MODEL)); + report(aligners.status(info.id)); try { - aligner = await getAligner(); + aligner = await getAligner(language); } finally { unsubscribe(); } @@ -681,6 +717,39 @@ async function forceAlign( return out; } +/** + * Shared post-ASR timing pass for Whisper and Parakeet. + * + * 1. VAD / loudness-envelope heuristic (language-agnostic fallback and batch seed) + * 2. CTC forced alignment when a model exists for `language` + * 3. Envelope edge expansion of peaky CTC spans + * 4. Disfluency placeholders (after alignment — "..." has no CTC spelling) + */ +async function refineWordTimestamps( + words: Word[], + speechFrames: boolean[], + audio: Float32Array, + duration: number, + language: TranscriptLanguage +): Promise { + let out = alignWordsToSpeech(words, speechFrames, { + duration, + audio, + sampleRate: VAD_SAMPLE_RATE, + }); + + if (alignModelFor(language)) { + try { + const measured = await forceAlign(out, audio, duration, language); + out = expandToAcoustics(measured, speechEnvelope(audio, VAD_SAMPLE_RATE)); + } catch (err) { + console.warn("Forced alignment unavailable; using VAD-corrected times.", err); + } + } + + return insertDisfluencyPlaceholders(out, speechFrames, { duration }); +} + async function diarize(audio: Float32Array): Promise { const { processor, model } = await getDiarizer(); const inputs = await processor(audio); @@ -785,9 +854,14 @@ async function finishWithDiarization( async function runParakeet( audio: Float32Array, - duration: number + duration: number, + transcriptLanguage: TranscriptLanguage ): Promise { + // Overlap diarizer (+ language-matched aligner) with Parakeet load. getDiarizer().catch(() => {}); + if (alignModelFor(transcriptLanguage)) { + getAligner(transcriptLanguage).catch(() => {}); + } const [loaded, vad] = await Promise.all([getParakeet(), getVad()]); let model = loaded; @@ -851,8 +925,13 @@ async function runParakeet( } const cleaned = cleanTranscript(rawWords); - // Same "..." filled-pause recovery as Whisper (no timestamp align for Parakeet). - const words = insertDisfluencyPlaceholders(cleaned, speechFrames, { duration }); + const words = await refineWordTimestamps( + cleaned, + speechFrames, + audio, + duration, + transcriptLanguage + ); return finishWithDiarization(words, audio); } @@ -860,12 +939,14 @@ async function runWhisper( audio: Float32Array, duration: number, choice: WhisperModel, - transcriptLanguage: WorkerRequest["language"] + transcriptLanguage: TranscriptLanguage ): Promise { - // Overlap Whisper + Silero downloads; diarizer and aligner warm in the - // background so both are cached by the time the transcript lands. + // Overlap Whisper + Silero downloads; diarizer and language-matched aligner + // warm in the background so both are cached by the time the transcript lands. getDiarizer().catch(() => {}); - getAligner().catch(() => {}); + if (alignModelFor(transcriptLanguage)) { + getAligner(transcriptLanguage).catch(() => {}); + } const [asr, vad] = await Promise.all([getAsr(choice), getVad()]); let transcriber = asr; @@ -1034,32 +1115,13 @@ async function runWhisper( // Post-process: collapse leftover n-gram loops and drop known hallucination // phrases ("I'm sorry", "thanks for watching", …) that slip past decoding. const cleaned = cleanTranscript(rawWords); - - // Whisper's DTW word timestamps run late by an amount that drifts across the - // clip. Correct them against the VAD as a first pass: forced alignment only - // needs them accurate enough to batch by, and this is what we fall back to if - // the acoustic model is unavailable. - let words = alignWordsToSpeech(cleaned, speechFrames, { - duration, + const words = await refineWordTimestamps( + cleaned, + speechFrames, audio, - sampleRate: VAD_SAMPLE_RATE, - }); - - // Then measure the boundaries properly. CTC forced alignment decides which - // audio each word occupies; the envelope sharpens the edges, which CTC leaves - // narrow because its emissions spike on each character and go blank around it. - try { - const measured = await forceAlign(words, audio, duration); - words = expandToAcoustics(measured, speechEnvelope(audio, VAD_SAMPLE_RATE)); - } catch (err) { - console.warn("Forced alignment unavailable; using VAD-corrected times.", err); - } - - // Recover filled pauses Whisper omitted as "..." so Remove fillers can cut - // them. After alignment, not before: this looks for speech no word covers, so - // it needs the measured boundaries — and a "..." has no spelling for CTC, so - // aligning it would only interpolate the placeholder away again. - words = insertDisfluencyPlaceholders(words, speechFrames, { duration }); + duration, + transcriptLanguage + ); return finishWithDiarization(words, audio); } @@ -1068,11 +1130,11 @@ self.onmessage = async (event: MessageEvent) => { const { audio, duration, model, language } = event.data; try { const choice: ModelId = model ?? "base"; - const transcriptLanguage = language ?? "en"; + const transcriptLanguage: TranscriptLanguage = language ?? "en"; let words: Word[]; if (isParakeetModel(choice)) { - words = await runParakeet(audio, duration); + words = await runParakeet(audio, duration, transcriptLanguage); } else if (isWhisperModel(choice)) { words = await runWhisper(audio, duration, choice, transcriptLanguage); } else { From b207e41063589fd066890d6fcc06f886c58c8271 Mon Sep 17 00:00:00 2001 From: Wassim Gharbi Date: Wed, 5 Aug 2026 19:25:55 -0700 Subject: [PATCH 6/7] Read the CTC blank and delimiter from the vocabulary, not by inference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both derivations in alignerModel() took a shortcut that is correct on two of the three shipped aligners and silently wrong on the third. Neither failure raises; both just cost timing accuracy on es/fr/de. blankId came from `pad_token_id`. For wav2vec2-base-960h and XLS-R-zh, *is* the CTC blank and sits at 0, so that happened to be right. The MMS forced-aligner declares them separately — at 0, at 1 — so MMS scored the lattice's blank states against a column the model never activates. Viterbi does not fail there, it avoids blank states, which smears every boundary it returns. delimiterId came from `encode("|")`, which cannot detect absence: a vocab without "|" maps it to , one id, indistinguishable from a real hit. MMS has no "|", so was threaded between every pair of words. The old comment ("MMS has none") described the intent correctly; the check did not implement it. Both now go through convert_tokens_to_ids, which yields undefined rather than for a token the vocab lacks. Moved out of the worker into lib/forcedAlign.ts as ctcVocabFromTokenizer() so it is unit-testable without loading a model; the new tests use fakes built from the ids the three real tokenizers actually report. Verified against the live tokenizers — all five languages now resolve blank 0, delimiter 4 for en/zh and none for es/fr/de, and no sample text falls through to . Before: es/fr/de got blank 1 and delimiter 3. Co-Authored-By: Claude Opus 5 (1M context) --- lib/forcedAlign.ts | 64 +++++++++++++++++++++++++++++++++ tests/forced-align-test.ts | 56 +++++++++++++++++++++++++++++ workers/transcription.worker.ts | 31 ++++------------ 3 files changed, 127 insertions(+), 24 deletions(-) diff --git a/lib/forcedAlign.ts b/lib/forcedAlign.ts index bab1a9a..7d0678c 100644 --- a/lib/forcedAlign.ts +++ b/lib/forcedAlign.ts @@ -47,6 +47,70 @@ export interface CtcVocab { normalizeMode?: CtcNormalizeMode; } +/** The subset of a transformers.js tokenizer {@link ctcVocabFromTokenizer} reads. */ +export interface CtcTokenizerLike { + pad_token_id?: number; + unk_token_id?: number; + /** Vocabulary lookup. Yields `undefined` for tokens the vocab does not have. */ + convert_tokens_to_ids?(tokens: string[]): (number | undefined)[]; + encode(text: string, opts?: { add_special_tokens?: boolean }): number[]; +} + +/** First id in `tokens` that the vocabulary actually contains. */ +function lookup( + tok: CtcTokenizerLike, + tokens: string[] +): number | undefined { + let ids: (number | undefined)[] | undefined; + try { + ids = tok.convert_tokens_to_ids?.(tokens); + } catch { + return undefined; + } + if (!Array.isArray(ids)) return undefined; + for (const id of ids) { + // An absent token yields undefined here rather than the unk id, which is + // the whole reason this goes through the vocabulary instead of `encode`. + if (typeof id === "number" && id >= 0 && id !== tok.unk_token_id) return id; + } + return undefined; +} + +/** + * Read a {@link CtcVocab} off a loaded tokenizer. + * + * Both halves have to consult the vocabulary directly, because the two obvious + * shortcuts are wrong on exactly one of the three aligners we ship — and wrong + * silently, costing timing accuracy rather than raising. + * + * **Blank.** For wav2vec2 and XLS-R, `` *is* the CTC blank and sits at id 0, + * so `pad_token_id` happens to be right. The MMS forced-aligner declares them + * separately — `` at 0, `` at 1 — so taking `pad_token_id` there + * scores the lattice's blank states against a column the model never activates. + * Viterbi does not fail; it just avoids blank states, which smears every + * boundary. So prefer an explicit `` and fall back to ``. + * + * **Delimiter.** `encode("|")` cannot detect absence: a vocab without `|` maps it + * to ``, one id, indistinguishable from a real hit. MMS has no `|`, so that + * route threads `` between every pair of words. A vocabulary lookup returns + * undefined instead, which is what {@link buildTokens} wants. + * + * Measured ids: wav2vec2-base-960h and XLS-R-zh give blank 0 / delimiter 4; + * mms-300m-1130-forced-aligner gives blank 0 / no delimiter. + */ +export function ctcVocabFromTokenizer( + tok: CtcTokenizerLike, + normalizeMode: CtcNormalizeMode = "latin-upper" +): CtcVocab { + return { + blankId: lookup(tok, [""]) ?? tok.pad_token_id ?? 0, + delimiterId: lookup(tok, ["|"]), + normalizeMode, + encode: (text) => + text ? tok.encode(text, { add_special_tokens: false }) : [], + }; +} + /** * Keep only what a character-level CTC vocabulary can represent. Digits and * punctuation have no acoustic spelling here, so a word made only of those diff --git a/tests/forced-align-test.ts b/tests/forced-align-test.ts index 09e883e..dd0ce6e 100644 --- a/tests/forced-align-test.ts +++ b/tests/forced-align-test.ts @@ -1,9 +1,11 @@ import { alignBatch, ctcViterbi, + ctcVocabFromTokenizer, groupWordsForAlignment, normalizeForCtc, type CtcEmission, + type CtcTokenizerLike, type CtcVocab, } from "../lib/forcedAlign"; import type { Word } from "../lib/types"; @@ -169,4 +171,58 @@ function word(id: number, text: string, start: number, end: number): Word { console.log(`batching: ok (${many.length} words -> ${batches.length} batches)`); } +/** + * Vocabulary reading, against fakes that reproduce what the three shipped + * aligners actually report. Both derivations have a plausible shortcut that is + * wrong on MMS and silent when it is wrong, so pin the real ids. + */ +{ + /** Build a tokenizer stand-in from a token→id table. */ + const fakeTokenizer = ( + vocab: Record, + padTokenId: number | undefined, + unkTokenId: number + ): CtcTokenizerLike => ({ + pad_token_id: padTokenId, + unk_token_id: unkTokenId, + convert_tokens_to_ids: (tokens) => tokens.map((t) => vocab[t]), + // Character-level, mapping anything absent to unk — as Wav2Vec2CTCTokenizer does. + encode: (text) => text.split("").map((c) => vocab[c] ?? unkTokenId), + }); + + // is the blank and sits at 0; "|" is a real token. + const wav2vec2 = fakeTokenizer({ "": 0, "": 1, "": 2, "": 3, "|": 4 }, 0, 3); + const v1 = ctcVocabFromTokenizer(wav2vec2, "latin-upper"); + assert(v1.blankId === 0, "wav2vec2 blank is at 0"); + assert(v1.delimiterId === 4, "wav2vec2 delimiter is | at 4"); + assert(v1.normalizeMode === "latin-upper", "normalize mode is carried through"); + + // MMS: and are distinct, and there is no "|" at all. + const mms = fakeTokenizer({ "": 0, "": 1, "": 2, "": 3 }, 1, 3); + const v2 = ctcVocabFromTokenizer(mms, "latin-lower"); + assert(v2.blankId === 0, "MMS blank is at 0, not pad_token_id (1)"); + assert( + v2.delimiterId === undefined, + "MMS has no |; encode() would have returned here" + ); + + // A tokenizer with no vocabulary lookup at all still has to yield something. + const bare: CtcTokenizerLike = { pad_token_id: 0, encode: () => [] }; + assert(ctcVocabFromTokenizer(bare).blankId === 0, "falls back to pad_token_id"); + assert(ctcVocabFromTokenizer(bare).delimiterId === undefined, "no lookup -> no delimiter"); + assert(ctcVocabFromTokenizer({ encode: () => [] }).blankId === 0, "last resort blank is 0"); + + // encode() skips specials and drops nothing else. + assert( + JSON.stringify(v1.encode("AB")) === JSON.stringify([3, 3]), + "encode passes through the tokenizer" + ); + assert(v1.encode("").length === 0, "empty text encodes to nothing"); + + console.log( + `vocab from tokenizer: ok (wav2vec2 blank ${v1.blankId}/delim ${v1.delimiterId}, ` + + `MMS blank ${v2.blankId}/delim ${v2.delimiterId})` + ); +} + console.log("ALL FORCED ALIGN TESTS PASSED"); diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index 50ae095..35c20cb 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -55,7 +55,9 @@ import { alignBatch, expandToAcoustics, groupWordsForAlignment, + ctcVocabFromTokenizer, type CtcEmission, + type CtcTokenizerLike, type CtcVocab, } from "@/lib/forcedAlign"; import type { TranscriptLanguage } from "@/lib/languages"; @@ -681,35 +683,16 @@ function alignerModel(info: AlignModelInfo): ModelDefinition { }), AutoTokenizer.from_pretrained(info.id, { progress_callback }), ]); - /* eslint-disable-next-line @typescript-eslint/no-explicit-any */ - const tok = tokenizer as any; - const vocab: CtcVocab = { - blankId: tok.pad_token_id ?? 0, - delimiterId: delimiterIdFromTokenizer(tok), - normalizeMode: info.normalize, - encode: (text) => - text ? (tok.encode(text, { add_special_tokens: false }) as number[]) : [], - }; + const vocab = ctcVocabFromTokenizer( + // Tokenizer internals are untyped in transformers.js. + tokenizer as unknown as CtcTokenizerLike, + info.normalize + ); return { processor, model, vocab }; }, }; } -/** `|` word delimiter when the tokenizer has one (wav2vec2); MMS has none. */ -function delimiterIdFromTokenizer(tok: { - encode?: (text: string, opts?: { add_special_tokens?: boolean }) => number[]; -}): number | undefined { - try { - const ids = tok.encode?.("|", { add_special_tokens: false }); - if (Array.isArray(ids) && ids.length === 1 && typeof ids[0] === "number") { - return ids[0]; - } - } catch { - // tokenizer rejected "|" - } - return undefined; -} - function getAligner(language: TranscriptLanguage): Promise { const info = alignModelFor(language); if (!info) { From 3899aea00ddfac07ee05bcc9428883e2e4cdc789 Mon Sep 17 00:00:00 2001 From: Wassim Gharbi Date: Wed, 5 Aug 2026 19:43:39 -0700 Subject: [PATCH 7/7] Lead every word by a fixed nudge at the end of the timing pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Words read consistently a hair late even with CTC alignment placing the boundaries. Measurement fixed the part of the old ALIGN_LEAD_S that was compensating for a late reference — Silero raises its speech flag ~40 ms after speech starts — but not the reason the lead existed in the first place: a highlight landing slightly early reads as in sync, one landing slightly late reads as lagging, and the two are not equally forgiving. So it comes back, as applyAlignLead() in lib/align.ts, applied in refineWordTimestamps() where Whisper and Parakeet, CTC and heuristic fallback, all converge. Applied last, after insertDisfluencyPlaceholders, for two reasons. A uniform shift over the finished list preserves every relative gap, so the "..." placeholders travel with the words around them instead of drifting against them. And nothing downstream can snap the nudge back onto a VAD onset, which is what the old placement inside alignWordsToSpeech had to work around. Both edges move together; leading only the starts would stretch every word and eventually overlap its neighbour. Starts clamp at 0. The tests bound the constant to (0, 0.15] rather than asserting a value — it is a perceptual setting tuned by ear, and pinning it would make retuning look like a regression. Co-Authored-By: Claude Opus 5 (1M context) --- lib/align.ts | 47 +++++++++++++++++++++++++++++++++ tests/align-test.ts | 44 ++++++++++++++++++++++++++++++ workers/transcription.worker.ts | 11 ++++++-- 3 files changed, 100 insertions(+), 2 deletions(-) diff --git a/lib/align.ts b/lib/align.ts index d5b08b0..1d7eec6 100644 --- a/lib/align.ts +++ b/lib/align.ts @@ -732,3 +732,50 @@ export function alignWordsToSpeech( const snapped = snapWordsToSpeech(warped, speechFrames, resolved); return repairCollapsedWords(snapped, env, options); } + +/** + * Fixed lead subtracted from every word at the end of the timing pipeline. + * + * This is a perceptual nudge, not a measured correction — against acoustic + * ground truth it slightly *increases* absolute error. A highlight that lands a + * touch early reads as in sync; one that lands a touch late reads as lagging, + * and the two are not equally forgiving. + * + * It survived the move to CTC alignment: measurement removed the part of the old + * lead that compensated for a late reference (Silero raises its speech flag + * ~40 ms after speech starts), but not the asymmetry in how early and late read. + * Words still came back reading a hair late once alignment was doing its job. + * + * Tuned by ear, so expect this number to move — the VAD-only pipeline used 80 ms. + * It is applied uniformly, so raising it never breaks relative timing; it only + * trades "reads late" for "reads early", and the tests bound it rather than pin + * it for that reason. + */ +export const ALIGN_LEAD_S = 0.08; + +/** + * Shift every word earlier by `leadS`, uniformly. + * + * Applied last, after disfluency placeholders, for two reasons: a uniform shift + * over the finished list preserves every relative gap (placeholders included), + * and nothing downstream can snap the nudge back onto an onset the way it could + * when this lived inside {@link alignWordsToSpeech}. + * + * Both edges move together — leading only the starts would stretch every word by + * `leadS` and eventually overlap its neighbour. + */ +export function applyAlignLead( + words: Word[], + leadS = ALIGN_LEAD_S, + options: AlignOptions = {} +): Word[] { + const { minWordS = 0.02, duration = 0 } = options; + const out = words.map((w) => ({ ...w })); + if (leadS === 0 || out.length === 0) return out; + for (const w of out) { + w.start -= leadS; + w.end -= leadS; + } + // Clamps starts back to 0 and keeps the ordering and minimum-length guards. + return normalizeWords(out, minWordS, duration); +} diff --git a/tests/align-test.ts b/tests/align-test.ts index 518aa77..6f40840 100644 --- a/tests/align-test.ts +++ b/tests/align-test.ts @@ -1,5 +1,7 @@ import { + ALIGN_LEAD_S, alignWordsToSpeech, + applyAlignLead, buildSpeechAnchors, correctionAt, estimateLagFromEnvelope, @@ -395,4 +397,46 @@ function word(id: number, start: number, end: number): Word { ); } +// The lead is a uniform shift, so every relative gap has to survive it. +{ + const words: Word[] = [ + { id: 0, text: "one", start: 1.0, end: 1.4, speaker: 0, deleted: false }, + { id: 1, text: "...", start: 1.4, end: 1.7, speaker: 0, deleted: false }, + { id: 2, text: "two", start: 1.9, end: 2.3, speaker: 0, deleted: false }, + ]; + const led = applyAlignLead(words, ALIGN_LEAD_S, { duration: 10 }); + assert( + led.every((w, i) => Math.abs(w.start - (words[i]!.start - ALIGN_LEAD_S)) < 1e-9), + "every start moves back by exactly the lead" + ); + assert( + led.every((w, i) => Math.abs(w.end - (words[i]!.end - ALIGN_LEAD_S)) < 1e-9), + "ends move with starts, so durations are unchanged" + ); + assert( + Math.abs((led[2]!.start - led[1]!.end) - (words[2]!.start - words[1]!.end)) < 1e-9, + "the gap before 'two' is preserved" + ); + assert(words[0]!.start === 1.0, "input is not mutated"); + + // A word already at the very start cannot be dragged negative. + const atZero = applyAlignLead( + [{ id: 0, text: "go", start: 0.02, end: 0.3, speaker: 0, deleted: false }], + ALIGN_LEAD_S, + { duration: 10 } + ); + assert(atZero[0]!.start === 0, "clamped to 0 rather than going negative"); + assert(atZero[0]!.end > atZero[0]!.start, "clamping keeps the word non-empty"); + + assert( + applyAlignLead(words, 0, { duration: 10 })[0]!.start === 1.0, + "a zero lead is a no-op" + ); + assert(applyAlignLead([], ALIGN_LEAD_S).length === 0, "empty input is fine"); + // Deliberately a range, not an equality: the lead is a perceptual setting and + // gets tuned by ear. Anything past ~150 ms stops reading as "a hair early". + assert(ALIGN_LEAD_S > 0 && ALIGN_LEAD_S <= 0.15, "lead stays a small nudge"); + console.log(`align lead: ok (${ALIGN_LEAD_S * 1000}ms, gaps and durations preserved)`); +} + console.log("ALL ALIGN TESTS PASSED"); diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index 35c20cb..b4b282b 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -44,7 +44,12 @@ import { type WhisperModel, } from "@/lib/models"; import { cleanTranscript } from "@/lib/hallucinations"; -import { alignWordsToSpeech, speechEnvelope } from "@/lib/align"; +import { + ALIGN_LEAD_S, + alignWordsToSpeech, + applyAlignLead, + speechEnvelope, +} from "@/lib/align"; import { ALIGN_MODELS, alignModelFor, @@ -830,7 +835,9 @@ async function refineWordTimestamps( } } - return insertDisfluencyPlaceholders(out, speechFrames, { duration }); + const withPauses = insertDisfluencyPlaceholders(out, speechFrames, { duration }); + // Last, so the placeholders move with the words around them. + return applyAlignLead(withPauses, ALIGN_LEAD_S, { duration }); } async function diarize(audio: Float32Array): Promise {