From aef394d962d8bab67d0911476d0f85bb82491b4c Mon Sep 17 00:00:00 2001 From: Wassim Gharbi Date: Wed, 5 Aug 2026 20:47:26 -0700 Subject: [PATCH 1/2] fix hallucination pass truncating output + faster transcription live view --- lib/hallucinations.ts | 35 +++++++++---- tests/hallucination-test.ts | 49 ++++++++++++++++++ workers/transcription.worker.ts | 88 ++++++++++++++++++++++++++++++--- 3 files changed, 156 insertions(+), 16 deletions(-) diff --git a/lib/hallucinations.ts b/lib/hallucinations.ts index 29d52fe..1f4429b 100644 --- a/lib/hallucinations.ts +++ b/lib/hallucinations.ts @@ -101,24 +101,39 @@ export function stripHallucinationPhrases(words: Word[]): Word[] { /** * Drop a trailing run of near-identical short words that often appears when * Whisper loops at the end of a clip (e.g. a dozen "um" / "you" / "." tokens). + * + * The run must reach the *last* word. An earlier version scanned backwards for + * the first degenerate window anywhere in the transcript and cut from there, + * which is only a tail trim when the transcript happens to end in a loop: on a + * verbatim transcript a mid-conversation "um uh um uh um uh" matched at minute + * three and took the remaining forty minutes with it. */ export function trimTrailingDegenerateTail(words: Word[]): Word[] { if (words.length < 8) return words; const keys = words.map(wordKey); const window = 6; - let cutFrom = words.length; - - for (let i = words.length - 1; i >= window; i--) { + /** Distinct non-empty tokens in the window ending at `hi`, inclusive. */ + const distinctAt = (hi: number) => { const unique = new Set(); - for (let j = i - window + 1; j <= i; j++) { + for (let j = hi - window + 1; j <= hi; j++) { if (keys[j]) unique.add(keys[j]!); } - // Degenerate: ≤2 distinct tokens in a 6-word window near the end. - if (unique.size > 0 && unique.size <= 2) { - cutFrom = i - window + 1; - } else if (cutFrom < words.length) { - break; - } + return unique.size; + }; + /** Degenerate: ≤2 distinct tokens in a 6-word window. */ + const degenerate = (hi: number) => { + const n = distinctAt(hi); + return n > 0 && n <= 2; + }; + + // Anchored at the end: if the transcript does not finish inside a loop there + // is no tail to trim, whatever happens earlier. + if (!degenerate(words.length - 1)) return words; + + let cutFrom = words.length - window; + for (let i = words.length - 2; i >= window - 1; i--) { + if (!degenerate(i)) break; + cutFrom = i - window + 1; } // Only trim if we'd drop a meaningful chunk (≥6 words) of the tail. if (words.length - cutFrom >= 6) return words.slice(0, cutFrom); diff --git a/tests/hallucination-test.ts b/tests/hallucination-test.ts index 6f0b357..fa20f95 100644 --- a/tests/hallucination-test.ts +++ b/tests/hallucination-test.ts @@ -2,6 +2,7 @@ import { cleanTranscript, collapseRepeatingNgrams, stripHallucinationPhrases, + trimTrailingDegenerateTail, } from "../lib/hallucinations"; import type { Word } from "../lib/types"; @@ -50,4 +51,52 @@ function texts(words: Word[]) { if (texts(out) !== "yes yes okay") throw new Error("over-collapsed short repeats"); } +{ + // A loop that actually ends the clip is still trimmed. + const words = "this is the real content of the talk you you you you you you you you" + .split(" ") + .map(w); + const out = trimTrailingDegenerateTail(words); + console.log("trailing loop trimmed:", texts(out)); + if (/you you you/.test(texts(out))) throw new Error("trailing loop not trimmed"); + if (!texts(out).startsWith("this is the real content")) { + throw new Error("trailing trim ate real speech"); + } +} + +{ + // Regression: a degenerate run in the *middle* must not truncate the tail. + // Verbatim models emit "um uh um uh" runs routinely, and the old backwards + // scan cut from the first such window to the end of the transcript. + const head = "so today we are going to talk about how browsers manage memory".split(" "); + const filler = ["um", "uh", "um", "uh", "um", "uh"]; + const tail = Array.from({ length: 4000 }, (_, k) => `word${k % 137}`); + const words = [...head, ...filler, ...tail].map(w); + const out = trimTrailingDegenerateTail(words); + console.log("mid-transcript filler:", words.length, "->", out.length); + if (out.length !== words.length) { + throw new Error( + `mid-transcript degenerate run truncated the transcript (dropped ${ + words.length - out.length + } words)` + ); + } +} + +{ + // The same shape, end to end, through cleanTranscript. + const head = "welcome back to the show today we have a great guest".split(" "); + const filler = ["um", "uh", "um", "uh", "um", "uh"]; + const tail = "and that is how the whole system fits together thanks everyone".split(" "); + const out = cleanTranscript([...head, ...filler, ...tail].map(w)); + console.log("cleanTranscript keeps tail:", texts(out)); + if (!texts(out).includes("how the whole system fits together")) { + throw new Error("cleanTranscript dropped everything after the filler run"); + } + // Ids stay a dense 0..n-1 range for the store / React keys. + out.forEach((word, i) => { + if (word.id !== i) throw new Error(`id ${word.id} at index ${i}`); + }); +} + console.log("ALL HALLUCINATION TESTS PASSED"); diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index b4b282b..648ed89 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -133,6 +133,78 @@ type AsrChunk = { text: string; timestamp: [number, number | null] }; const post = (msg: WorkerResponse, transfer: Transferable[] = []) => (self as unknown as Worker).postMessage(msg, transfer); +/** + * Live progress / partial-text updates are coalesced onto a timer. + * + * Both fire once per decoded token — tens of thousands of times on a long + * recording — and each one costs a structured clone across the worker + * boundary plus a store write and a React render on the main thread. For the + * partial text that clone is of the whole transcript so far, so the cost grows + * with the transcript and the total work is quadratic: an hour of speech moved + * hundreds of megabytes of short-lived strings for a preview nobody can read + * at that rate. 10 updates a second looks identical and is O(n). + */ +const LIVE_POST_INTERVAL_MS = 100; +/** + * The preview is a "something is happening" affordance pinned above the + * progress bar, not a readable document — only the last few lines are ever on + * screen. Sending the tail keeps each message a fixed size no matter how long + * the recording is. + */ +const PARTIAL_TAIL_CHARS = 1200; + +let pendingProgress: WorkerResponse | null = null; +let pendingPartial: WorkerResponse | null = null; +let liveTimer: ReturnType | null = null; + +function flushLive() { + if (liveTimer !== null) { + clearTimeout(liveTimer); + liveTimer = null; + } + if (pendingProgress) { + post(pendingProgress); + pendingProgress = null; + } + if (pendingPartial) { + post(pendingPartial); + pendingPartial = null; + } +} + +/** + * Drop queued updates without sending them. Used before a terminal message, + * which supersedes anything still in flight. + */ +function cancelLive() { + if (liveTimer !== null) { + clearTimeout(liveTimer); + liveTimer = null; + } + pendingProgress = null; + pendingPartial = null; +} + +/** Queue a coalescing update; the newest value for each type wins. */ +function postLive(msg: WorkerResponse) { + if (msg.type === "partial") pendingPartial = msg; + else pendingProgress = msg; + if (liveTimer === null) { + liveTimer = setTimeout(flushLive, LIVE_POST_INTERVAL_MS); + } +} + +/** Queue the streaming transcript preview, trimmed to its tail. */ +function postPartial(text: string) { + postLive({ + type: "partial", + text: + text.length > PARTIAL_TAIL_CHARS + ? `…${text.slice(-PARTIAL_TAIL_CHARS)}` + : text, + }); +} + /** Device the current ASR pipeline is running on. */ let asrDevice: "webgpu" | "wasm" = "wasm"; @@ -1005,13 +1077,13 @@ async function runParakeet( const piece = (result.utterance_text ?? "").trim(); if (piece) { partial = partial ? `${partial} ${piece}` : piece; - post({ type: "partial", text: partial }); + postPartial(partial); } speechDone += segmentSamples; const value = speechSamples > 0 ? Math.min(1, speechDone / speechSamples) : 1; - post({ type: "progress", message: "Transcribing…", value }); + postLive({ type: "progress", message: "Transcribing…", value }); } const cleaned = cleanTranscript(rawWords); @@ -1085,7 +1157,7 @@ async function runWhisper( chunkFloor = next; chunkTokens = 0; transcribed = next; - post({ type: "progress", message: "Transcribing…", value: transcribed }); + postLive({ type: "progress", message: "Transcribing…", value: transcribed }); }; /** Nudge the bar forward between chunk boundaries as tokens stream in. */ @@ -1097,7 +1169,7 @@ async function runWhisper( const interpolated = Math.min(0.999, chunkFloor + frac * avgChunkDelta); if (interpolated > transcribed) { transcribed = interpolated; - post({ type: "progress", message: "Transcribing…", value: transcribed }); + postLive({ type: "progress", message: "Transcribing…", value: transcribed }); } }; @@ -1159,7 +1231,7 @@ async function runWhisper( }, callback_function: (text: string) => { partial += text; - post({ type: "partial", text: partial }); + postPartial(partial); interpolateProgress(); }, }); @@ -1183,7 +1255,7 @@ async function runWhisper( transcribed = progressBefore.transcribed; chunkFloor = progressBefore.chunkFloor; chunkTokens = progressBefore.chunkTokens; - post({ type: "partial", text: partial }); + postPartial(partial); await fallbackAsrToWasm(); transcriber = await getAsr(choice); chunks = await runSlice(); @@ -1239,9 +1311,13 @@ self.onmessage = async (event: MessageEvent) => { throw new Error(`Unknown speech model: ${String(choice)}`); } + // Drop anything still queued: a stale "Transcribing… 99%" landing after + // "complete" would put the UI back into its busy state. + cancelLive(); post({ type: "complete", words }); } catch (err) { console.error(err); + cancelLive(); post({ type: "error", message: isWebGpuDeviceLostError(err) From 6a13832e7decf1d42cb45b7a3d93492339853a85 Mon Sep 17 00:00:00 2001 From: Wassim Gharbi Date: Wed, 5 Aug 2026 21:29:08 -0700 Subject: [PATCH 2/2] Bound transcription memory: window diarization, drop the retained PCM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Long recordings reload the tab on WebKit ("This webpage was reloaded because it was using significant memory") and lose speakers everywhere else. The causes were all allocations that scale with duration. Diarization ran one forward pass over the entire file. That is a 230 MB input tensor for an hour, SincNet activations several times larger again inside the onnxruntime heap, and a post-processing step that calls logits.tolist() to turn ~213k frames into as many JS arrays. Chromium usually threw and fell back to a single speaker; WebKit killed the tab. lib/diarize.ts windows the audio into bounded passes and stitches pyannote's per-pass class indices together by shared activity in the overlap, so peak memory no longer grows with duration and long files get speakers at all. The decoded PCM was also kept on the main thread for the waveform, which only ever wants one min/max pair per pixel column. lib/waveform.ts builds a bounded envelope instead and the buffer is transferred to the worker rather than copied, removing two full-length copies. Alongside those: cap ASR segments so continuous speech is not decoded as one file-length slice, and cap the undo stack, which pinned a distinct words array per step. Three passes were quadratic in transcript length and only bite on long files. Each rewrite is verified output-identical to the original by differential fuzzing, not by inspection — the disfluency one caught a real behaviour change on the first attempt. - align.ts nested scans over sorted arrays -> binary search: 735ms -> 20ms at one hour, and near-linear where it was quadratic - disfluencies.ts per-run filter().pop()/find() -> two forward cursors - assignSpeakers O(words x segments) -> single cursor; pre-existing, but only reachable now that diarization survives a long file Co-Authored-By: Claude Opus 5 (1M context) --- components/ExportDialog.tsx | 2 +- components/Timeline.tsx | 24 +-- hooks/useTranscriber.ts | 11 +- lib/align.ts | 94 +++++++++--- lib/diarize.ts | 262 ++++++++++++++++++++++++++++++++ lib/disfluencies.ts | 35 ++++- lib/store.ts | 59 +++++-- lib/vad.ts | 97 +++++++++++- lib/waveform.ts | 101 ++++++++++++ tests/diarize-test.ts | 148 ++++++++++++++++++ tests/vad-test.ts | 58 +++++++ tests/waveform-test.ts | 97 ++++++++++++ workers/transcription.worker.ts | 89 ++++++++--- 13 files changed, 991 insertions(+), 86 deletions(-) create mode 100644 lib/diarize.ts create mode 100644 lib/waveform.ts create mode 100644 tests/diarize-test.ts create mode 100644 tests/waveform-test.ts diff --git a/components/ExportDialog.tsx b/components/ExportDialog.tsx index 74c5c4d..2289c2b 100644 --- a/components/ExportDialog.tsx +++ b/components/ExportDialog.tsx @@ -65,7 +65,7 @@ export default function ExportDialog() { const duration = useEditorStore((s) => s.duration); const words = useEditorStore((s) => s.words); const speakers = useEditorStore((s) => s.speakers); - const hasAudioTrack = useEditorStore((s) => s.audio !== null); + const hasAudioTrack = useEditorStore((s) => s.hasAudio); const status = useEditorStore((s) => s.status); const setStatus = useEditorStore((s) => s.setStatus); const exportUrl = useEditorStore((s) => s.exportUrl); diff --git a/components/Timeline.tsx b/components/Timeline.tsx index 52837bd..8329750 100644 --- a/components/Timeline.tsx +++ b/components/Timeline.tsx @@ -38,6 +38,7 @@ import { import type { ClipSegment, Word } from "@/lib/types"; import { isDisfluencyPlaceholder } from "@/lib/disfluencies"; import { VAD_SAMPLE_RATE } from "@/lib/vad"; +import { peakBetween } from "@/lib/waveform"; import { useCutRanges } from "@/hooks/useCutRanges"; import { useIsDark } from "@/hooks/useIsDark"; @@ -94,7 +95,7 @@ type DragKind = | { type: "trim"; edge: "in" | "out"; time: number; lo: number; hi: number }; export default function Timeline() { - const audio = useEditorStore((s) => s.audio); + const waveform = useEditorStore((s) => s.waveform); const words = useEditorStore((s) => s.words); const sceneBoundaries = useEditorStore((s) => s.sceneBoundaries); const duration = useEditorStore((s) => s.duration); @@ -319,32 +320,23 @@ export default function Timeline() { ctx.restore(); }); - // Waveform - if (!audio) return; + // Waveform. Drawn from the precomputed min/max envelope rather than the + // decoded PCM — see lib/waveform.ts. peakBetween already clamps the + // overshoot hot sources produce, so the bar cannot spill into the wordbar. + if (!waveform) return; const samplesPerPx = VAD_SAMPLE_RATE / pps; - const stride = Math.max(1, Math.floor(samplesPerPx / 40)); for (let x = 0; x < width; x++) { const t = (scrollLeft + x) / pps; if (t > duration) break; const i0 = Math.floor(t * VAD_SAMPLE_RATE); - const i1 = Math.min(audio.length, Math.floor(i0 + samplesPerPx) + 1); - let min = 0; - let max = 0; - for (let i = i0; i < i1; i += stride) { - const v = audio[i]; - if (v < min) min = v; - if (v > max) max = v; - } + const peak = peakBetween(waveform, i0, Math.floor(i0 + samplesPerPx) + 1); const inCut = cuts.some((c) => t >= c.start && t < c.end); ctx.fillStyle = inCut ? "#fca5a5" : "#818cf8"; - // Resampling overshoot on hot sources puts samples past ±1, so clamp to - // full scale rather than letting the bar spill into the wordbar lane. - const peak = Math.min(1, (max - min) / 2); const h = Math.max(1, peak * trackH * WAVE_LANE_FILL); ctx.fillRect(x, midY - h / 2, 1, h); } }, [ - audio, + waveform, cuts, clips, duration, diff --git a/hooks/useTranscriber.ts b/hooks/useTranscriber.ts index d3dcc65..3ccbeb2 100644 --- a/hooks/useTranscriber.ts +++ b/hooks/useTranscriber.ts @@ -85,11 +85,14 @@ export function useTranscriber() { ); }; - // Transfer a copy so the original stays available for the waveform. - const copy = audio.slice(); + // Transfer, not copy: the worker takes ownership of the PCM and `audio` is + // detached here. Nothing on the main thread reads it afterwards — the + // waveform draws from the envelope the store built in setAudio — and on a + // long recording the copy this replaces was hundreds of megabytes held for + // the length of the run. workerRef.current.postMessage( - { audio: copy, duration, model, language: transcriptLanguage }, - [copy.buffer] + { audio, duration, model, language: transcriptLanguage }, + [audio.buffer] ); }, []); diff --git a/lib/align.ts b/lib/align.ts index 1d7eec6..2e1a8c4 100644 --- a/lib/align.ts +++ b/lib/align.ts @@ -308,8 +308,7 @@ function landmarkScoreAtLag( let score = 0; const vote = (times: number[], candidates: number[]) => { for (const t of times) { - let bestDist = Infinity; - for (const c of candidates) bestDist = Math.min(bestDist, Math.abs(t - lag - c)); + const bestDist = nearestDistance(candidates, t - lag); if (bestDist < tol) score += 1 - bestDist / tol; } }; @@ -318,6 +317,34 @@ function landmarkScoreAtLag( return score; } +/** + * Index of the first element of `sorted` that is >= `target`, or `length`. + * + * The arrays this file searches — VAD edges, anchors — are all ascending and + * all get probed once per word, per lag candidate, or both. Scanning them + * linearly makes those passes quadratic in the length of the recording, which + * is unnoticeable on a clip and hundreds of milliseconds on an hour. + */ +function lowerBound(sorted: number[], target: number): number { + let lo = 0; + let hi = sorted.length; + while (lo < hi) { + const mid = (lo + hi) >> 1; + if (sorted[mid] < target) lo = mid + 1; + else hi = mid; + } + return lo; +} + +/** Distance from `target` to the closest value in ascending `sorted`. */ +function nearestDistance(sorted: number[], target: number): number { + if (sorted.length === 0) return Infinity; + const i = lowerBound(sorted, target); + const after = i < sorted.length ? sorted[i] - target : Infinity; + const before = i > 0 ? target - sorted[i - 1] : Infinity; + return Math.min(after, before); +} + /** Candidate shifts, ordered outward from 0 so the smallest wins any tie. */ function lagCandidates(maxLagS: number, lagStepS: number): number[] { const steps = Math.floor(maxLagS / lagStepS); @@ -508,16 +535,23 @@ export function buildSpeechAnchors( const anchors: SpeechAnchor[] = []; words.forEach((w, i) => { if (i > 0 && w.start - words[i - 1].end < minGapS) return; - let onset: number | null = null; - let bestDist = Infinity; - for (const o of onsets) { - const d = Math.abs(w.start - lag - o); - if (d < bestDist) { - bestDist = d; - onset = o; - } - } - if (onset === null || bestDist > landmarkTolS) return; + // Nearest onset to the de-lagged word start. Ties go to the earlier onset, + // matching the linear scan this replaces (which kept its first best). + const target = w.start - lag; + const at = lowerBound(onsets, target); + const before = at > 0 ? onsets[at - 1] : null; + const after = at < onsets.length ? onsets[at] : null; + const onset = + before === null + ? after + : after === null + ? before + : after - target < target - before + ? after + : before; + if (onset === null) return; + const bestDist = Math.abs(target - onset); + if (bestDist > landmarkTolS) return; const prev = anchors[anchors.length - 1]; if (!prev) { @@ -545,14 +579,21 @@ export function correctionAt(anchors: SpeechAnchor[], fallbackLag: number, t: nu if (t <= first.from) return first.from - first.to; const last = anchors[anchors.length - 1]; if (t >= last.from) return last.from - last.to; - for (let k = 0; k + 1 < anchors.length; k++) { - const a = anchors[k]; - const b = anchors[k + 1]; - if (t >= a.from && t <= b.from) { - const da = a.from - a.to; - const db = b.from - b.to; - return da + ((db - da) * (t - a.from)) / (b.from - a.from); - } + // Anchors are strictly increasing in `from` (buildSpeechAnchors guarantees + // it), so binary search the bracketing pair — this runs twice per word. + let lo = 0; + let hi = anchors.length - 1; + while (hi - lo > 1) { + const mid = (lo + hi) >> 1; + if (anchors[mid].from <= t) lo = mid; + else hi = mid; + } + const a = anchors[lo]; + const b = anchors[hi]; + if (t >= a.from && t <= b.from) { + const da = a.from - a.to; + const db = b.from - b.to; + return da + ((db - da) * (t - a.from)) / (b.from - a.from); } return fallbackLag; } @@ -567,15 +608,20 @@ function nearestEdge( ): number | null { let best: number | null = null; let bestDist = Infinity; - for (const v of sorted) { - if (v < target - maxDist) continue; - if (v > target + maxDist) break; // sorted ascending - if (v < lo || v > hi) continue; + // Seek to the window rather than scanning from the start: only the handful of + // edges within `maxDist` can win, and `maxDist` is a snap tolerance of a few + // tens of milliseconds. + const consider = (v: number) => { + if (v < lo || v > hi) return; const d = Math.abs(v - target); if (d < bestDist) { bestDist = d; best = v; } + }; + const at = lowerBound(sorted, target - maxDist); + for (let i = at; i < sorted.length && sorted[i] <= target + maxDist; i++) { + consider(sorted[i]); } return best; } diff --git a/lib/diarize.ts b/lib/diarize.ts new file mode 100644 index 0000000..5ab2c9b --- /dev/null +++ b/lib/diarize.ts @@ -0,0 +1,262 @@ +/** + * Windowing and speaker stitching for pyannote segmentation. + * + * pyannote-segmentation-3.0 is a *local* model: it was trained on 10 s windows + * and emits a powerset class per frame, where the mapping from class index to + * a person is arbitrary and only stable within one forward pass. Running it + * over a whole recording in one pass — which is what the worker used to do — + * is therefore not only unnecessary, it is the single largest allocation in + * the app: + * + * - the feature extractor builds one `[1, 1, num_samples]` tensor (230 MB of + * float32 for an hour at 16 kHz), + * - the SincNet frontend convolves that into activations several times larger + * again inside the onnxruntime heap, + * - and `post_process_speaker_diarization` calls `logits.tolist()`, turning + * ~213k frames into as many small JS arrays before softmaxing each one. + * + * Chromium usually throws somewhere in there and the worker falls back to a + * single speaker; WebKit kills the tab instead ("This webpage was reloaded + * because it was using significant memory"). Either way nobody gets speakers + * on a long file today. + * + * So the audio is windowed. Each window costs a bounded, small amount of + * memory, and the cost of windowing is that class 2 in one window is not class + * 2 in the next. Consecutive windows therefore overlap, and the overlap is used + * to match this window's classes onto the ones already emitted — the classic + * permutation-matching stitch. Everything here is pure and takes plain segment + * lists, so it is tested without loading a model. + */ + +/** One contiguous run of a single pyannote powerset class. */ +export interface DiarizationSegment { + /** Powerset class index. 0 is "no speaker"; higher indices are speakers. */ + id: number; + start: number; + end: number; + confidence: number; +} + +/** One window's raw result, with times relative to the window start. */ +export interface DiarizationWindow { + /** Media-time offset of this window's first sample. */ + offsetS: number; + /** Length of the audio actually fed to the model. */ + durationS: number; + segments: DiarizationSegment[]; +} + +export interface WindowSpan { + startSample: number; + endSample: number; +} + +/** + * Window length in seconds. + * + * Well above the model's 10 s training window (longer context segments more + * consistently) but short enough that one forward pass stays small: 30 s is + * 480k samples, ~2 MB of input against 230 MB for an hour. + */ +export const DIARIZE_WINDOW_S = 30; + +/** + * Overlap between consecutive windows, in seconds. + * + * This is the only evidence available for matching one window's classes onto + * the previous window's, so it has to be long enough to contain speech from + * the speakers who span the boundary. Five seconds covers a normal + * conversational turn without adding much recomputation (a sixth of each + * window is decoded twice). + */ +export const DIARIZE_OVERLAP_S = 5; + +/** + * Ignore a candidate match supported by less than this many seconds of shared + * activity. Below it the "match" is usually a stray frame or two, and inventing + * a new speaker is the better failure: a wrong merge silently attributes one + * person's words to another, whereas a spurious extra speaker is visible and + * fixable from the transcript's speaker labels. + */ +const MIN_MATCH_S = 0.25; + +/** + * Split `totalSamples` into overlapping windows. + * + * The final window is snapped back so it ends exactly at the audio end rather + * than being short — a runt window carries too little context for the model to + * segment well, and stitching it is exactly where a short window fails. + */ +export function diarizationWindows( + totalSamples: number, + sampleRate: number, + { windowS = DIARIZE_WINDOW_S, overlapS = DIARIZE_OVERLAP_S } = {} +): WindowSpan[] { + if (totalSamples <= 0) return []; + const windowSamples = Math.max(1, Math.round(windowS * sampleRate)); + const overlapSamples = Math.max(0, Math.round(overlapS * sampleRate)); + if (totalSamples <= windowSamples) { + return [{ startSample: 0, endSample: totalSamples }]; + } + + const step = Math.max(1, windowSamples - overlapSamples); + const spans: WindowSpan[] = []; + for (let start = 0; start < totalSamples; start += step) { + const end = Math.min(totalSamples, start + windowSamples); + // Snap the tail window back so it is full length instead of a runt. + const snappedStart = end === totalSamples ? Math.max(0, end - windowSamples) : start; + spans.push({ startSample: snappedStart, endSample: end }); + if (end === totalSamples) break; + } + return spans; +} + +type Interval = [number, number]; + +/** Total length of the intersection of two interval lists. */ +function overlapDuration(a: Interval[], b: Interval[]): number { + let total = 0; + let i = 0; + let j = 0; + while (i < a.length && j < b.length) { + const lo = Math.max(a[i][0], b[j][0]); + const hi = Math.min(a[i][1], b[j][1]); + if (hi > lo) total += hi - lo; + if (a[i][1] < b[j][1]) i++; + else j++; + } + return total; +} + +/** Intervals of `segments` belonging to `id`, clipped to [lo, hi]. */ +function intervalsFor( + segments: DiarizationSegment[], + id: number, + lo: number, + hi: number +): Interval[] { + const out: Interval[] = []; + for (const s of segments) { + if (s.id !== id) continue; + const start = Math.max(s.start, lo); + const end = Math.min(s.end, hi); + if (end > start) out.push([start, end]); + } + return out; +} + +/** + * Stitch per-window results into one timeline with globally consistent ids. + * + * Each window's classes are matched onto the already-emitted ones by how much + * activity they share in the region the two windows have in common, greedily, + * best pair first. A class with no good match becomes a new speaker. Windows + * contribute segments only from the midpoint of their overlap with the previous + * window, so every instant is described exactly once, by whichever window has + * more context around it. + */ +export function stitchDiarizationWindows( + windows: DiarizationWindow[] +): DiarizationSegment[] { + const emitted: DiarizationSegment[] = []; + let nextGlobalId = 1; + let prevWindowEnd = -Infinity; + + for (let k = 0; k < windows.length; k++) { + const w = windows[k]; + const winStart = w.offsetS; + const winEnd = w.offsetS + w.durationS; + + // Absolute-time, speaker-only segments for this window. + const local = w.segments + .filter((s) => s.id !== 0 && s.end > s.start) + .map((s) => ({ + ...s, + start: s.start + winStart, + end: Math.min(s.end + winStart, winEnd), + })) + .filter((s) => s.end > s.start); + + // Region shared with the previous window: the only place the two agree + // about what happened, and so the only usable matching evidence. + const shareLo = winStart; + const shareHi = Math.min(prevWindowEnd, winEnd); + const hasShared = k > 0 && shareHi > shareLo; + // Emit from the middle of the shared region, so each side contributes the + // half it saw with more surrounding context. + const emitFrom = hasShared ? (shareLo + shareHi) / 2 : winStart; + + const localIds = [...new Set(local.map((s) => s.id))]; + const mapping = new Map(); + + if (hasShared) { + const globalIds = [...new Set(emitted.map((s) => s.id))]; + const localIntervals = new Map( + localIds.map((id) => [id, intervalsFor(local, id, shareLo, shareHi)]) + ); + const globalIntervals = new Map( + globalIds.map((id) => [id, intervalsFor(emitted, id, shareLo, shareHi)]) + ); + + const candidates: Array<{ local: number; global: number; score: number }> = []; + for (const localId of localIds) { + for (const globalId of globalIds) { + const score = overlapDuration( + localIntervals.get(localId)!, + globalIntervals.get(globalId)! + ); + if (score >= MIN_MATCH_S) candidates.push({ local: localId, global: globalId, score }); + } + } + // Greedy best-first. With at most a handful of classes per window this is + // equivalent to an optimal assignment in every realistic case, without + // the machinery. + candidates.sort((a, b) => b.score - a.score); + const takenGlobal = new Set(); + for (const c of candidates) { + if (mapping.has(c.local) || takenGlobal.has(c.global)) continue; + mapping.set(c.local, c.global); + takenGlobal.add(c.global); + } + } + + for (const id of localIds) { + if (!mapping.has(id)) mapping.set(id, nextGlobalId++); + } + + for (const s of local) { + const start = Math.max(s.start, emitFrom); + if (s.end <= start) continue; + emitted.push({ ...s, id: mapping.get(s.id)!, start, end: s.end }); + } + prevWindowEnd = winEnd; + } + + emitted.sort((a, b) => a.start - b.start || a.end - b.end); + return mergeAdjacent(emitted); +} + +/** + * Join runs of the same speaker split only by a window boundary. Confidence is + * averaged by duration so a merged run reports the confidence of the whole span + * rather than of whichever piece happened to come last. + */ +function mergeAdjacent(segments: DiarizationSegment[]): DiarizationSegment[] { + const out: DiarizationSegment[] = []; + for (const s of segments) { + const prev = out[out.length - 1]; + if (prev && prev.id === s.id && s.start - prev.end <= 1e-6) { + const prevLen = prev.end - prev.start; + const thisLen = s.end - s.start; + const total = prevLen + thisLen; + prev.confidence = + total > 0 + ? (prev.confidence * prevLen + s.confidence * thisLen) / total + : prev.confidence; + prev.end = Math.max(prev.end, s.end); + } else { + out.push({ ...s }); + } + } + return out; +} diff --git a/lib/disfluencies.ts b/lib/disfluencies.ts index 0ad0204..508653c 100644 --- a/lib/disfluencies.ts +++ b/lib/disfluencies.ts @@ -93,20 +93,45 @@ export function insertDisfluencyPlaceholders( const runs = uncoveredSpeechRuns(sorted, speechFrames, frameS); if (runs.length === 0) return words; + const kept = sorted.filter((w) => !w.deleted); + // Ends, ascending, each tagged with its position in `kept`. Ends are not + // monotonic in start order — a long word can end after a shorter one that + // starts later — so finding "the last word ending by T" needs this second + // ordering rather than a scan of `kept`. + const byEnd = kept + .map((w, i) => ({ end: w.end, index: i })) + .sort((a, b) => a.end - b.end); + + // `runs` is ascending and both thresholds below derive from it, so each + // lookup only ever moves forward: one pointer each, walked once across the + // transcript. Rescanning `kept` per run instead (a filter plus a find) is + // quadratic, and allocates a full-length array per run. + let endPtr = 0; + /** Greatest position in `kept` among words ending at or before the threshold. */ + let prevIndex = -1; + let nextPtr = 0; + const placeholders: Word[] = []; for (const run of runs) { let start = run.start; let end = Math.min(run.end, maxT); // Hangover after the preceding word is not a filled pause. - const prev = sorted.filter((w) => !w.deleted && w.end <= start + frameS).pop(); - if (prev && start - prev.end <= frameS + 1e-4) { + const endLimit = start + frameS; + while (endPtr < byEnd.length && byEnd[endPtr].end <= endLimit) { + if (byEnd[endPtr].index > prevIndex) prevIndex = byEnd[endPtr].index; + endPtr++; + } + const prevEnd = prevIndex >= 0 ? kept[prevIndex].end : null; + if (prevEnd !== null && start - prevEnd <= frameS + 1e-4) { start = Math.min(end, start + HANGOVER_TRIM_S); } - // Don't spill into the next word (alignment / frame rounding). - const next = sorted.find((w) => !w.deleted && w.start >= start - frameS); - if (next) end = Math.min(end, next.start); + // Don't spill into the next word (alignment / frame rounding). `start` may + // have just moved later, which only ever advances this bound too. + const startLimit = start - frameS; + while (nextPtr < kept.length && kept[nextPtr].start < startLimit) nextPtr++; + if (nextPtr < kept.length) end = Math.min(end, kept[nextPtr].start); if (end - start < minDurationS - 1e-4) continue; if (start >= maxT - 1e-4) continue; diff --git a/lib/store.ts b/lib/store.ts index 8293dc5..4c7e867 100644 --- a/lib/store.ts +++ b/lib/store.ts @@ -37,6 +37,7 @@ import { type TranscriptLanguage, } from "./languages"; import { detectMediaKind, type MediaKind } from "./media"; +import { buildWaveformPeaks, type WaveformPeaks } from "./waveform"; import { deleteProject, fileFromProject, @@ -66,8 +67,16 @@ interface EditorState { /** Whether the loaded file is video or audio-only. */ mediaKind: MediaKind | null; duration: number; - /** Mono 16 kHz PCM of the media's audio track (used for waveform + ASR). */ - audio: Float32Array | null; + /** + * Min/max envelope of the media's audio track, for the timeline waveform. + * + * The decoded PCM itself is deliberately not kept: it is hundreds of + * megabytes on a long recording and the worker takes ownership of it (see + * useTranscriber). Null when the file has no audio track. + */ + waveform: WaveformPeaks | null; + /** Whether the media has an audio track at all. */ + hasAudio: boolean; /** Transcript source selected on the upload screen (speech model or import). */ source: TranscriptSource; /** Language hint sent to Whisper when transcribing (Parakeet auto-detects). */ @@ -144,6 +153,11 @@ interface EditorState { setTranscriptLanguage: (language: TranscriptLanguage) => void; setPendingTranscript: (t: PendingTranscript | null) => void; setDuration: (d: number) => void; + /** + * Hand the decoded PCM to the store. Only the waveform envelope is retained; + * the caller keeps ownership of the buffer itself (and transfers it to the + * transcription worker). + */ setAudio: (a: Float32Array | null) => void; setStatus: (s: EditorStatus) => void; setProgress: (p: ProgressInfo) => void; @@ -228,6 +242,23 @@ function bumpAutosave() { void import("./autosave").then((m) => m.scheduleProjectAutosave()); } +/** + * How many undo steps to keep. + * + * Snapshots share structure with the live state, but every edit replaces the + * words array wholesale, so each entry pins a distinct copy — on an hour-long + * transcript that is a few megabytes per step. Unbounded, a long editing + * session grows without limit and never gives any of it back. A hundred steps + * is far more than anyone walks back through interactively. + */ +const MAX_UNDO_STEPS = 100; + +/** Append to the undo stack, dropping the oldest entries past the cap. */ +function pushHistory(past: EditSnapshot[], entry: EditSnapshot): EditSnapshot[] { + const next = [...past, entry]; + return next.length > MAX_UNDO_STEPS ? next.slice(next.length - MAX_UNDO_STEPS) : next; +} + function snapshotOf(s: { words: Word[]; speakers: SpeakerInfo[]; @@ -283,7 +314,7 @@ function pushEdit( set({ future: [], ...next }); } else { set({ - past: [...s.past, snapshotOf(s)], + past: pushHistory(s.past, snapshotOf(s)), future: [], ...next, }); @@ -296,7 +327,8 @@ export const useEditorStore = create((set, get) => ({ mediaUrl: null, mediaKind: null, duration: 0, - audio: null, + waveform: null, + hasAudio: false, source: "base", transcriptLanguage: DEFAULT_TRANSCRIPT_LANGUAGE, pendingTranscript: null, @@ -369,7 +401,8 @@ export const useEditorStore = create((set, get) => ({ error: null, currentTime: 0, exportUrl: null, - audio: null, + waveform: null, + hasAudio: false, duration: 0, }); // Funnel step between opening the app and getting a transcript. `kind` and @@ -421,7 +454,8 @@ export const useEditorStore = create((set, get) => ({ playing: false, exportUrl: null, exportOpen: false, - audio: null, + waveform: null, + hasAudio: false, }); }, @@ -449,7 +483,11 @@ export const useEditorStore = create((set, get) => ({ set({ duration }); if (get().status === "ready") bumpAutosave(); }, - setAudio: (audio) => set({ audio }), + setAudio: (audio) => + set({ + waveform: audio && audio.length > 0 ? buildWaveformPeaks(audio) : null, + hasAudio: audio !== null, + }), setStatus: (status) => { set({ status }); if (status === "ready") bumpAutosave(); @@ -836,7 +874,7 @@ export const useEditorStore = create((set, get) => ({ if (s.gestureActive) return; set({ gestureActive: true, - past: [...s.past, snapshotOf(s)], + past: pushHistory(s.past, snapshotOf(s)), future: [], }); }, @@ -886,7 +924,7 @@ export const useEditorStore = create((set, get) => ({ manualCuts: next.manualCuts, sceneBoundaries: next.sceneBoundaries, future: future.slice(1), - past: [...past, { words, speakers, manualCuts, sceneBoundaries }], + past: pushHistory(past, { words, speakers, manualCuts, sceneBoundaries }), selectedClipIndex: null, selectedCutIndex: null, selectedWordIds: [], @@ -933,7 +971,8 @@ export const useEditorStore = create((set, get) => ({ mediaUrl: null, mediaKind: null, duration: 0, - audio: null, + waveform: null, + hasAudio: false, source: loadModelPreference(), transcriptLanguage: loadTranscriptLanguagePreference(), pendingTranscript: null, diff --git a/lib/vad.ts b/lib/vad.ts index 831db05..f03feb8 100644 --- a/lib/vad.ts +++ b/lib/vad.ts @@ -28,17 +28,73 @@ export interface SpeechSegmentOptions { padS?: number; /** Drop segments shorter than this (seconds). Default 0.15. */ minSpeechS?: number; + /** + * Longest segment handed to ASR, in seconds. Default {@link MAX_SEGMENT_S}. + * Longer merged runs are split at their quietest interior point. + */ + maxSegmentS?: number; frameSize?: number; sampleRate?: number; } +/** + * Longest speech segment produced, in seconds. + * + * `maxGapS` alone puts no ceiling on segment length: continuous speech with no + * pause longer than 1.5 s — a lecture, a podcast, most talking-head video — is + * one run from the first word to the last. The caller then copies that whole + * span into a padded buffer before decoding it, which duplicates the entire + * recording in memory (230 MB for an hour) and makes the ASR pipeline + * accumulate its per-chunk state across the entire file in a single call. + * + * Two minutes is far longer than the 29 s window the ASR pipeline chunks to + * internally, so splitting here costs no decoding context that the pipeline was + * keeping anyway, and it bounds the copy at ~4 MB. + */ +export const MAX_SEGMENT_S = 120; + +/** + * Pick a frame to cut a long run at, as close to `target` as possible while + * preferring silence. Searches a window either side of the target for the + * longest non-speech stretch and returns its middle, so a split lands between + * words rather than through one. Falls back to the target itself. + */ +function quietestSplit( + speechFrames: boolean[], + lo: number, + hi: number, + target: number, + searchFrames: number +): number { + const from = Math.max(lo + 1, target - searchFrames); + const to = Math.min(hi - 1, target + searchFrames); + let bestStart = -1; + let bestLen = 0; + for (let i = from; i < to; ) { + if (speechFrames[i]) { + i++; + continue; + } + let j = i + 1; + while (j < to && !speechFrames[j]) j++; + if (j - i > bestLen) { + bestLen = j - i; + bestStart = i; + } + i = j; + } + if (bestStart < 0) return target; + return Math.floor(bestStart + bestLen / 2); +} + /** * Build contiguous speech segments from per-frame speech flags. * * 1. Collect raw speech runs * 2. Expand each run by `padS` * 3. Merge runs whose gap is shorter than `maxGapS` - * 4. Drop runs shorter than `minSpeechS` + * 4. Split runs longer than `maxSegmentS` at their quietest interior point + * 5. Drop runs shorter than `minSpeechS` */ export function speechSegmentsFromFrames( speechFrames: boolean[], @@ -47,6 +103,7 @@ export function speechSegmentsFromFrames( maxGapS = 1.5, padS = 0.25, minSpeechS = 0.15, + maxSegmentS = MAX_SEGMENT_S, frameSize = VAD_FRAME_SIZE, sampleRate = VAD_SAMPLE_RATE, }: SpeechSegmentOptions = {} @@ -85,7 +142,43 @@ export function speechSegmentsFromFrames( } } - return merged.flatMap(([startFrame, endFrame]) => { + // Split anything longer than the ceiling, cutting at the quietest frame near + // each boundary so a split lands between words. + const maxSegmentFrames = Math.max( + 1, + Math.round((maxSegmentS * sampleRate) / frameSize) + ); + // Look a tenth of a segment either way for silence to cut in. + const searchFrames = Math.max(1, Math.round(maxSegmentFrames / 10)); + const bounded: Array<[number, number]> = []; + for (const [start, end] of merged) { + if (end - start <= maxSegmentFrames) { + bounded.push([start, end]); + continue; + } + // Greedy from the last cut actually made, rather than from an even + // division of the run: the search for silence moves a boundary by up to + // `searchFrames`, and measuring the next target from the ideal position + // lets that drift accumulate past the ceiling. + let cut = start; + while (end - cut > maxSegmentFrames) { + const next = quietestSplit( + speechFrames, + cut, + end, + cut + maxSegmentFrames, + searchFrames + ); + // Only ever move the boundary earlier, so the piece cannot exceed the cap. + const nextCut = Math.min(next, cut + maxSegmentFrames); + if (nextCut <= cut) break; + bounded.push([cut, nextCut]); + cut = nextCut; + } + bounded.push([cut, end]); + } + + return bounded.flatMap(([startFrame, endFrame]) => { const startSample = startFrame * frameSize; const endSample = endFrame >= n ? totalSamples : Math.min(totalSamples, endFrame * frameSize); diff --git a/lib/waveform.ts b/lib/waveform.ts new file mode 100644 index 0000000..cc5c704 --- /dev/null +++ b/lib/waveform.ts @@ -0,0 +1,101 @@ +/** + * Downsampled min/max envelope of the audio, for drawing the timeline waveform. + * + * The timeline is the only thing on the main thread that ever wanted the raw + * PCM, and it does not really want it: at every zoom level it collapses a span + * of samples down to one min/max pair per pixel column. Keeping the decoded + * Float32Array around to recompute that is expensive in the one place we can + * least afford it — an hour of mono 16 kHz float32 is 230 MB, held for the whole + * session, alongside the model weights and the onnxruntime heap. WebKit reloads + * the tab well before that adds up. + * + * So the envelope is computed once, the raw buffer is handed to the worker, and + * the main thread keeps a few megabytes instead of a few hundred. + */ + +export interface WaveformPeaks { + /** Source samples summarised by each min/max pair. */ + bucketSize: number; + /** Length of the audio this was built from, in samples. */ + sampleCount: number; + /** Per-bucket extremes, quantised to signed bytes. */ + min: Int8Array; + max: Int8Array; +} + +/** + * Envelope resolution ceiling — about 4 MB at two bytes per bucket. + * + * This is comfortably finer than the timeline can render. The finest span the + * timeline ever asks for is one pixel at maximum zoom, i.e. + * `sampleCount / (trackWidthPx * MAX_ZOOM)` samples; with a 256x zoom cap that + * stays coarser than `sampleCount / 2e6` for any track narrower than ~7800 px. + * Media short enough to fall under the cap keeps full sample resolution. + */ +const MAX_BUCKETS = 2_000_000; + +/** Quantise [-1, 1] to a signed byte, clamping the overshoot hot sources have. */ +function quantise(v: number): number { + const clamped = v < -1 ? -1 : v > 1 ? 1 : v; + return Math.round(clamped * 127); +} + +/** + * Summarise `audio` into a min/max envelope. + * + * One pass, no allocation beyond the two output arrays. `bucketSize` is 1 — + * i.e. lossless in time, only quantised in amplitude — whenever the audio is + * short enough for that to fit under {@link MAX_BUCKETS}. + */ +export function buildWaveformPeaks( + audio: Float32Array, + maxBuckets = MAX_BUCKETS +): WaveformPeaks { + const sampleCount = audio.length; + const bucketSize = Math.max(1, Math.ceil(sampleCount / Math.max(1, maxBuckets))); + const buckets = Math.ceil(sampleCount / bucketSize); + const min = new Int8Array(buckets); + const max = new Int8Array(buckets); + + for (let b = 0; b < buckets; b++) { + const from = b * bucketSize; + const to = Math.min(sampleCount, from + bucketSize); + let lo = 0; + let hi = 0; + for (let i = from; i < to; i++) { + const v = audio[i]; + if (v < lo) lo = v; + else if (v > hi) hi = v; + } + min[b] = quantise(lo); + max[b] = quantise(hi); + } + return { bucketSize, sampleCount, min, max }; +} + +/** + * Peak amplitude over `[startSample, endSample)`, as a 0..1 fraction of full + * scale — half the peak-to-peak swing, which is what the timeline draws. + * + * Buckets are inclusive at both ends, so a span never reads as silent just + * because it fell between two of them. + */ +export function peakBetween( + peaks: WaveformPeaks, + startSample: number, + endSample: number +): number { + const { bucketSize, min, max } = peaks; + const buckets = min.length; + if (buckets === 0) return 0; + const from = Math.max(0, Math.min(buckets - 1, Math.floor(startSample / bucketSize))); + const to = Math.max(from, Math.min(buckets - 1, Math.floor((endSample - 1) / bucketSize))); + + let lo = 0; + let hi = 0; + for (let b = from; b <= to; b++) { + if (min[b] < lo) lo = min[b]; + if (max[b] > hi) hi = max[b]; + } + return Math.min(1, (hi - lo) / 2 / 127); +} diff --git a/tests/diarize-test.ts b/tests/diarize-test.ts new file mode 100644 index 0000000..262ed9e --- /dev/null +++ b/tests/diarize-test.ts @@ -0,0 +1,148 @@ +import { + diarizationWindows, + stitchDiarizationWindows, + DIARIZE_WINDOW_S, + DIARIZE_OVERLAP_S, + type DiarizationSegment, + type DiarizationWindow, +} from "../lib/diarize"; + +const SR = 16_000; + +function seg(id: number, start: number, end: number): DiarizationSegment { + return { id, start, end, confidence: 0.9 }; +} + +/** Which global id covers `t`, or 0 for none. */ +function speakerAt(segments: DiarizationSegment[], t: number): number { + const hit = segments.find((s) => t >= s.start && t < s.end); + return hit ? hit.id : 0; +} + +{ + // Windows tile the audio with the requested overlap and no gaps. + const spans = diarizationWindows(SR * 100, SR); + console.log("windows for 100s:", spans.map((s) => `${s.startSample / SR}-${s.endSample / SR}`).join(" ")); + if (spans[0].startSample !== 0) throw new Error("first window must start at 0"); + if (spans[spans.length - 1].endSample !== SR * 100) { + throw new Error("last window must reach the end of the audio"); + } + for (let i = 1; i < spans.length; i++) { + if (spans[i].startSample > spans[i - 1].endSample) { + throw new Error(`gap between window ${i - 1} and ${i}`); + } + } + for (const s of spans) { + const len = (s.endSample - s.startSample) / SR; + if (len > DIARIZE_WINDOW_S + 1e-6) throw new Error(`window too long: ${len}s`); + } +} + +{ + // Audio shorter than one window is a single pass. + const spans = diarizationWindows(SR * 5, SR); + if (spans.length !== 1 || spans[0].endSample !== SR * 5) { + throw new Error("short audio should produce exactly one full-length window"); + } +} + +{ + // Every window is bounded, whatever the file length — this is the property + // that keeps peak memory flat instead of scaling with duration. + const spans = diarizationWindows(SR * 3600, SR); + const longest = Math.max(...spans.map((s) => (s.endSample - s.startSample) / SR)); + console.log("1h ->", spans.length, "windows, longest", longest, "s"); + if (longest > DIARIZE_WINDOW_S + 1e-6) throw new Error("window grew with duration"); +} + +{ + // The core stitch: pyannote's class indices are arbitrary per window, so the + // same person is class 1 in the first window and class 2 in the second. + // Both windows are 30 s with 5 s of overlap (25..30). + const windows: DiarizationWindow[] = [ + { + offsetS: 0, + durationS: 30, + segments: [seg(1, 0, 12), seg(2, 12, 24), seg(1, 24, 30)], + }, + { + // Local time 0 == media time 25. Alice (global 1) holds 25..35, + // then Bob (global 2) takes over — but the model labelled them 2 and 1. + offsetS: 25, + durationS: 30, + segments: [seg(2, 0, 10), seg(1, 10, 30)], + }, + ]; + const out = stitchDiarizationWindows(windows); + console.log( + "stitched:", + out.map((s) => `${s.id}@${s.start.toFixed(1)}-${s.end.toFixed(1)}`).join(" ") + ); + + // Alice speaks at 26 s (window 1 class 1) and at 33 s (window 2 class 2). + // The stitch must call those the same person. + if (speakerAt(out, 26) !== speakerAt(out, 33)) { + throw new Error("speaker identity was not carried across the window boundary"); + } + // And the speaker who takes over at 35 s must be someone else. + if (speakerAt(out, 40) === speakerAt(out, 33)) { + throw new Error("distinct speakers were merged across the boundary"); + } + // The timeline is covered exactly once, in order. + for (let i = 1; i < out.length; i++) { + if (out[i].start < out[i - 1].end - 1e-6) { + throw new Error("stitched segments overlap"); + } + } + if (out.some((s) => s.end <= s.start)) throw new Error("empty segment emitted"); + const last = out[out.length - 1]; + if (Math.abs(last.end - 55) > 1e-6) throw new Error(`timeline ends at ${last.end}, want 55`); +} + +{ + // A speaker who only appears in the later window gets a new id rather than + // being folded into whoever was talking before. + const windows: DiarizationWindow[] = [ + { offsetS: 0, durationS: 30, segments: [seg(1, 0, 30)] }, + { offsetS: 25, durationS: 30, segments: [seg(1, 0, 5), seg(3, 5, 30)] }, + ]; + const out = stitchDiarizationWindows(windows); + const ids = new Set(out.map((s) => s.id)); + console.log("new speaker:", [...ids].join(","), out.length, "segments"); + if (ids.size !== 2) throw new Error(`expected 2 speakers, got ${ids.size}`); + if (speakerAt(out, 10) === speakerAt(out, 50)) { + throw new Error("a genuinely new speaker was merged into the previous one"); + } +} + +{ + // Same speaker straight through: the boundary must not show up as a split, + // and must not invent a second speaker. + const windows: DiarizationWindow[] = [ + { offsetS: 0, durationS: 30, segments: [seg(1, 0, 30)] }, + { offsetS: 25, durationS: 30, segments: [seg(1, 0, 30)] }, + { offsetS: 50, durationS: 30, segments: [seg(2, 0, 30)] }, + ]; + const out = stitchDiarizationWindows(windows); + console.log("continuous speaker ->", out.length, "segment(s)"); + if (out.length !== 1) throw new Error(`boundary leaked into output: ${out.length} segments`); + if (Math.abs(out[0].end - 80) > 1e-6) throw new Error("continuous run truncated"); +} + +{ + // Silence-only windows contribute nothing and break nothing. + const windows: DiarizationWindow[] = [ + { offsetS: 0, durationS: 30, segments: [seg(0, 0, 30)] }, + { offsetS: 25, durationS: 30, segments: [] }, + ]; + const out = stitchDiarizationWindows(windows); + if (out.length !== 0) throw new Error("silence produced speaker segments"); +} + +{ + if (DIARIZE_OVERLAP_S >= DIARIZE_WINDOW_S) { + throw new Error("overlap must be shorter than the window or windows cannot advance"); + } +} + +console.log("ALL DIARIZE TESTS PASSED"); diff --git a/tests/vad-test.ts b/tests/vad-test.ts index 0a02983..db0258d 100644 --- a/tests/vad-test.ts +++ b/tests/vad-test.ts @@ -1,4 +1,5 @@ import { + MAX_SEGMENT_S, VAD_FRAME_SIZE, energySpeechFrames, speechSegmentsFromFrames, @@ -70,4 +71,61 @@ function assert(cond: boolean, msg: string) { console.log("all silence empty: ok"); } +{ + // Continuous speech with no long pause used to come back as one segment + // spanning the whole recording, which the caller then copies wholesale. + const sr = 16_000; + const framesPerSecond = sr / VAD_FRAME_SIZE; + // 20 minutes of speech, with a brief breath every 10 s (well under maxGapS, + // so none of them split the run on their own). + const frames: boolean[] = []; + for (let s = 0; s < 20 * 60; s++) { + for (let f = 0; f < framesPerSecond; f++) { + // ~0.2 s of quiet at the top of every tenth second. + frames.push(!(s % 10 === 0 && f < framesPerSecond * 0.2)); + } + } + const total = frames.length * VAD_FRAME_SIZE; + const segs = speechSegmentsFromFrames(frames, total, { maxGapS: 1.5, padS: 0 }); + + const longest = Math.max(...segs.map((s) => (s.endSample - s.startSample) / sr)); + console.log(`20min continuous -> ${segs.length} segments, longest ${longest.toFixed(1)}s`); + assert(segs.length > 1, "continuous speech should be split into bounded segments"); + assert(longest <= MAX_SEGMENT_S + 1, `segment of ${longest}s exceeds the ceiling`); + + // Splitting must not lose or reorder audio: segments stay ordered and their + // union still covers the speech the un-capped version covered. + for (let i = 1; i < segs.length; i++) { + assert(segs[i].startSample >= segs[i - 1].endSample, "segments overlap or are unordered"); + } + const covered = segs.reduce((n, s) => n + (s.endSample - s.startSample), 0); + const uncapped = speechSegmentsFromFrames(frames, total, { + maxGapS: 1.5, + padS: 0, + maxSegmentS: Infinity, + }); + const uncappedCovered = uncapped.reduce((n, s) => n + (s.endSample - s.startSample), 0); + assert(uncapped.length === 1, "control: uncapped should still be a single run"); + assert( + Math.abs(covered - uncappedCovered) < sr * 0.05, + `capping changed coverage by ${(uncappedCovered - covered) / sr}s` + ); + + // Cuts should land in the quiet gaps, not through speech. + const cutsInSilence = segs + .slice(1) + .filter((s) => !frames[Math.floor(s.startSample / VAD_FRAME_SIZE)]).length; + console.log(`${cutsInSilence}/${segs.length - 1} splits landed in silence`); + assert(cutsInSilence === segs.length - 1, "a split cut through speech"); +} + +{ + // A segment already under the ceiling is untouched. + const frames = Array(200).fill(true); + const total = frames.length * VAD_FRAME_SIZE; + const segs = speechSegmentsFromFrames(frames, total, { maxGapS: 1.5, padS: 0 }); + assert(segs.length === 1, "short continuous speech should stay one segment"); + assert(segs[0].endSample === total, "segment should reach the end of the audio"); +} + console.log("ALL VAD TESTS PASSED"); diff --git a/tests/waveform-test.ts b/tests/waveform-test.ts new file mode 100644 index 0000000..37dbe84 --- /dev/null +++ b/tests/waveform-test.ts @@ -0,0 +1,97 @@ +import { buildWaveformPeaks, peakBetween } from "../lib/waveform"; + +/** Reference: what the timeline used to compute straight from the PCM. */ +function rawPeak(audio: Float32Array, from: number, to: number): number { + let lo = 0; + let hi = 0; + for (let i = from; i < Math.min(to, audio.length); i++) { + if (audio[i] < lo) lo = audio[i]; + if (audio[i] > hi) hi = audio[i]; + } + return Math.min(1, (hi - lo) / 2); +} + +const SR = 16_000; + +/** A 1 s sine at `freq`, scaled by an envelope over the whole clip. */ +function tone(seconds: number, freq: number, amp: (t: number) => number): Float32Array { + const out = new Float32Array(Math.round(seconds * SR)); + for (let i = 0; i < out.length; i++) { + const t = i / SR; + out[i] = Math.sin(2 * Math.PI * freq * t) * amp(t); + } + return out; +} + +{ + // Short audio keeps full time resolution. + const audio = tone(2, 220, () => 0.8); + const peaks = buildWaveformPeaks(audio); + console.log("2s bucketSize:", peaks.bucketSize, "buckets:", peaks.min.length); + if (peaks.bucketSize !== 1) throw new Error("short audio should not be downsampled"); + if (peaks.sampleCount !== audio.length) throw new Error("sampleCount mismatch"); +} + +{ + // Long audio is capped: this is the property that bounds memory. + const maxBuckets = 1000; + const audio = tone(4, 300, () => 0.9); + const peaks = buildWaveformPeaks(audio, maxBuckets); + console.log("capped buckets:", peaks.min.length, "bucketSize:", peaks.bucketSize); + if (peaks.min.length > maxBuckets) throw new Error("bucket cap exceeded"); + const bytes = peaks.min.length + peaks.max.length; + if (bytes >= audio.length * 4) throw new Error("envelope is not smaller than the PCM"); +} + +{ + // The envelope must agree with the old raw computation at the resolution the + // timeline actually draws at. + const audio = tone(20, 180, (t) => 0.2 + 0.75 * Math.abs(Math.sin(t / 3))); + const peaks = buildWaveformPeaks(audio, 20_000); + let worst = 0; + // 1200 pixel columns across the clip, as the timeline would. + const columns = 1200; + const samplesPerColumn = audio.length / columns; + for (let x = 0; x < columns; x++) { + const from = Math.floor(x * samplesPerColumn); + const to = Math.floor(from + samplesPerColumn) + 1; + const err = Math.abs(peakBetween(peaks, from, to) - rawPeak(audio, from, to)); + if (err > worst) worst = err; + } + console.log("worst per-column error:", worst.toFixed(5)); + // Quantisation to signed bytes is 1/127; allow a shade over for bucket edges. + if (worst > 0.02) throw new Error(`envelope drifts from the PCM by ${worst}`); +} + +{ + // A short loud transient must survive downsampling — min/max buckets exist + // precisely so that averaging cannot swallow it. + const audio = new Float32Array(SR * 4); + for (let i = SR * 2; i < SR * 2 + 50; i++) audio[i] = 0.95; + const peaks = buildWaveformPeaks(audio, 2000); + const at = peakBetween(peaks, SR * 2 - 10, SR * 2 + 60); + console.log("transient peak:", at.toFixed(3)); + if (at < 0.4) throw new Error("downsampling swallowed a transient"); + if (peakBetween(peaks, 0, SR) > 0.01) throw new Error("silence should read as silent"); +} + +{ + // Out-of-range and empty queries are clamped, not crashes. + const peaks = buildWaveformPeaks(tone(1, 200, () => 0.5)); + if (peakBetween(peaks, -500, 10) < 0) throw new Error("negative start broke the query"); + if (peakBetween(peaks, 1e9, 1e9 + 10) !== peakBetween(peaks, 1e9, 1e9 + 10)) { + throw new Error("out-of-range query returned NaN"); + } + const empty = buildWaveformPeaks(new Float32Array(0)); + if (peakBetween(empty, 0, 10) !== 0) throw new Error("empty audio should read as silent"); +} + +{ + // Hot sources overshoot full scale; the bar must not spill past the lane. + const audio = new Float32Array(SR); + for (let i = 0; i < audio.length; i++) audio[i] = i % 2 === 0 ? 1.8 : -1.8; + const peaks = buildWaveformPeaks(audio); + if (peakBetween(peaks, 0, SR) > 1) throw new Error("overshoot was not clamped"); +} + +console.log("ALL WAVEFORM TESTS PASSED"); diff --git a/workers/transcription.worker.ts b/workers/transcription.worker.ts index 648ed89..7a574dd 100644 --- a/workers/transcription.worker.ts +++ b/workers/transcription.worker.ts @@ -56,6 +56,12 @@ import { type AlignModelInfo, } from "@/lib/alignModels"; import { insertDisfluencyPlaceholders } from "@/lib/disfluencies"; +import { + diarizationWindows, + stitchDiarizationWindows, + type DiarizationSegment, + type DiarizationWindow, +} from "@/lib/diarize"; import { alignBatch, expandToAcoustics, @@ -144,14 +150,14 @@ const post = (msg: WorkerResponse, transfer: Transferable[] = []) => * hundreds of megabytes of short-lived strings for a preview nobody can read * at that rate. 10 updates a second looks identical and is O(n). */ -const LIVE_POST_INTERVAL_MS = 100; +const LIVE_POST_INTERVAL_MS = 50; /** * The preview is a "something is happening" affordance pinned above the * progress bar, not a readable document — only the last few lines are ever on * screen. Sending the tail keeps each message a fixed size no matter how long * the recording is. */ -const PARTIAL_TAIL_CHARS = 1200; +const PARTIAL_TAIL_CHARS = 550; let pendingProgress: WorkerResponse | null = null; let pendingPartial: WorkerResponse | null = null; @@ -546,13 +552,6 @@ async function fallbackAsrToWasm() { }); } -interface DiarizationSegment { - id: number; - start: number; - end: number; - confidence: number; -} - type Diarizer = { processor: Awaited>; model: Awaited>; @@ -912,10 +911,17 @@ async function refineWordTimestamps( return applyAlignLead(withPauses, ALIGN_LEAD_S, { duration }); } +/** + * Segment the whole recording, one bounded window at a time. + * + * Feeding the model the entire file was the app's largest allocation by a wide + * margin — see the note at the top of lib/diarize.ts. Windowing keeps every + * forward pass the same size whatever the duration; the price is that pyannote's + * class indices are only meaningful within a pass, which is what the overlap and + * `stitchDiarizationWindows` are for. + */ async function diarize(audio: Float32Array): Promise { const { processor, model } = await getDiarizer(); - const inputs = await processor(audio); - const { logits } = await model(inputs); // post_process_speaker_diarization is specific to the PyAnnote processor // and is not part of the generic Processor typings. const pyannote = processor as unknown as { @@ -924,8 +930,28 @@ async function diarize(audio: Float32Array): Promise { numSamples: number ) => DiarizationSegment[][]; }; - const result = pyannote.post_process_speaker_diarization(logits, audio.length); - return result[0] ?? []; + + const spans = diarizationWindows(audio.length, VAD_SAMPLE_RATE); + const windows: DiarizationWindow[] = []; + for (let i = 0; i < spans.length; i++) { + const { startSample, endSample } = spans[i]; + // Fresh buffer rather than a subarray: non-zero byteOffset views have + // produced wrong results from onnxruntime-web elsewhere in this worker. + const slice = audio.slice(startSample, endSample); + const inputs = await processor(slice); + const { logits } = await model(inputs); + windows.push({ + offsetS: startSample / VAD_SAMPLE_RATE, + durationS: slice.length / VAD_SAMPLE_RATE, + segments: pyannote.post_process_speaker_diarization(logits, slice.length)[0] ?? [], + }); + postLive({ + type: "progress", + message: "Identifying speakers…", + value: (i + 1) / spans.length, + }); + } + return stitchDiarizationWindows(windows); } /** Assign a speaker to each word from the diarization segments. */ @@ -936,19 +962,34 @@ function assignSpeakers(words: Word[], segments: DiarizationSegment[]) { for (const w of words) w.speaker = 0; return; } + // Both lists run in time order, so a single cursor walks them together. + // Rescanning every segment per word is O(words x segments) — fine on a clip, + // but an hour of speech is thousands of each and this used to be unreachable + // only because diarizing a file that long failed outright. + const byStart = [...speech].sort((a, b) => a.start - b.start); + let cursor = 0; + const idMap = new Map(); // pyannote id -> sequential index for (const w of words) { const mid = (w.start + w.end) / 2; - let seg = speech.find((s) => mid >= s.start && mid < s.end); - if (!seg) { - // Fall back to the nearest speech segment. - let best = Infinity; - for (const s of speech) { - const d = mid < s.start ? s.start - mid : mid - s.end; - if (d < best) { - best = d; - seg = s; - } + // Advance past segments that end before this word and can no longer be the + // containing one. Words are in time order, so this never rewinds. + while (cursor + 1 < byStart.length && byStart[cursor].end <= mid) cursor++; + + let seg: DiarizationSegment | undefined; + let best = Infinity; + // The containing segment, or failing that the nearest, is at the cursor or + // immediately beside it — a constant-size neighbourhood, not a full scan. + for (let i = Math.max(0, cursor - 1); i < byStart.length && i <= cursor + 1; i++) { + const s = byStart[i]; + if (mid >= s.start && mid < s.end) { + seg = s; + break; + } + const d = mid < s.start ? s.start - mid : mid - s.end; + if (d < best) { + best = d; + seg = s; } } const raw = seg ? seg.id : -1; @@ -1005,7 +1046,7 @@ async function finishWithDiarization( audio: Float32Array ): Promise { try { - post({ type: "progress", message: "Identifying speakers…", value: null }); + post({ type: "progress", message: "Identifying speakers…", value: 0 }); const segments = await diarize(audio); assignSpeakers(words, segments); } catch (err) {