54 lines
2.3 KiB
TypeScript
54 lines
2.3 KiB
TypeScript
/**
|
|
* Transcript-translation speed on onnxruntime-node (CPU), D-WS4b.
|
|
*
|
|
* flock /tmp/whisper-heavy.lock bun scripts/bench-translate.ts [batchSizes…]
|
|
*
|
|
* Translates the reference transcripts in tests/audio_files/ with the same
|
|
* Opus-MT loader and translation loop the browser worker uses, once per
|
|
* batch size, and prints wall-clock time, pieces, and whether the output
|
|
* depends on the batch size (padding). Downloads ~113 MB per direction into
|
|
* the Transformers.js cache on first run. Not browser numbers: WASM in a
|
|
* tab is slower than onnxruntime-node.
|
|
*/
|
|
|
|
import { join } from "node:path";
|
|
import { buildTranslationUnits, getTranslationModel, translateUnits } from "../src/lib/translation";
|
|
import { loadOpusMt } from "../src/worker/opusMt";
|
|
|
|
const AUDIO = join(import.meta.dir, "../tests/audio_files");
|
|
const batchSizes = process.argv.slice(2).map(Number).filter((n) => n > 0);
|
|
const sizes = batchSizes.length > 0 ? batchSizes : [1, 8];
|
|
|
|
for (const [file, source, target] of [
|
|
["en_alien.txt", "en", "da"],
|
|
["da_festuge.txt", "da", "en"],
|
|
] as const) {
|
|
const model = getTranslationModel(source, target)!;
|
|
const text = await Bun.file(join(AUDIO, file)).text();
|
|
const units = buildTranslationUnits({ text, chunks: [] }).map((unit) => unit.text);
|
|
const loadStart = performance.now();
|
|
const translator = await loadOpusMt(model.id);
|
|
const loadMs = performance.now() - loadStart;
|
|
console.log(`\n${file} (${text.length} chars, ${units.length} lines) ${model.label} — load ${loadMs.toFixed(0)} ms`);
|
|
let baseline: string[] | null = null;
|
|
for (const batchSize of sizes) {
|
|
let pieces = 0;
|
|
const start = performance.now();
|
|
const out = await translateUnits(units, {
|
|
batchSize,
|
|
countTokens: translator.countTokens,
|
|
translateBatch: (batch) => translator.translateBatch(batch),
|
|
onProgress: (_done, total) => (pieces = total),
|
|
});
|
|
const seconds = (performance.now() - start) / 1000;
|
|
const same = baseline === null ? "baseline" : JSON.stringify(out) === JSON.stringify(baseline) ? "identical" : "differs";
|
|
baseline ??= out;
|
|
console.log(
|
|
` batch ${String(batchSize).padStart(2)}: ${seconds.toFixed(2)} s, ${pieces} pieces, ` +
|
|
`${(text.length / seconds).toFixed(0)} chars/s, output ${same}`,
|
|
);
|
|
}
|
|
console.log(` sample: ${baseline!.slice(0, 3).join(" | ")}`);
|
|
await translator.dispose();
|
|
}
|