whisper_web/tests/translation.test.ts

import { describe, expect, test } from "bun:test";
import { join } from "node:path";
import {
  TRANSLATION_MODELS,
  assembleTranslation,
  buildTranslationUnits,
  estimateTokens,
  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", () => {
  test("offers da→en and en→da Opus-MT ONNX exports", () => {
    expect(TRANSLATION_MODELS.map((m) => m.id)).toEqual([
      "Xenova/opus-mt-da-en",
      "Xenova/opus-mt-en-da",
    ]);
    expect(getTranslationModel("da", "en")?.id).toBe("Xenova/opus-mt-da-en");
    expect(getTranslationModel("en", "da")?.id).toBe("Xenova/opus-mt-en-da");
    for (const model of TRANSLATION_MODELS) {
      // q8 encoder + merged decoder + tokenizer, measured on the Hub.
      expect(model.approxBytes).toBeGreaterThan(100e6);
      expect(model.approxBytes).toBeLessThan(130e6);
    }
  });
});

describe("buildTranslationUnits", () => {
  test("merges chunks until a sentence ends, keeping the merged time span", () => {
    const chunks: TranscriptChunk[] = [
      { text: " So the plan is", timestamp: [0, 2] },
      { text: " to ship on Friday.", timestamp: [2, 4.5] },
      { text: " Any questions?", timestamp: [5, 6] },
    ];
    expect(buildTranslationUnits({ text: "", chunks })).toEqual([
      { text: "So the plan is to ship on Friday.", timestamp: [0, 4.5] },
      { text: "Any questions?", timestamp: [5, 6] },
    ]);
  });

  test("an unterminated tail still becomes a unit, with its null end kept", () => {
    const chunks: TranscriptChunk[] = [
      { text: " Done.", timestamp: [0, 1] },
      { text: " and then", timestamp: [1, null] },
    ];
    expect(buildTranslationUnits({ text: "", chunks })).toEqual([
      { text: "Done.", timestamp: [0, 1] },
      { text: "and then", timestamp: [1, null] },
    ]);
  });

  test("unpunctuated chunks are capped so cues stay subtitle-sized", () => {
    const chunks: TranscriptChunk[] = Array.from({ length: 20 }, (_, i) => ({
      text: ` word${i}`,
      timestamp: [i, i + 1] as [number, number],
    }));
    const units = buildTranslationUnits({ text: "", chunks }, 8);
    expect(units.map((u) => u.timestamp)).toEqual([
      [0, 8],
      [8, 16],
      [16, 20],
    ]);
  });

  test("empty chunks are skipped", () => {
    const chunks: TranscriptChunk[] = [
      { text: "  ", timestamp: [0, 1] },
      { text: " Hi.", timestamp: [1, 2] },
    ];
    expect(buildTranslationUnits({ text: "", chunks })).toEqual([
      { text: "Hi.", timestamp: [1, 2] },
    ]);
  });

  test("without chunks, each non-empty line is an untimed unit", () => {
    expect(buildTranslationUnits({ text: "Hej.\n\nHvordan går det?\n", chunks: [] })).toEqual([
      { text: "Hej.", timestamp: null },
      { text: "Hvordan går det?", timestamp: null },
    ]);
  });
});

describe("splitForTranslation", () => {
  test("every sentence is its own piece (Opus-MT drops trailing sentences otherwise)", () => {
    expect(
      splitForTranslation("My name is Lucas and I live in Copenhagen. How are you? Fine!"),
    ).toEqual(["My name is Lucas and I live in Copenhagen.", "How are you?", "Fine!"]);
  });

  test("abbreviation-only fragments stay with the following sentence", () => {
    expect(splitForTranslation("Mr. Smith went home. He slept.")).toEqual([
      "Mr. Smith went home.",
      "He slept.",
    ]);
  });

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

  test("an unpunctuated run-on is wrapped under the budget, words intact and in order", () => {
    const runOn = Array.from({ length: 300 }, (_, i) => `ord${i}`).join(" ");
    const pieces = splitForTranslation(runOn, { maxTokens: 40, countTokens: (s) => words(s).length });
    expect(pieces.length).toBeGreaterThan(1);
    for (const piece of pieces) expect(words(piece).length).toBeLessThanOrEqual(40);
    expect(pieces.join(" ")).toBe(runOn);
  });

  test("a long sentence prefers comma boundaries before word wrapping", () => {
    const clause = "one two three four five six";
    const sentence = `${clause}, ${clause}, ${clause}, ${clause}.`;
    const pieces = splitForTranslation(sentence, { maxTokens: 14, countTokens: (s) => words(s).length });
    expect(pieces).toEqual([`${clause}, ${clause},`, `${clause}, ${clause}.`]);
  });

  test("a transcript far beyond Opus-MT's 512-token input is split with nothing dropped, in order (Review Focus 5)", async () => {
    const alien = await Bun.file(join(import.meta.dir, "audio_files/en_alien.txt")).text();
    expect(estimateTokens(alien)).toBeGreaterThan(512);
    const pieces = splitForTranslation(alien, { maxTokens: 200 });
    for (const piece of pieces) expect(estimateTokens(piece)).toBeLessThanOrEqual(200);
    expect(words(pieces.join(" "))).toEqual(words(alien));
  });
});

describe("estimateTokens", () => {
  test("over-estimates rather than under-estimates SentencePiece counts", () => {
    // Measured with the Xenova/opus-mt-en-da tokenizer: this sentence is 12 tokens (incl. </s>).
    expect(estimateTokens("My name is Lucas and I live in Copenhagen.")).toBeGreaterThanOrEqual(12);
    expect(estimateTokens("")).toBe(0);
  });
});

describe("assembleTranslation", () => {
  test("pairs translated units with their time spans", () => {
    const units = [
      { text: "Hej.", timestamp: [0, 1] as [number, number] },
      { text: "Farvel.", timestamp: [2, null] as [number, null] },
    ];
    expect(assembleTranslation(units, [" Hello. ", "Goodbye."])).toEqual({
      text: "Hello. Goodbye.",
      chunks: [
        { text: "Hello.", timestamp: [0, 1] },
        { text: "Goodbye.", timestamp: [2, null] },
      ],
    });
  });

  test("untimed units join as lines and produce no chunks", () => {
    const units = [
      { text: "Hej.", timestamp: null },
      { text: "Farvel.", timestamp: null },
    ];
    expect(assembleTranslation(units, ["Hello.", "Goodbye."])).toEqual({
      text: "Hello.\nGoodbye.",
      chunks: [],
    });
  });

  test("a count mismatch is an error, never a silently shifted transcript", () => {
    expect(() => assembleTranslation([{ text: "a", timestamp: null }], [])).toThrow();
  });
});

describe("guessSourceLanguage", () => {
  test("trusts an explicit Whisper language", () => {
    expect(guessSourceLanguage("whatever", "da")).toBe("da");
    expect(guessSourceLanguage("whatever", "en")).toBe("en");
  });

  test("detects Danish from its letters and common words", () => {
    expect(guessSourceLanguage("På Aarhus Festuges sekretariat har vi løbet rigtig stærkt.")).toBe("da");
    expect(guessSourceLanguage("Det er ikke noget med at vi skal")).toBe("da");
  });

  test("defaults to English", () => {
    expect(guessSourceLanguage("Guys, I have been addicted to this game.")).toBe("en");
    expect(guessSourceLanguage("", "auto")).toBe("en");
  });
});