From a0748b5fa6851d4df61c23768871085cd67a0b99 Mon Sep 17 00:00:00 2001 From: Lucas Winther Date: Thu, 29 Oct 2026 22:10:51 +0100 Subject: [PATCH] Translate transcripts with Opus-MT in a dedicated translation worker Co-Authored-By: Claude Opus 5.5 --- scripts/bench-translate.ts | 53 +++++++++++++++ scripts/build.ts | 15 +++- src/index.ts | 10 ++- src/lib/translation.ts | 63 ++++++++++++++++- src/worker/opusMt.ts | 59 ++++++++++++++++ src/worker/translate.worker.ts | 98 ++++++++++++++++++++++++++ src/worker/translateProtocol.ts | 106 +++++++++++++++++++++++++++++ tests/translate-real-model.test.ts | 97 ++++++++++++++++++++++++++ tests/translateProtocol.test.ts | 49 +++++++++++++ tests/translation.test.ts | 70 +++++++++++++++++++ 10 files changed, 615 insertions(+), 5 deletions(-) create mode 100644 scripts/bench-translate.ts create mode 100644 src/worker/opusMt.ts create mode 100644 src/worker/translate.worker.ts create mode 100644 src/worker/translateProtocol.ts create mode 100644 tests/translate-real-model.test.ts create mode 100644 tests/translateProtocol.test.ts diff --git a/scripts/bench-translate.ts b/scripts/bench-translate.ts new file mode 100644 index 0000000..12f8ac6 --- /dev/null +++ b/scripts/bench-translate.ts @@ -0,0 +1,53 @@ +/** + * 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(); +} diff --git a/scripts/build.ts b/scripts/build.ts index 702ed38..0bd50ef 100644 --- a/scripts/build.ts +++ b/scripts/build.ts @@ -3,7 +3,8 @@ * * Emits a self-contained `dist/`: * 1. `bun build src/index.html` → hashed app-shell assets - * 2. inference worker bundle at the stable /inference.worker.js URL + * 2. inference + translate worker bundles at stable /inference.worker.js and + * /translate.worker.js URLs * 3. manifest + icons + onnxruntime .wasm/.mjs (dist/ort/, offline inference) * 4. sw.js with the precache manifest + a content-hash version injected * @@ -53,6 +54,15 @@ const workerBuild = await Bun.build({ }); const workerScript = await workerBuild.outputs[0]!.text(); await Bun.write(join(DIST, "inference.worker.js"), workerScript); +// Transcript-translation worker (D-WS4b), same mechanism. +const translateWorkerBuild = await Bun.build({ + entrypoints: [join(ROOT, "src/worker/translate.worker.ts")], + target: "browser", + format: "esm", + minify: true, +}); +const translateWorkerScript = await translateWorkerBuild.outputs[0]!.text(); +await Bun.write(join(DIST, "translate.worker.js"), translateWorkerScript); // 3. Static PWA assets + onnxruntime wasm binaries (offline inference — the // worker loads them from /ort/ instead of the jsdelivr CDN). The manifest and @@ -89,6 +99,7 @@ const precachePaths = precachePathsFromOutputs( shellBuild.outputs.map((output) => relative(DIST, output.path)), [ "inference.worker.js", + "translate.worker.js", "manifest.webmanifest", ...ICON_SIZES.map((size) => `icons/icon-${size}.png`), ...fontFiles.map((name) => `fonts/${name}`), @@ -98,7 +109,7 @@ const fontBytes = await Promise.all( fontFiles.map((name) => Bun.file(join(FONTS_DIR, name)).bytes()), ); const version = Bun.hash( - JSON.stringify(precachePaths) + swSource + workerScript + fontBytes.map((b) => Bun.hash(b)).join(), + JSON.stringify(precachePaths) + swSource + workerScript + translateWorkerScript + fontBytes.map((b) => Bun.hash(b)).join(), ).toString(36); await Bun.write(join(DIST, "sw.js"), injectPrecacheManifest(swSource, precachePaths, version)); diff --git a/src/index.ts b/src/index.ts index 328d7f3..e657557 100644 --- a/src/index.ts +++ b/src/index.ts @@ -223,7 +223,8 @@ async function productionServeOptions(): Promise { * Bun's bundler does not currently follow `new Worker(...)` graphs (verified * with Bun 1.3: both `bun build` and the dev server leave the call untouched), * so worker-like entrypoints are bundled at startup and served from stable - * URLs. The UI spawns `new Worker("/inference.worker.js")`; the service + * URLs. The UI spawns `new Worker("/inference.worker.js")` (and + * "/translate.worker.js" for transcript translation, D-WS4b); the service * worker registers as "/sw.js". */ async function bundleEntrypoint(relativePath: string): Promise { @@ -249,6 +250,8 @@ async function developmentServeOptions(): Promise { const { default: index } = await import("./index.html"); const workerScript = await bundleEntrypoint("./worker/inference.worker.ts"); + // Transcript translation runs in its own worker (D-WS4b). + const translateWorkerScript = await bundleEntrypoint("./worker/translate.worker.ts"); // Dev SW: empty precache list — assets are unhashed here, so the shell is // runtime-cached network-first instead (see src/pwa/sw.ts). const swScript = injectPrecacheManifest( @@ -283,6 +286,11 @@ async function developmentServeOptions(): Promise { if (pathname.endsWith("/inference.worker.js")) { return withIsolationHeaders(new Response(workerScript, { headers: scriptHeaders })); } + if (pathname.endsWith("/translate.worker.js")) { + return withIsolationHeaders( + new Response(translateWorkerScript, { headers: scriptHeaders }), + ); + } if (pathname.endsWith("/sw.js")) { return withIsolationHeaders( new Response(swScript, { headers: { ...scriptHeaders, "Cache-Control": "no-cache" } }), diff --git a/src/lib/translation.ts b/src/lib/translation.ts index eecd042..b568734 100644 --- a/src/lib/translation.ts +++ b/src/lib/translation.ts @@ -154,8 +154,13 @@ export interface SplitOptions { countTokens?: (text: string) => number; } -/** "Mr.", "Dr.", "St." — sentence-segmenter false breaks. */ -const ABBREVIATION = /^\p{Lu}\p{Ll}{0,3}\.$/u; +/** + * Title abbreviations (en + da) and single-letter initials ("J.") — the + * sentence segmenter's false breaks. Kept to an explicit list: a broad + * "short capitalised word" rule would glue real one-word sentences + * ("Yes.", "Okay.") onto the next one. + */ +const ABBREVIATION = /^(?:Mr|Mrs|Ms|Dr|Prof|St|Jr|Sr|Hr|Fru|Nr|\p{Lu})\.$/u; function sentences(text: string): string[] { const out: string[] = []; @@ -215,6 +220,60 @@ export function splitForTranslation(text: string, options: SplitOptions = {}): s ); } +// --------------------------------------------------------------------------- +// The translation loop (model injected: worker, benchmark and tests share it) +// --------------------------------------------------------------------------- + +export interface TranslateUnitsDeps extends SplitOptions { + /** Translate a batch of pieces; must answer one string per piece, in order. */ + translateBatch: (pieces: string[]) => Promise; + /** Pieces per model call (padding makes larger batches cheaper per piece). */ + batchSize?: number; + /** Called with (done, total) pieces: once at 0, then after every batch. */ + onProgress?: (done: number, total: number) => void; +} + +/** + * One piece per model call. Batching would be ~4x faster, but Transformers.js + * 4.2's batched Marian generation is broken: every row of a padded batch + * comes back as row 0's translation (measured 2026-10-29 on + * onnxruntime-node, q8 and fp32, any graph optimisation level — D-WS4b). + * Re-measure with `scripts/bench-translate.ts` after a Transformers.js + * upgrade before raising this. + */ +export const DEFAULT_TRANSLATE_BATCH_SIZE = 1; + +/** + * Translate unit texts: split each into sentence pieces, translate all + * pieces in order-preserving batches, rejoin each unit's pieces with spaces. + * Returns one translated string per input unit. + */ +export async function translateUnits( + units: readonly string[], + deps: TranslateUnitsDeps, +): Promise { + const batchSize = Math.max(1, deps.batchSize ?? DEFAULT_TRANSLATE_BATCH_SIZE); + const perUnit = units.map((unit) => splitForTranslation(unit, deps)); + const pieces = perUnit.flat(); + const translated: string[] = []; + deps.onProgress?.(0, pieces.length); + for (let start = 0; start < pieces.length; start += batchSize) { + const batch = pieces.slice(start, start + batchSize); + const out = await deps.translateBatch(batch); + if (out.length !== batch.length) { + throw new Error(`Translation returned ${out.length} sentences for ${batch.length}.`); + } + translated.push(...out.map((text) => text.trim())); + deps.onProgress?.(translated.length, pieces.length); + } + let next = 0; + return perUnit.map((unitPieces) => { + const joined = translated.slice(next, next + unitPieces.length).join(" "); + next += unitPieces.length; + return joined; + }); +} + // --------------------------------------------------------------------------- // Assembly // --------------------------------------------------------------------------- diff --git a/src/worker/opusMt.ts b/src/worker/opusMt.ts new file mode 100644 index 0000000..8bbee00 --- /dev/null +++ b/src/worker/opusMt.ts @@ -0,0 +1,59 @@ +/** + * Opus-MT (MarianMT) through Transformers.js — the model half of transcript + * translation (D-WS4b). Shared by the browser translate worker and the + * onnxruntime-node measurement script / real-model test, so all three run + * the exact same pipeline options. + */ + +import { pipeline, type ProgressInfo, type TranslationPipeline } from "@huggingface/transformers"; + +/** Hard input limit of every Opus-MT checkpoint (max_position_embeddings). */ +export const OPUS_MT_MAX_POSITIONS = 512; + +export interface OpusMtTranslator { + /** One translated string per piece, same order. */ + translateBatch(pieces: string[]): Promise; + /** Real SentencePiece token count (incl. ) for the splitter's budget. */ + countTokens(text: string): number; + dispose(): Promise; +} + +export interface LoadOpusMtOptions { + /** "wasm" in the browser; omit under onnxruntime-node (defaults to CPU). */ + device?: "wasm"; + progressCallback?: (info: ProgressInfo) => void; +} + +export async function loadOpusMt( + modelId: string, + options: LoadOpusMtOptions = {}, +): Promise { + const translator = (await pipeline("translation", modelId, { + // q8: 113 MB per direction (encoder + merged decoder + tokenizer); the + // fp32 export is ~4x larger for no audible gain on transcripts. + dtype: "q8", + ...(options.device && { device: options.device }), + // D-014: the bundled onnxruntime-web's level-2 QDQ transform rejects + // 8-bit decoders with a tied embedding weight — Marian ties them too. + session_options: { graphOptimizationLevel: "basic" as const }, + ...(options.progressCallback && { progress_callback: options.progressCallback }), + })) as TranslationPipeline; + + const countTokens = (text: string) => translator.tokenizer.encode(text).length; + + return { + countTokens, + async translateBatch(pieces) { + if (pieces.length === 0) return []; + // Cap output length relative to the longest input: a translation is + // rarely more than ~2x its source in tokens, and the cap stops a rare + // repetition loop from running to 512 tokens. + const longest = Math.max(...pieces.map(countTokens)); + const output = (await translator(pieces, { + max_new_tokens: Math.min(OPUS_MT_MAX_POSITIONS, 3 * longest + 16), + })) as { translation_text: string }[]; + return output.map((entry) => entry.translation_text); + }, + dispose: () => translator.dispose(), + }; +} diff --git a/src/worker/translate.worker.ts b/src/worker/translate.worker.ts new file mode 100644 index 0000000..6e0ba4e --- /dev/null +++ b/src/worker/translate.worker.ts @@ -0,0 +1,98 @@ +/** + * Transcript-translation Web Worker (D-WS4b): Opus-MT through + * Transformers.js, off the main thread (D-003). + * + * A dedicated worker rather than a message type on the Whisper worker: + * - its own onnxruntime instance gets the full WASM thread pool. ort fixes + * the pool at first init (D-023), and a Whisper worker that initialised + * ort for WebGPU/WebNN runs WASM single-threaded for life; + * - cancelling a translation terminates only this worker — the loaded + * Whisper pipeline survives, and a Whisper cancel/restart (D-018/D-025) + * never kills a translation; + * - the model is downloaded only when the user asks for a translation and + * cached by Transformers.js in the Cache API like Whisper weights + * (D-005/D-020), so a cached model works offline. + * WASM only for now: Opus-MT is small, and WebGPU behaviour can't be + * measured in this environment. + */ + +import { env, type ProgressInfo } from "@huggingface/transformers"; +import { getTranslationModelById, translateUnits } from "../lib/translation"; +import { loadOpusMt, type OpusMtTranslator } from "./opusMt"; +import { + DownloadTally, + isTranslateRequest, + type TranslateRequest, + type TranslateWorkerMessage, +} from "./translateProtocol"; + +// Same-origin onnxruntime binaries (offline, COEP) — see inference.worker.ts. +if (env.backends.onnx.wasm) { + env.backends.onnx.wasm.wasmPaths = "./ort/"; + // D-023: with cross-origin isolation, use most of the machine. + env.backends.onnx.wasm.numThreads = globalThis.crossOriginIsolated + ? Math.max(1, Math.min((navigator.hardwareConcurrency || 4) - 1, 12)) + : 1; +} + +interface DedicatedWorkerScope { + postMessage(message: TranslateWorkerMessage): void; + addEventListener(type: "message", listener: (event: MessageEvent) => void): void; +} +const scope = self as unknown as DedicatedWorkerScope; + +/** One loaded model at a time; switching direction disposes the other. */ +let active: { modelId: string; promise: Promise } | null = null; + +function load(modelId: string, onProgress: (info: ProgressInfo) => void): Promise { + if (active?.modelId === modelId) return active.promise; + const previous = active; + const promise = (async () => { + if (previous) await (await previous.promise.catch(() => null))?.dispose(); + return loadOpusMt(modelId, { device: "wasm", progressCallback: onProgress }); + })(); + active = { modelId, promise }; + promise.catch(() => { + if (active?.promise === promise) active = null; // allow a clean retry + }); + return promise; +} + +async function handle(request: TranslateRequest): Promise { + const { requestId } = request; + if (!getTranslationModelById(request.modelId)) { + throw new Error(`Unknown translation model: ${request.modelId}`); + } + const tally = new DownloadTally(); + const translator = await load(request.modelId, (info) => { + if (info.status === "progress") { + const { loaded, total } = tally.update(info.file, info.loaded, info.total); + scope.postMessage({ type: "translate-download", requestId, loaded, total }); + } + }); + const texts = await translateUnits(request.units, { + translateBatch: (pieces) => translator.translateBatch(pieces), + countTokens: translator.countTokens, + onProgress: (done, total) => + scope.postMessage({ type: "translate-progress", requestId, done, total }), + }); + scope.postMessage({ type: "translate-result", requestId, texts }); +} + +// Requests run one at a time, in arrival order: the UI only ever has one in +// flight, and serialising keeps two passes from sharing the thread pool. +let queue: Promise = Promise.resolve(); + +scope.addEventListener("message", (event) => { + const data = event.data; + if (!isTranslateRequest(data)) return; + queue = queue.then(() => + handle(data).catch((error: unknown) => { + scope.postMessage({ + type: "translate-error", + requestId: data.requestId, + message: error instanceof Error ? error.message : String(error), + }); + }), + ); +}); diff --git a/src/worker/translateProtocol.ts b/src/worker/translateProtocol.ts new file mode 100644 index 0000000..e26c105 --- /dev/null +++ b/src/worker/translateProtocol.ts @@ -0,0 +1,106 @@ +/** + * Typed postMessage protocol between the UI and the transcript-translation + * worker (`src/worker/translate.worker.ts`, D-WS4b). Kept apart from the + * Whisper worker's protocol: separate worker, separate lifetime. Free of + * browser globals so `bun test` covers the guards. + */ + +/** UI → worker: translate these unit texts with this Opus-MT model. */ +export interface TranslateRequest { + type: "translate"; + /** Correlates every reply; the UI ignores replies for stale ids. */ + requestId: string; + /** One of `TRANSLATION_MODEL_IDS` (src/lib/translation.ts). */ + modelId: string; + /** Unit texts (`buildTranslationUnits`), answered 1:1 in `texts`. */ + units: string[]; +} + +/** Model files being fetched (first use) — bytes summed across files. */ +export interface TranslateDownloadMessage { + type: "translate-download"; + requestId: string; + loaded: number; + total: number; +} + +/** Model ready; `done` of `total` sentence pieces translated so far. */ +export interface TranslateProgressMessage { + type: "translate-progress"; + requestId: string; + done: number; + total: number; +} + +export interface TranslateResultMessage { + type: "translate-result"; + requestId: string; + /** One translated string per request unit, same order. */ + texts: string[]; +} + +export interface TranslateErrorMessage { + type: "translate-error"; + requestId: string; + message: string; +} + +export type TranslateWorkerMessage = + | TranslateDownloadMessage + | TranslateProgressMessage + | TranslateResultMessage + | TranslateErrorMessage; + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +const isNumber = (value: unknown): value is number => + typeof value === "number" && Number.isFinite(value); + +export function isTranslateRequest(value: unknown): value is TranslateRequest { + return ( + isRecord(value) && + value.type === "translate" && + typeof value.requestId === "string" && + typeof value.modelId === "string" && + Array.isArray(value.units) && + value.units.every((unit) => typeof unit === "string") + ); +} + +export function isTranslateWorkerMessage(value: unknown): value is TranslateWorkerMessage { + if (!isRecord(value) || typeof value.requestId !== "string") return false; + switch (value.type) { + case "translate-download": + return isNumber(value.loaded) && isNumber(value.total); + case "translate-progress": + return isNumber(value.done) && isNumber(value.total); + case "translate-result": + return Array.isArray(value.texts) && value.texts.every((t) => typeof t === "string"); + case "translate-error": + return typeof value.message === "string"; + default: + return false; + } +} + +/** + * Sums Transformers.js per-file download progress into one byte count for + * the UI. A file's total is only known once its download starts, so the + * total grows as files start — fine for a percentage that never goes back. + */ +export class DownloadTally { + private readonly files = new Map(); + + update(file: string, loaded: number, total: number): { loaded: number; total: number } { + this.files.set(file, { loaded, total }); + let sumLoaded = 0; + let sumTotal = 0; + for (const entry of this.files.values()) { + sumLoaded += entry.loaded; + sumTotal += entry.total; + } + return { loaded: sumLoaded, total: sumTotal }; + } +} diff --git a/tests/translate-real-model.test.ts b/tests/translate-real-model.test.ts new file mode 100644 index 0000000..b2b156b --- /dev/null +++ b/tests/translate-real-model.test.ts @@ -0,0 +1,97 @@ +import { afterAll, beforeAll, describe, expect, test } from "bun:test"; +import { join } from "node:path"; +import { + assembleTranslation, + buildTranslationUnits, + getTranslationModel, + translateUnits, +} from "../src/lib/translation"; +import { loadOpusMt, type OpusMtTranslator } from "../src/worker/opusMt"; + +/** + * Real Opus-MT translation on onnxruntime-node (D-WS4b), opt-in: + * + * MT_INTEGRATION=1 flock /tmp/whisper-heavy.lock bun test tests/translate-real-model.test.ts + * + * Downloads Xenova/opus-mt-en-da and -da-en (~113 MB each, q8) into the + * Transformers.js cache on first run; ~3 min on a 16-core CPU. Same loader + * and loop as the browser worker. + */ +const ENABLED = process.env.MT_INTEGRATION === "1"; +const AUDIO = join(import.meta.dir, "audio_files"); +const words = (text: string) => text.split(/\s+/).filter(Boolean); + +describe.skipIf(!ENABLED)("Opus-MT on a real transcript", () => { + let enDa: OpusMtTranslator; + let daEn: OpusMtTranslator; + + beforeAll(async () => { + enDa = await loadOpusMt(getTranslationModel("en", "da")!.id); + daEn = await loadOpusMt(getTranslationModel("da", "en")!.id); + }, 600_000); + + afterAll(async () => { + await enDa?.dispose(); + await daEn?.dispose(); + }); + + const run = (translator: OpusMtTranslator, units: string[]) => + translateUnits(units, { + countTokens: translator.countTokens, + translateBatch: (pieces) => translator.translateBatch(pieces), + }); + + test( + "all of en_alien → Danish: one translation per line, none empty, none copied from a neighbour", + async () => { + const text = await Bun.file(join(AUDIO, "en_alien.txt")).text(); + const units = buildTranslationUnits({ text, chunks: [] }); + const started = performance.now(); + const out = await run(enDa, units.map((unit) => unit.text)); + console.info(`[mt] en_alien en→da: ${((performance.now() - started) / 1000).toFixed(1)} s`); + const assembled = assembleTranslation(units, out); + expect(out).toHaveLength(units.length); + for (const line of out) expect(line.trim()).not.toBe(""); + // Guards the batched-generation bug (every row = row 0's output). + const distinctSources = new Set(units.map((unit) => unit.text)).size; + expect(new Set(out).size).toBeGreaterThan(distinctSources * 0.9); + // Nothing dropped: a translation keeps roughly the source's length. + const ratio = words(assembled.text).length / words(text).length; + expect(ratio).toBeGreaterThan(0.6); + expect(ratio).toBeLessThan(1.5); + }, + 600_000, + ); + + test( + "an unpunctuated unit far beyond 512 tokens is translated whole (Review Focus 5)", + async () => { + const text = await Bun.file(join(AUDIO, "en_alien.txt")).text(); + // Live-mode-like run-on: no sentence punctuation, one single unit. + const runOn = words(text.replace(/[.!?…]/g, "")).slice(0, 900).join(" "); + expect(enDa.countTokens(runOn)).toBeGreaterThan(512); + const [out] = await run(enDa, [runOn]); + const ratio = words(out ?? "").length / words(runOn).length; + expect(ratio).toBeGreaterThan(0.6); + expect(ratio).toBeLessThan(1.5); + }, + 600_000, + ); + + test( + "da_festuge → English keeps names and dates, timestamps carried over", + async () => { + const text = await Bun.file(join(AUDIO, "da_festuge.txt")).text(); + const chunks = text + .split("\n") + .filter((line) => line.trim() !== "") + .map((line, index) => ({ text: ` ${line}`, timestamp: [index * 5, index * 5 + 5] as [number, number] })); + const units = buildTranslationUnits({ text, chunks }); + const result = assembleTranslation(units, await run(daEn, units.map((unit) => unit.text))); + expect(result.text).toContain("Aarhus"); + expect(result.text).toMatch(/August/); + expect(result.chunks.map((chunk) => chunk.timestamp)).toEqual(chunks.map((chunk) => chunk.timestamp)); + }, + 600_000, + ); +}); diff --git a/tests/translateProtocol.test.ts b/tests/translateProtocol.test.ts new file mode 100644 index 0000000..2f88f01 --- /dev/null +++ b/tests/translateProtocol.test.ts @@ -0,0 +1,49 @@ +import { describe, expect, test } from "bun:test"; +import { + DownloadTally, + isTranslateRequest, + isTranslateWorkerMessage, +} from "../src/worker/translateProtocol"; + +describe("isTranslateRequest", () => { + test("accepts a well-formed request", () => { + expect( + isTranslateRequest({ type: "translate", requestId: "t-1", modelId: "Xenova/opus-mt-da-en", units: ["Hej."] }), + ).toBe(true); + }); + + test("rejects missing ids and non-string units", () => { + expect(isTranslateRequest({ type: "translate", modelId: "m", units: [] })).toBe(false); + expect(isTranslateRequest({ type: "translate", requestId: "t", modelId: "m", units: [1] })).toBe(false); + expect(isTranslateRequest({ type: "load", requestId: "t", modelId: "m", units: [] })).toBe(false); + expect(isTranslateRequest(null)).toBe(false); + }); +}); + +describe("isTranslateWorkerMessage", () => { + test("accepts each reply type", () => { + for (const message of [ + { type: "translate-download", requestId: "t", loaded: 1, total: 2 }, + { type: "translate-progress", requestId: "t", done: 0, total: 3 }, + { type: "translate-result", requestId: "t", texts: ["Hello."] }, + { type: "translate-error", requestId: "t", message: "boom" }, + ]) { + expect(isTranslateWorkerMessage(message)).toBe(true); + } + }); + + test("rejects malformed replies and Whisper-worker messages", () => { + expect(isTranslateWorkerMessage({ type: "translate-progress", requestId: "t", done: "1", total: 3 })).toBe(false); + expect(isTranslateWorkerMessage({ type: "translate-result", texts: [] })).toBe(false); + expect(isTranslateWorkerMessage({ type: "result", requestId: "t", text: "", chunks: [] })).toBe(false); + }); +}); + +describe("DownloadTally", () => { + test("sums bytes across files, updating each file in place", () => { + const tally = new DownloadTally(); + expect(tally.update("onnx/encoder.onnx", 10, 100)).toEqual({ loaded: 10, total: 100 }); + expect(tally.update("onnx/decoder.onnx", 5, 50)).toEqual({ loaded: 15, total: 150 }); + expect(tally.update("onnx/encoder.onnx", 100, 100)).toEqual({ loaded: 105, total: 150 }); + }); +}); diff --git a/tests/translation.test.ts b/tests/translation.test.ts index dd96443..06e3ca2 100644 --- a/tests/translation.test.ts +++ b/tests/translation.test.ts @@ -8,11 +8,73 @@ import { getTranslationModel, guessSourceLanguage, splitForTranslation, + translateUnits, } from "../src/lib/translation"; import type { TranscriptChunk } from "../src/worker/protocol"; // Pure planning for transcript translation (D-WS4b) — no model, no network. +describe("translateUnits", () => { + test("translates every piece in order, batched, and rejoins per unit", async () => { + const batches: string[][] = []; + const out = await translateUnits(["One. Two. Three.", "Four.", "Five. Six."], { + batchSize: 4, + translateBatch: async (pieces) => { + batches.push(pieces); + return pieces.map((p) => p.toUpperCase()); + }, + }); + expect(out).toEqual(["ONE. TWO. THREE.", "FOUR.", "FIVE. SIX."]); + expect(batches).toEqual([["One.", "Two.", "Three.", "Four."], ["Five.", "Six."]]); + }); + + test("defaults to one piece per call (batched Marian generation is broken in 4.2)", async () => { + const sizes: number[] = []; + await translateUnits(["One. Two. Three."], { + translateBatch: async (pieces) => { + sizes.push(pieces.length); + return pieces; + }, + }); + expect(sizes).toEqual([1, 1, 1]); + }); + + test("reports progress in pieces, starting at zero", async () => { + const progress: [number, number][] = []; + await translateUnits(["Yes. No. Maybe."], { + batchSize: 2, + translateBatch: async (pieces) => pieces, + onProgress: (done, total) => progress.push([done, total]), + }); + expect(progress).toEqual([[0, 3], [2, 3], [3, 3]]); + }); + + test("an empty unit stays empty without a model call", async () => { + let calls = 0; + const out = await translateUnits(["", "Hi."], { + translateBatch: async (pieces) => { + calls += 1; + return pieces; + }, + }); + expect(out).toEqual(["", "Hi."]); + expect(calls).toBe(1); + }); + + test("a batch answering with the wrong count is an error", async () => { + await expect( + translateUnits(["Yes. No."], { batchSize: 2, translateBatch: async () => ["only one"] }), + ).rejects.toThrow(); + }); + + test("the whole of en_alien round-trips through an echo translator with nothing dropped (Review Focus 5)", async () => { + const alien = await Bun.file(join(import.meta.dir, "audio_files/en_alien.txt")).text(); + const units = buildTranslationUnits({ text: alien, chunks: [] }).map((u) => u.text); + const out = await translateUnits(units, { translateBatch: async (pieces) => pieces }); + expect(words(out.join(" "))).toEqual(words(alien)); + }); +}); + const words = (text: string) => text.split(/\s+/).filter(Boolean); describe("TRANSLATION_MODELS", () => { @@ -100,6 +162,14 @@ describe("splitForTranslation", () => { ]); }); + test("one-word sentences are not mistaken for abbreviations", () => { + expect(splitForTranslation("Yes. Okay. J. K. Rowling wrote it.")).toEqual([ + "Yes.", + "Okay.", + "J. K. Rowling wrote it.", + ]); + }); + test("blank input yields no pieces", () => { expect(splitForTranslation(" ")).toEqual([]); });