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,
} from "../src/lib/translation";
import type { TranscriptChunk } from "../src/worker/protocol";

// Pure planning for transcript translation (D-WS4b) — no model, no network.

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