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 bd7eac5..2d136f5 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -7,16 +7,19 @@ * 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, AutoProcessor, AutoModel, AutoModelForAudioFrameClassification, + AutoModelForCTC, + AutoTokenizer, WhisperTextStreamer, Tensor, env, @@ -25,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 { @@ -37,8 +42,15 @@ import { type WhisperModel, } from "@/lib/models"; import { cleanTranscript } from "@/lib/hallucinations"; -import { alignWordsToSpeech } from "@/lib/align"; +import { alignWordsToSpeech, speechEnvelope } from "@/lib/align"; import { insertDisfluencyPlaceholders } from "@/lib/disfluencies"; +import { + alignBatch, + expandToAcoustics, + groupWordsForAlignment, + type CtcEmission, + type CtcVocab, +} from "@/lib/forcedAlign"; import { VAD_FRAME_SIZE, VAD_SAMPLE_RATE, @@ -59,6 +71,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. */ @@ -248,6 +272,18 @@ models.subscribe((snap) => { }); }); +/** + * 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); } @@ -259,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(); @@ -318,6 +355,12 @@ interface DiarizationSegment { confidence: number; } +type Aligner = { + processor: Awaited>; + model: Awaited>; + vocab: CtcVocab; +}; + type Diarizer = { processor: Awaited>; model: Awaited>; @@ -498,6 +541,146 @@ function getDiarizer(): Promise { return diarizerPromise; } +/** + * 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. + */ +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, { 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; + 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 }; + }, + }; +} + +function getAligner(): Promise { + return aligners.load(ALIGN_MODEL); +} + +/** 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 { + // 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 report = (rec: ReturnType) => { + if (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, + }); + }; + 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 { + unsubscribe(); + } + // 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; + + 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); @@ -679,8 +862,10 @@ async function runWhisper( choice: WhisperModel, transcriptLanguage: WorkerRequest["language"] ): Promise { - // Overlap Whisper + Silero downloads; diarizer warms in the background. + // Overlap Whisper + Silero downloads; diarizer and aligner warm in the + // background so both are cached by the time the transcript lands. getDiarizer().catch(() => {}); + getAligner().catch(() => {}); const [asr, vad] = await Promise.all([getAsr(choice), getVad()]); let transcriber = asr; @@ -850,13 +1035,31 @@ async function runWhisper( // 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 aligned = 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 { + 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. - const words = insertDisfluencyPlaceholders(aligned, speechFrames, { duration }); + // 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 }); return finishWithDiarization(words, audio); }