Run Silero VAD and Whisper decodes for the long-form pipeline on Transformers.js
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Lucas WintherandClaude Opus 5.5 committed
1 parent
19318a56f2
commit
0c3a83c23c
3 files changed
+356
No files matched your search
@@ -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 };
|
||||
}
|
||||
Reference in new issue
Block a user