60 lines
2.4 KiB
TypeScript
60 lines
2.4 KiB
TypeScript
/**
|
|
* 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<string[]>;
|
|
/** Real SentencePiece token count (incl. </s>) for the splitter's budget. */
|
|
countTokens(text: string): number;
|
|
dispose(): Promise<void>;
|
|
}
|
|
|
|
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<OpusMtTranslator> {
|
|
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(),
|
|
};
|
|
}
|