103 lines
4.3 KiB
TypeScript
103 lines
4.3 KiB
TypeScript
/**
|
||
* Parakeet vs Whisper Large v3 Turbo benchmark (D-WS7a) — CPU only.
|
||
*
|
||
* flock /tmp/whisper-heavy.lock bun scripts/bench-parakeet.ts [parakeet-node|parakeet-web|turbo ...]
|
||
*
|
||
* Whole-file WER (same normalization as tests/live-real-model.test.ts) and
|
||
* wall-clock per English clip in tests/audio_files. Parakeet runs the
|
||
* worker's own `createParakeetEngine` on onnxruntime-node (CPU EP) and on
|
||
* onnxruntime-web (WASM, the browser engine's baseline); Turbo runs the
|
||
* Transformers.js pipeline with the worker's file-flow options on
|
||
* onnxruntime-node at its WASM-default q4 dtype. Weights come from / go to
|
||
* node_modules/@huggingface/transformers/.cache (Parakeet must already be
|
||
* there — `PARAKEET_REAL=1 bun test tests/parakeet-real-model.test.ts`
|
||
* downloads it). Wall-clock excludes model load.
|
||
*/
|
||
|
||
import { join } from "node:path";
|
||
import ffmpegPath from "ffmpeg-static";
|
||
import { pipeline } from "@huggingface/transformers";
|
||
import { PARAKEET_MODEL_ID, resolveDefaultDtype } from "../src/lib/models";
|
||
import { PARAKEET_VOCAB_FILE, parakeetFileName } from "../src/lib/parakeet";
|
||
import { createParakeetEngine, parakeetSessionOptions, type OrtLike } from "../src/worker/parakeetEngine";
|
||
import { wordErrorRate } from "../tests/dsp";
|
||
|
||
const ROOT = join(import.meta.dir, "..");
|
||
const AUDIO = join(ROOT, "tests/audio_files");
|
||
const PARAKEET_DIR = join(ROOT, "node_modules/@huggingface/transformers/.cache", PARAKEET_MODEL_ID);
|
||
const CLIPS = ["en_aiproblem.m4a", "en_alien.opus"];
|
||
const THREADS = 12;
|
||
|
||
const normalize = (text: string) =>
|
||
text.toLowerCase().replace(/['’]/g, "").replace(/[.,!?;:"“”‘‚«»()\-–—…]/g, " ").replace(/\s+/g, " ").trim();
|
||
|
||
async function decode(file: string): Promise<Float32Array> {
|
||
const proc = Bun.spawn(
|
||
[ffmpegPath as unknown as string, "-hide_banner", "-loglevel", "error", "-i", join(AUDIO, file),
|
||
"-ac", "1", "-ar", "16000", "-f", "f32le", "-"],
|
||
{ stdout: "pipe" },
|
||
);
|
||
return new Float32Array(await new Response(proc.stdout).arrayBuffer());
|
||
}
|
||
|
||
type Transcribe = (audio: Float32Array) => Promise<string>;
|
||
|
||
async function parakeet(runtime: "node" | "web"): Promise<Transcribe> {
|
||
let ort: OrtLike;
|
||
let options = parakeetSessionOptions("wasm");
|
||
if (runtime === "web") {
|
||
const web = await import("onnxruntime-web");
|
||
web.env.wasm.numThreads = THREADS;
|
||
ort = web as unknown as OrtLike;
|
||
} else {
|
||
ort = (await import("onnxruntime-node")) as unknown as OrtLike;
|
||
const cpu = { executionProviders: ["cpu"], intraOpNumThreads: THREADS, graphOptimizationLevel: "all" };
|
||
options = { encoder: cpu, decoder: cpu };
|
||
}
|
||
const read = async (file: string) => new Uint8Array(await Bun.file(join(PARAKEET_DIR, file)).arrayBuffer());
|
||
const engine = await createParakeetEngine(
|
||
ort,
|
||
{
|
||
encoder: await read(parakeetFileName("encoder", "q8")),
|
||
decoder: await read(parakeetFileName("decoder", "q8")),
|
||
vocab: await Bun.file(join(PARAKEET_DIR, PARAKEET_VOCAB_FILE)).text(),
|
||
},
|
||
options,
|
||
);
|
||
return async (audio) => (await engine.transcribe(audio)).text;
|
||
}
|
||
|
||
async function turbo(): Promise<Transcribe> {
|
||
const id = "onnx-community/whisper-large-v3-turbo";
|
||
const asr = await pipeline("automatic-speech-recognition", id, { dtype: resolveDefaultDtype(id, "wasm") });
|
||
return async (audio) => {
|
||
const output = await asr(audio, {
|
||
chunk_length_s: 30, stride_length_s: 5, return_timestamps: true, max_new_tokens: 440,
|
||
language: "en", task: "transcribe",
|
||
});
|
||
return ((Array.isArray(output) ? output[0] : output) as { text?: string }).text ?? "";
|
||
};
|
||
}
|
||
|
||
const engines: Record<string, () => Promise<Transcribe>> = {
|
||
"parakeet-node": () => parakeet("node"),
|
||
"parakeet-web": () => parakeet("web"),
|
||
turbo,
|
||
};
|
||
|
||
for (const name of process.argv.slice(2).length ? process.argv.slice(2) : Object.keys(engines)) {
|
||
const transcribe = await engines[name]!();
|
||
for (const clip of CLIPS) {
|
||
const audio = await decode(clip);
|
||
const reference = await Bun.file(join(AUDIO, clip.replace(/\.[^.]+$/, ".txt"))).text();
|
||
const start = performance.now();
|
||
const text = await transcribe(audio);
|
||
const seconds = (performance.now() - start) / 1000;
|
||
const wer = wordErrorRate(normalize(reference), normalize(text));
|
||
console.log(
|
||
`${name}\t${clip}\taudio ${(audio.length / 16000).toFixed(1)} s\twall ${seconds.toFixed(1)} s\tWER ${wer.toFixed(3)}`,
|
||
);
|
||
}
|
||
}
|
||
process.exit(0);
|