Translate transcripts with Opus-MT in a dedicated translation worker

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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,
+  );
+});

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 });
+  });
+});

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([]);
   });