From 0c3a83c23ce56643f6e04729b3aea9efffa0c4be Mon Sep 17 00:00:00 2001 From: Lucas Winther Date: Thu, 29 Oct 2026 22:02:47 +0100 Subject: [PATCH] Run Silero VAD and Whisper decodes for the long-form pipeline on Transformers.js Co-Authored-By: Claude Opus 5.5 --- src/worker/longform/transformers.ts | 232 ++++++++++++++++++++++++++++ tests/clips.ts | 49 ++++++ tests/longform-transformers.test.ts | 75 +++++++++ 3 files changed, 356 insertions(+) create mode 100644 src/worker/longform/transformers.ts create mode 100644 tests/clips.ts create mode 100644 tests/longform-transformers.test.ts diff --git a/src/worker/longform/transformers.ts b/src/worker/longform/transformers.ts new file mode 100644 index 0000000..419dffe --- /dev/null +++ b/src/worker/longform/transformers.ts @@ -0,0 +1,232 @@ +/** + * `LongFormDeps` over the existing Transformers.js runtime (D-WS3a/D-WS3b): + * Silero VAD runs as a plain ONNX model on the same onnxruntime instance the + * Whisper pipeline uses, and Whisper decodes go through the already-loaded + * ASR pipeline — no second inference runtime. Free of worker globals so + * `bun test` drives it on onnxruntime-node. + */ + +import { + AutoModel, + LogitsProcessor, + LogitsProcessorList, + Tensor, + WhisperTextStreamer, + type AutomaticSpeechRecognitionPipeline, + type PretrainedModelOptions, +} from "@huggingface/transformers"; +import { getDownloadedModels } from "../../lib/modelCache"; +import type { TranscriptChunk } from "../protocol"; +import type { DecodeRequest, DecodeResult, LongFormDeps } from "./types"; +import { SAMPLE_RATE, VAD_WINDOW_SAMPLES } from "./vad"; + +/** + * Silero VAD v5, ONNX export by onnx-community (MIT; `onnx/model.onnx`, + * fp32, 2.2 MB). Fetched from the Hub like Whisper weights — CORS-enabled, + * so COEP allows it — and kept in the same Cache API store (D-005, D-WS3b). + */ +export const VAD_MODEL_ID = "onnx-community/silero-vad"; + +/** Samples of the previous window prepended to each one (Silero v5 wrapper). */ +const VAD_CONTEXT_SAMPLES = 64; + +export interface SileroVad { + speechProbabilities(audio: Float32Array): Promise; + dispose(): Promise; +} + +export interface LoadSileroVadOptions { + /** + * Never touch the network: load only if the model is already cached, + * otherwise reject (offline, or never downloaded — Review Focus 3). + */ + localOnly: boolean; + /** onnxruntime device; the worker passes "wasm" (a 2 MB model gains nothing from a GPU). */ + device?: PretrainedModelOptions["device"]; + /** Is the model in the cache? Defaults to the browser Cache API probe (modelCache.ts). */ + isCached?: () => Promise; +} + +async function vadInBrowserCache(): Promise { + return (await getDownloadedModels([VAD_MODEL_ID])).has(VAD_MODEL_ID); +} + +export async function loadSileroVad(options: LoadSileroVadOptions): Promise { + if (options.localOnly && !(await (options.isCached ?? vadInBrowserCache)())) { + throw new Error("The speech-detection model is not downloaded (D-WS3b)."); + } + const model = await AutoModel.from_pretrained(VAD_MODEL_ID, { + // The repo ships no config.json: the model is a bare ONNX graph. + config: { model_type: "custom" } as unknown as PretrainedModelOptions["config"], + dtype: "fp32", + ...(options.device !== undefined && { device: options.device }), + }); + return { + async speechProbabilities(audio) { + const windows = Math.ceil(audio.length / VAD_WINDOW_SAMPLES); + const probs = new Float32Array(windows); + const sr = new Tensor("int64", [BigInt(SAMPLE_RATE)], []); + let state = new Tensor("float32", new Float32Array(2 * 128), [2, 1, 128]); + const input = new Float32Array(VAD_CONTEXT_SAMPLES + VAD_WINDOW_SAMPLES); + for (let i = 0; i < windows; i++) { + // [64 samples of context | 512-sample window], zero-padded at the end. + input.copyWithin(0, VAD_WINDOW_SAMPLES); + const window = audio.subarray(i * VAD_WINDOW_SAMPLES, (i + 1) * VAD_WINDOW_SAMPLES); + input.set(window, VAD_CONTEXT_SAMPLES); + input.fill(0, VAD_CONTEXT_SAMPLES + window.length); + const out = (await model({ + input: new Tensor("float32", input, [1, input.length]), + sr, + state, + })) as { output: Tensor; stateN: Tensor }; + probs[i] = (out.output.data as Float32Array)[0]!; + state = out.stateN; + } + return probs; + }, + dispose: async () => { + await model.dispose(); + }, + }; +} + +// --------------------------------------------------------------------------- +// Whisper decodes +// --------------------------------------------------------------------------- + +/** + * Records the log-probability of each generated token and the token count. + * Kept last in the processor list (after Whisper's timestamp rules), so it + * sees the distribution the token is actually picked from. A step's token is + * known on the next call (it is the last input id), so the final token (EOS + * as a rule) is not averaged in — close enough for the −1.0 threshold. + */ +class LogProbRecorder extends LogitsProcessor { + steps = 0; + private sum = 0; + private count = 0; + private previous: { data: Float32Array; logSumExp: number } | null = null; + + override _call(inputIds: bigint[][], logits: Tensor): Tensor { + const ids = inputIds[0]!; + if (this.previous) { + const token = Number(ids[ids.length - 1]); + this.sum += this.previous.data[token]! - this.previous.logSumExp; + this.count += 1; + } + const data = (logits as unknown as Tensor[])[0]!.data as Float32Array; + let max = -Infinity; + for (const v of data) if (v > max) max = v; + let total = 0; + for (const v of data) total += Math.exp(v - max); + this.previous = { data, logSumExp: max + Math.log(total) }; + this.steps += 1; + return logits; + } + + average(): number | null { + return this.count > 0 ? this.sum / this.count : null; + } +} + +/** A processor list that keeps `last` at the end whatever Whisper pushes. */ +class ListWithLast extends LogitsProcessorList { + constructor(last: LogitsProcessor) { + super(); + this.processors.push(last); + } + + override push(item: LogitsProcessor): void { + this.processors.splice(this.processors.length - 1, 0, item); + } +} + +/** Only language tokens may be generated — one-step language detection. */ +class LanguageOnly extends LogitsProcessor { + constructor(private readonly allowed: Set) { + super(); + } + + override _call(_inputIds: bigint[][], logits: Tensor): Tensor { + const data = (logits as unknown as Tensor[])[0]!.data as Float32Array; + for (let i = 0; i < data.length; i++) if (!this.allowed.has(i)) data[i] = -Infinity; + return logits; + } +} + +interface WhisperGenerationConfig { + decoder_start_token_id: number; + lang_to_id?: Record; +} + +export function createWhisperDeps( + asr: AutomaticSpeechRecognitionPipeline, + options: { multilingual: boolean }, +): Pick { + const decode = async (audio: Float32Array, request: DecodeRequest): Promise => { + const recorder = new LogProbRecorder(); + const generate: Record = { + return_timestamps: true, + // One window, single-pass generate (D-021's seek-loop note applies). + max_new_tokens: 440, + logits_processor: new ListWithLast(recorder), + }; + if (options.multilingual) { + if (request.language !== undefined) generate.language = request.language; + if (request.task !== undefined) generate.task = request.task; + } + if (request.noRepeatNgramSize > 0) generate.no_repeat_ngram_size = request.noRepeatNgramSize; + if (request.repetitionPenalty !== 1) generate.repetition_penalty = request.repetitionPenalty; + if (request.temperature > 0) { + generate.do_sample = true; + generate.temperature = request.temperature; + } + if (request.onPartial || request.onTimestamp) { + generate.streamer = new WhisperTextStreamer( + asr.tokenizer as ConstructorParameters[0], + { + skip_prompt: true, + callback_function: (text: string) => request.onPartial?.(text), + on_chunk_start: (t: number) => request.onTimestamp?.(t), + on_chunk_end: (t: number) => request.onTimestamp?.(t), + }, + ); + } + const output = await asr(audio, generate); + const single = (Array.isArray(output) ? output[0] : output) as { + text?: string; + chunks?: TranscriptChunk[]; + }; + return { + text: single.text ?? "", + chunks: single.chunks ?? [], + generatedTokens: recorder.steps, + avgLogProb: recorder.average(), + }; + }; + + if (!options.multilingual) return { decode }; + + const detectLanguage = async (audio: Float32Array): Promise => { + const model = asr.model as unknown as { + generation_config: WhisperGenerationConfig; + generate(args: Record): Promise; + }; + const config = model.generation_config; + const languages = Object.entries(config.lang_to_id ?? {}); + const { input_features } = (await asr.processor(audio)) as { input_features: Tensor }; + const list = new LogitsProcessorList(); + list.push(new LanguageOnly(new Set(languages.map(([, id]) => id)))); + const sequences = await model.generate({ + inputs: input_features, + decoder_input_ids: [config.decoder_start_token_id], + max_new_tokens: 1, + logits_processor: list, + }); + const tokenId = Number((sequences.tolist() as bigint[][])[0]!.at(-1)); + const token = languages.find(([, id]) => id === tokenId)?.[0]; + if (!token) throw new Error(`Language detection produced a non-language token (${tokenId}).`); + return token.slice(2, -2); // "<|da|>" → "da" + }; + return { decode, detectLanguage }; +} diff --git a/tests/clips.ts b/tests/clips.ts new file mode 100644 index 0000000..b52a0ca --- /dev/null +++ b/tests/clips.ts @@ -0,0 +1,49 @@ +/** + * Test-time access to the recordings in tests/audio_files: decode with + * ffmpeg-static (trusted devDependency) and the WER text normalization used + * by tests/live-real-model.test.ts. + */ +import { join } from "node:path"; +import ffmpegPath from "ffmpeg-static"; + +export const AUDIO_DIR = join(import.meta.dir, "audio_files"); + +/** Decode `tests/audio_files/` to mono float32 PCM at `rate` (16 kHz by default). */ +export async function decodeClip(name: string, rate = 16_000): Promise { + const path = join(AUDIO_DIR, name); + const proc = Bun.spawn( + [ + ffmpegPath as unknown as string, + ...["-hide_banner", "-loglevel", "error"], + ...["-i", path], + ...["-ac", "1", "-ar", String(rate), "-f", "f32le", "-"], + ], + { stdout: "pipe", stderr: "pipe" }, + ); + const [bytes, errors, exitCode] = await Promise.all([ + new Response(proc.stdout).arrayBuffer(), + new Response(proc.stderr).text(), + proc.exited, + ]); + if (exitCode !== 0) throw new Error(`ffmpeg failed on ${path}: ${errors}`); + return new Float32Array(bytes); +} + +/** Reference transcript of `tests/audio_files/.txt`. */ +export function readReference(base: string): Promise { + return Bun.file(join(AUDIO_DIR, `${base}.txt`)).text(); +} + +/** + * WER normalization: casing, punctuation, and apostrophes are not + * transcription accuracy; "Århus"/"Aarhus" unify (same as the real-model test). + */ +export function normalize(text: string): string { + return text + .toLowerCase() + .replaceAll("århus", "aarhus") + .replace(/['’]/g, "") + .replace(/[.,!?;:"“”‘‚«»()\-–—…]/g, " ") + .replace(/\s+/g, " ") + .trim(); +} diff --git a/tests/longform-transformers.test.ts b/tests/longform-transformers.test.ts new file mode 100644 index 0000000..03da66c --- /dev/null +++ b/tests/longform-transformers.test.ts @@ -0,0 +1,75 @@ +import { afterEach, describe, expect, test } from "bun:test"; +import { existsSync } from "node:fs"; +import { join } from "node:path"; +import { env } from "@huggingface/transformers"; +import { loadSileroVad, VAD_MODEL_ID } from "../src/worker/longform/transformers"; +import { getSpeechTimestamps, SAMPLE_RATE } from "../src/worker/longform/vad"; +import { decodeClip } from "./clips"; + +const realFetch = globalThis.fetch; +afterEach(() => { + globalThis.fetch = realFetch; +}); + +describe("loadSileroVad offline (plan 001-0 Review Focus 3)", () => { + test("localOnly load with nothing cached rejects without calling fetch", async () => { + let fetches = 0; + globalThis.fetch = (async () => { + fetches += 1; + throw new Error("network must not be touched"); + }) as unknown as typeof fetch; + await expect( + loadSileroVad({ localOnly: true, isCached: async () => false }), + ).rejects.toThrow(/not downloaded/); + expect(fetches).toBe(0); + }); +}); + +/** Cached locally, or downloadable right now — otherwise the real-VAD tests skip. */ +async function vadAvailable(): Promise { + if (env.cacheDir && existsSync(join(env.cacheDir, VAD_MODEL_ID))) return true; + try { + const response = await fetch( + `https://huggingface.co/${VAD_MODEL_ID}/resolve/main/onnx/model.onnx`, + { method: "HEAD", signal: AbortSignal.timeout(5_000) }, + ); + return response.ok; + } catch { + return false; + } +} +const available = await vadAvailable(); + +describe.skipIf(!available)("Silero VAD on the Transformers.js runtime (real model)", () => { + const seconds = 10; + const tone = new Float32Array(seconds * SAMPLE_RATE).map( + (_, i) => 0.3 * Math.sin((2 * Math.PI * 440 * i) / SAMPLE_RATE), + ); + let seed = 1; + const noise = new Float32Array(seconds * SAMPLE_RATE).map(() => { + seed = (seed * 1_103_515_245 + 12_345) % 2 ** 31; + return 0.2 * (seed / 2 ** 30 - 1); + }); + + for (const [name, audio] of [ + ["silence", new Float32Array(seconds * SAMPLE_RATE)], + ["a 440 Hz tone", tone], + ["white noise", noise], + ] as const) { + test(`${name} contains no speech`, async () => { + const vad = await loadSileroVad({ localOnly: false }); + const probs = await vad.speechProbabilities(audio); + expect(probs.length).toBe(Math.ceil(audio.length / 512)); + expect(getSpeechTimestamps(probs, audio.length)).toEqual([]); + }); + } + + test("en_aiproblem: speech is found and covers most of the clip", async () => { + const audio = await decodeClip("en_aiproblem.m4a"); + const vad = await loadSileroVad({ localOnly: false }); + const spans = getSpeechTimestamps(await vad.speechProbabilities(audio), audio.length); + expect(spans.length).toBeGreaterThan(0); + const covered = spans.reduce((n, s) => n + s.end - s.start, 0); + expect(covered / audio.length).toBeGreaterThan(0.8); + }); +});