333 lines
12 KiB
TypeScript
333 lines
12 KiB
TypeScript
/**
|
||
* Text translation of a finished transcript (D-WS4b) — the pure half.
|
||
*
|
||
* Opus-MT (MarianMT) runs in `src/worker/translate.worker.ts`; everything
|
||
* here is planning that needs no model: which checkpoint serves a language
|
||
* pair, how a transcript becomes translation units that keep their
|
||
* timestamps, and how a unit is cut into sentence pieces under the model's
|
||
* input budget. No browser globals, so `bun test` covers it.
|
||
*
|
||
* Why per-sentence pieces: Opus-MT is trained on sentence pairs and, given
|
||
* several sentences at once, routinely drops the later ones (measured
|
||
* 2026-10-29: "…I live in Copenhagen. How are you?" came back as one
|
||
* sentence). Its encoder also tops out at 512 positions.
|
||
*/
|
||
|
||
import type { TranscriptChunk } from "../worker/protocol";
|
||
|
||
export type TranslationLanguage = "da" | "en";
|
||
|
||
export interface TranslationModel {
|
||
/** Hugging Face Hub id of the Transformers.js ONNX export. */
|
||
id: string;
|
||
source: TranslationLanguage;
|
||
target: TranslationLanguage;
|
||
/** Direction as shown in the picker, e.g. "Danish → English". */
|
||
label: string;
|
||
/** Plain name of the target language, e.g. "English". */
|
||
targetLabel: string;
|
||
/**
|
||
* One-time download at the q8 dtype the worker loads: encoder + merged
|
||
* decoder + tokenizer/configs, summed from the Hub file sizes.
|
||
*/
|
||
approxBytes: number;
|
||
}
|
||
|
||
/**
|
||
* Xenova's ONNX exports of Helsinki-NLP's Opus-MT checkpoints (Apache-2.0
|
||
* upstream). Verified to load and translate under Transformers.js 4.2.0.
|
||
*/
|
||
export const TRANSLATION_MODELS: readonly TranslationModel[] = [
|
||
{
|
||
id: "Xenova/opus-mt-da-en",
|
||
source: "da",
|
||
target: "en",
|
||
label: "Danish → English",
|
||
targetLabel: "English",
|
||
approxBytes: 113_000_000,
|
||
},
|
||
{
|
||
id: "Xenova/opus-mt-en-da",
|
||
source: "en",
|
||
target: "da",
|
||
label: "English → Danish",
|
||
targetLabel: "Danish",
|
||
approxBytes: 113_000_000,
|
||
},
|
||
];
|
||
|
||
export const TRANSLATION_MODEL_IDS: readonly string[] = TRANSLATION_MODELS.map((m) => m.id);
|
||
|
||
export function getTranslationModel(
|
||
source: TranslationLanguage,
|
||
target: TranslationLanguage,
|
||
): TranslationModel | undefined {
|
||
return TRANSLATION_MODELS.find((m) => m.source === source && m.target === target);
|
||
}
|
||
|
||
/** Lookup by Hub id (storage panel names, worker request validation). */
|
||
export function getTranslationModelById(id: string): TranslationModel | undefined {
|
||
return TRANSLATION_MODELS.find((m) => m.id === id);
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Units: transcript → sentence-ending groups of chunks with their time span
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/** A stretch of transcript translated as one cue. */
|
||
export interface TranslationUnit {
|
||
text: string;
|
||
/** Merged [start, end] of the source chunks; null for untimed text. */
|
||
timestamp: [number, number | null] | null;
|
||
}
|
||
|
||
const SENTENCE_END = /[.!?…]["'”’)\]]*$/;
|
||
|
||
/** Default cap on chunks merged into one unit, so an unpunctuated transcript
|
||
* still yields subtitle-sized cues rather than one giant one. */
|
||
export const MAX_CHUNKS_PER_UNIT = 8;
|
||
|
||
/**
|
||
* Group a transcript into translation units. Whisper's segments often end
|
||
* mid-sentence; translating fragments hurts quality, so consecutive chunks
|
||
* are merged until one ends a sentence (or `maxChunksPerUnit` is reached)
|
||
* and the unit takes the merged time span — `.srt` keeps working, with
|
||
* slightly longer cues. Without chunks (nothing timestamped), each non-empty
|
||
* line is an untimed unit.
|
||
*/
|
||
export function buildTranslationUnits(
|
||
result: { text: string; chunks: TranscriptChunk[] },
|
||
maxChunksPerUnit: number = MAX_CHUNKS_PER_UNIT,
|
||
): TranslationUnit[] {
|
||
const chunks = result.chunks.filter((chunk) => chunk.text.trim() !== "");
|
||
if (chunks.length === 0) {
|
||
return result.text
|
||
.split(/\n+/)
|
||
.map((line) => line.trim())
|
||
.filter((line) => line !== "")
|
||
.map((text) => ({ text, timestamp: null }));
|
||
}
|
||
const units: TranslationUnit[] = [];
|
||
let group: TranscriptChunk[] = [];
|
||
const flush = () => {
|
||
const first = group[0];
|
||
const last = group[group.length - 1];
|
||
if (!first || !last) return;
|
||
units.push({
|
||
text: group.map((chunk) => chunk.text.trim()).join(" "),
|
||
timestamp: [first.timestamp[0], last.timestamp[1]],
|
||
});
|
||
group = [];
|
||
};
|
||
for (const chunk of chunks) {
|
||
group.push(chunk);
|
||
if (SENTENCE_END.test(chunk.text.trim()) || group.length >= maxChunksPerUnit) flush();
|
||
}
|
||
flush();
|
||
return units;
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Pieces: unit text → sentences under the model's token budget
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/**
|
||
* Cheap upper-bound token estimate for Opus-MT's SentencePiece vocab, used
|
||
* when the real tokenizer isn't at hand (tests, planning). Measured ~0.34
|
||
* tokens/char on en_alien and ~0.26 on da_festuge with the real tokenizers;
|
||
* 0.5/char plus the end-of-sentence token stays above both.
|
||
*/
|
||
export function estimateTokens(text: string): number {
|
||
return text.length === 0 ? 0 : Math.ceil(text.length / 2) + 1;
|
||
}
|
||
|
||
/**
|
||
* Default per-piece budget, well inside Opus-MT's 512 positions: quality
|
||
* falls off on very long inputs long before the hard limit, and normal
|
||
* sentences are 10–60 tokens, so this only bites on unpunctuated run-ons.
|
||
*/
|
||
export const DEFAULT_MAX_PIECE_TOKENS = 200;
|
||
|
||
export interface SplitOptions {
|
||
maxTokens?: number;
|
||
/** Token counter — the worker passes the real tokenizer. */
|
||
countTokens?: (text: string) => number;
|
||
}
|
||
|
||
/**
|
||
* 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[] = [];
|
||
let carry = "";
|
||
for (const { segment } of new Intl.Segmenter("en", { granularity: "sentence" }).segment(text)) {
|
||
const trimmed = segment.trim();
|
||
if (trimmed === "") continue;
|
||
if (ABBREVIATION.test(trimmed)) {
|
||
carry = carry === "" ? trimmed : `${carry} ${trimmed}`;
|
||
continue;
|
||
}
|
||
out.push(carry === "" ? trimmed : `${carry} ${trimmed}`);
|
||
carry = "";
|
||
}
|
||
if (carry !== "") out.push(carry);
|
||
return out;
|
||
}
|
||
|
||
/** Greedily pack `parts` (in order) into space-joined pieces under budget. */
|
||
function pack(parts: string[], maxTokens: number, count: (text: string) => number): string[] {
|
||
const pieces: string[] = [];
|
||
let current = "";
|
||
for (const part of parts) {
|
||
const candidate = current === "" ? part : `${current} ${part}`;
|
||
if (current !== "" && count(candidate) > maxTokens) {
|
||
pieces.push(current);
|
||
current = part;
|
||
} else {
|
||
current = candidate;
|
||
}
|
||
}
|
||
if (current !== "") pieces.push(current);
|
||
return pieces;
|
||
}
|
||
|
||
/** Split one over-budget sentence: commas first, then words. A single word
|
||
* longer than the budget is kept whole (nothing is ever dropped). */
|
||
function wrap(sentence: string, maxTokens: number, count: (text: string) => number): string[] {
|
||
const clauses = sentence.split(/(?<=,)\s+/);
|
||
const parts = clauses.flatMap((clause) =>
|
||
count(clause) <= maxTokens ? [clause] : pack(clause.split(/\s+/), maxTokens, count),
|
||
);
|
||
return pack(parts, maxTokens, count);
|
||
}
|
||
|
||
/**
|
||
* Cut a unit's text into the pieces actually fed to Opus-MT: one per
|
||
* sentence, an over-budget sentence wrapped on commas then words. Order is
|
||
* preserved and every word survives — joining the pieces with spaces gives
|
||
* the input back up to whitespace.
|
||
*/
|
||
export function splitForTranslation(text: string, options: SplitOptions = {}): string[] {
|
||
const maxTokens = options.maxTokens ?? DEFAULT_MAX_PIECE_TOKENS;
|
||
const count = options.countTokens ?? estimateTokens;
|
||
return sentences(text).flatMap((sentence) =>
|
||
count(sentence) <= maxTokens ? [sentence] : wrap(sentence, maxTokens, count),
|
||
);
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 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<string[]>;
|
||
/** 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<string[]> {
|
||
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
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/**
|
||
* Pair translated unit texts back with their units' time spans, in the
|
||
* `{ text, chunks }` shape the transcript view and exports already take.
|
||
*/
|
||
export function assembleTranslation(
|
||
units: readonly TranslationUnit[],
|
||
translated: readonly string[],
|
||
): { text: string; chunks: TranscriptChunk[] } {
|
||
if (units.length !== translated.length) {
|
||
throw new Error(
|
||
`Translation returned ${translated.length} parts for ${units.length} — refusing to misalign the transcript.`,
|
||
);
|
||
}
|
||
const texts = translated.map((text) => text.trim());
|
||
const chunks: TranscriptChunk[] = [];
|
||
units.forEach((unit, index) => {
|
||
if (unit.timestamp) chunks.push({ text: texts[index] ?? "", timestamp: unit.timestamp });
|
||
});
|
||
const untimed = units.every((unit) => unit.timestamp === null);
|
||
return { text: texts.join(untimed ? "\n" : " "), chunks };
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Source language guess (default for the direction picker)
|
||
// ---------------------------------------------------------------------------
|
||
|
||
const DANISH_WORDS = new Set([
|
||
"og", "ikke", "det", "jeg", "vi", "på", "til", "med", "som", "har", "af", "der",
|
||
"kan", "skal", "vil", "også", "men", "hvad", "hvordan", "være", "noget",
|
||
]);
|
||
const ENGLISH_WORDS = new Set([
|
||
"the", "and", "is", "to", "of", "that", "you", "it", "with", "have", "this", "not",
|
||
"what", "how", "be", "was", "are",
|
||
]);
|
||
|
||
/**
|
||
* Best guess at which language a transcript is in, for the direction
|
||
* picker's default (the user can always change it). An explicit Whisper
|
||
* language wins; otherwise Danish letters and function words vs English
|
||
* ones decide, defaulting to English.
|
||
*/
|
||
export function guessSourceLanguage(text: string, whisperLanguage?: string): TranslationLanguage {
|
||
if (whisperLanguage === "da" || whisperLanguage === "en") return whisperLanguage;
|
||
const lower = text.toLowerCase();
|
||
let danish = (lower.match(/[æøå]/g)?.length ?? 0) * 2;
|
||
let english = 0;
|
||
for (const word of lower.match(/\p{L}+/gu) ?? []) {
|
||
if (DANISH_WORDS.has(word)) danish += 1;
|
||
if (ENGLISH_WORDS.has(word)) english += 1;
|
||
}
|
||
return danish > english ? "da" : "en";
|
||
}
|