99 lines
3.9 KiB
TypeScript
99 lines
3.9 KiB
TypeScript
/**
|
|
* Transcript-translation Web Worker (D-045): Opus-MT through
|
|
* Transformers.js, off the main thread (D-003).
|
|
*
|
|
* A dedicated worker rather than a message type on the Whisper worker:
|
|
* - its own onnxruntime instance gets the full WASM thread pool. ort fixes
|
|
* the pool at first init (D-023), and a Whisper worker that initialised
|
|
* ort for WebGPU/WebNN runs WASM single-threaded for life;
|
|
* - cancelling a translation terminates only this worker — the loaded
|
|
* Whisper pipeline survives, and a Whisper cancel/restart (D-018/D-025)
|
|
* never kills a translation;
|
|
* - the model is downloaded only when the user asks for a translation and
|
|
* cached by Transformers.js in the Cache API like Whisper weights
|
|
* (D-005/D-020), so a cached model works offline.
|
|
* WASM only for now: Opus-MT is small, and WebGPU behaviour can't be
|
|
* measured in this environment.
|
|
*/
|
|
|
|
import { env, type ProgressInfo } from "@huggingface/transformers";
|
|
import { getTranslationModelById, translateUnits } from "../lib/translation";
|
|
import { loadOpusMt, type OpusMtTranslator } from "./opusMt";
|
|
import {
|
|
DownloadTally,
|
|
isTranslateRequest,
|
|
type TranslateRequest,
|
|
type TranslateWorkerMessage,
|
|
} from "./translateProtocol";
|
|
|
|
// Same-origin onnxruntime binaries (offline, COEP) — see inference.worker.ts.
|
|
if (env.backends.onnx.wasm) {
|
|
env.backends.onnx.wasm.wasmPaths = "./ort/";
|
|
// D-023: with cross-origin isolation, use most of the machine.
|
|
env.backends.onnx.wasm.numThreads = globalThis.crossOriginIsolated
|
|
? Math.max(1, Math.min((navigator.hardwareConcurrency || 4) - 1, 12))
|
|
: 1;
|
|
}
|
|
|
|
interface DedicatedWorkerScope {
|
|
postMessage(message: TranslateWorkerMessage): void;
|
|
addEventListener(type: "message", listener: (event: MessageEvent<unknown>) => void): void;
|
|
}
|
|
const scope = self as unknown as DedicatedWorkerScope;
|
|
|
|
/** One loaded model at a time; switching direction disposes the other. */
|
|
let active: { modelId: string; promise: Promise<OpusMtTranslator> } | null = null;
|
|
|
|
function load(modelId: string, onProgress: (info: ProgressInfo) => void): Promise<OpusMtTranslator> {
|
|
if (active?.modelId === modelId) return active.promise;
|
|
const previous = active;
|
|
const promise = (async () => {
|
|
if (previous) await (await previous.promise.catch(() => null))?.dispose();
|
|
return loadOpusMt(modelId, { device: "wasm", progressCallback: onProgress });
|
|
})();
|
|
active = { modelId, promise };
|
|
promise.catch(() => {
|
|
if (active?.promise === promise) active = null; // allow a clean retry
|
|
});
|
|
return promise;
|
|
}
|
|
|
|
async function handle(request: TranslateRequest): Promise<void> {
|
|
const { requestId } = request;
|
|
if (!getTranslationModelById(request.modelId)) {
|
|
throw new Error(`Unknown translation model: ${request.modelId}`);
|
|
}
|
|
const tally = new DownloadTally();
|
|
const translator = await load(request.modelId, (info) => {
|
|
if (info.status === "progress") {
|
|
const { loaded, total } = tally.update(info.file, info.loaded, info.total);
|
|
scope.postMessage({ type: "translate-download", requestId, loaded, total });
|
|
}
|
|
});
|
|
const texts = await translateUnits(request.units, {
|
|
translateBatch: (pieces) => translator.translateBatch(pieces),
|
|
countTokens: translator.countTokens,
|
|
onProgress: (done, total) =>
|
|
scope.postMessage({ type: "translate-progress", requestId, done, total }),
|
|
});
|
|
scope.postMessage({ type: "translate-result", requestId, texts });
|
|
}
|
|
|
|
// Requests run one at a time, in arrival order: the UI only ever has one in
|
|
// flight, and serialising keeps two passes from sharing the thread pool.
|
|
let queue: Promise<void> = Promise.resolve();
|
|
|
|
scope.addEventListener("message", (event) => {
|
|
const data = event.data;
|
|
if (!isTranslateRequest(data)) return;
|
|
queue = queue.then(() =>
|
|
handle(data).catch((error: unknown) => {
|
|
scope.postMessage({
|
|
type: "translate-error",
|
|
requestId: data.requestId,
|
|
message: error instanceof Error ? error.message : String(error),
|
|
});
|
|
}),
|
|
);
|
|
});
|