Run Silero VAD and Whisper decodes for the long-form pipeline on Transformers.js
Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
Lucas Winther
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
19318a56f2
commit
0c3a83c23c
@@ -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<Float32Array>;
|
||||
dispose(): Promise<void>;
|
||||
}
|
||||
|
||||
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<boolean>;
|
||||
}
|
||||
|
||||
async function vadInBrowserCache(): Promise<boolean> {
|
||||
return (await getDownloadedModels([VAD_MODEL_ID])).has(VAD_MODEL_ID);
|
||||
}
|
||||
|
||||
export async function loadSileroVad(options: LoadSileroVadOptions): Promise<SileroVad> {
|
||||
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<number>) {
|
||||
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<string, number>;
|
||||
}
|
||||
|
||||
export function createWhisperDeps(
|
||||
asr: AutomaticSpeechRecognitionPipeline,
|
||||
options: { multilingual: boolean },
|
||||
): Pick<LongFormDeps, "decode" | "detectLanguage"> {
|
||||
const decode = async (audio: Float32Array, request: DecodeRequest): Promise<DecodeResult> => {
|
||||
const recorder = new LogProbRecorder();
|
||||
const generate: Record<string, unknown> = {
|
||||
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<typeof WhisperTextStreamer>[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<string> => {
|
||||
const model = asr.model as unknown as {
|
||||
generation_config: WhisperGenerationConfig;
|
||||
generate(args: Record<string, unknown>): Promise<Tensor>;
|
||||
};
|
||||
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 };
|
||||
}
|
||||
@@ -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/<name>` to mono float32 PCM at `rate` (16 kHz by default). */
|
||||
export async function decodeClip(name: string, rate = 16_000): Promise<Float32Array> {
|
||||
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/<base>.txt`. */
|
||||
export function readReference(base: string): Promise<string> {
|
||||
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();
|
||||
}
|
||||
@@ -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<boolean> {
|
||||
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);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user