Files

54 lines
2.3 KiB
TypeScript

/**
* Transcript-translation speed on onnxruntime-node (CPU), D-045.
*
* 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();
}