342 lines
12 KiB
TypeScript
342 lines
12 KiB
TypeScript
/**
|
||
* Inference Web Worker (Decisions.md D-003 — inference never runs on the main
|
||
* thread). Hosts the Transformers.js `automatic-speech-recognition` pipeline.
|
||
*
|
||
* M2: hardware backend detection with automatic fallback — "auto" resolves
|
||
* WebNN → WebGPU → WASM (D-004, D-017) — catalog-driven dtype defaults
|
||
* (D-006), model switching with pipeline disposal, and language/task options
|
||
* for multilingual models.
|
||
*/
|
||
|
||
import {
|
||
env,
|
||
pipeline,
|
||
WhisperTextStreamer,
|
||
type AutomaticSpeechRecognitionPipeline,
|
||
type ProgressInfo,
|
||
} from "@huggingface/transformers";
|
||
import { resolveDefaultDtype } from "../lib/models";
|
||
import {
|
||
deviceAttempts,
|
||
isUiToWorkerMessage,
|
||
TranscribeProgressTracker,
|
||
type DeviceCapabilities,
|
||
type DevicePreference,
|
||
type InferenceDevice,
|
||
type LoadMessage,
|
||
type ResultMessage,
|
||
type TranscribeMessage,
|
||
type TranscriptChunk,
|
||
type WorkerToUiMessage,
|
||
} from "./protocol";
|
||
|
||
// Load onnxruntime's .wasm binaries from our own origin instead of the
|
||
// jsdelivr CDN default — required for offline use (M3). The server exposes
|
||
// them at /ort/* (src/index.ts; copied into dist/ort by scripts/build.ts)
|
||
// and the service worker runtime-caches them cache-first (src/pwa/sw.ts).
|
||
if (env.backends.onnx.wasm) {
|
||
env.backends.onnx.wasm.wasmPaths = "./ort/";
|
||
}
|
||
|
||
// `lib: ["DOM"]` types `self` as Window; give it the worker-scope shape we use.
|
||
interface DedicatedWorkerScope {
|
||
postMessage(message: WorkerToUiMessage): void;
|
||
addEventListener(
|
||
type: "message",
|
||
listener: (event: MessageEvent<unknown>) => void,
|
||
): void;
|
||
}
|
||
const scope = self as unknown as DedicatedWorkerScope;
|
||
|
||
function post(message: WorkerToUiMessage): void {
|
||
scope.postMessage(message);
|
||
}
|
||
|
||
function errorMessage(error: unknown): string {
|
||
return error instanceof Error ? error.message : String(error);
|
||
}
|
||
|
||
/**
|
||
* Forward Transformers.js progress events to the UI as load-progress messages.
|
||
* The library's own "ready" event is swallowed — the worker posts its own
|
||
* `ready` (with the selected device) once the pipeline actually resolved,
|
||
* which also covers the device fallback retries (D-004, D-017).
|
||
*/
|
||
function forwardProgress(info: ProgressInfo): void {
|
||
switch (info.status) {
|
||
case "initiate":
|
||
case "download":
|
||
case "done":
|
||
post({ type: "load-progress", status: info.status, file: info.file });
|
||
break;
|
||
case "progress":
|
||
post({
|
||
type: "load-progress",
|
||
status: "progress",
|
||
file: info.file,
|
||
loaded: info.loaded,
|
||
total: info.total,
|
||
progress: info.progress,
|
||
});
|
||
break;
|
||
default:
|
||
// "ready" (posted by the worker itself, with the device) and
|
||
// "progress_total" — the UI derives totals per-file itself.
|
||
break;
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Device selection (D-004, D-017): "auto" resolves WebNN → WebGPU → WASM.
|
||
// The probe only answers "does the backend exist?"; the attempt ordering and
|
||
// fallback policy live in the pure `deviceAttempts` (protocol.ts, unit-tested).
|
||
// ---------------------------------------------------------------------------
|
||
|
||
interface NavigatorWithHardware {
|
||
gpu?: { requestAdapter(): Promise<unknown | null> };
|
||
ml?: unknown;
|
||
}
|
||
|
||
async function probeCapabilities(
|
||
preference: DevicePreference,
|
||
): Promise<DeviceCapabilities> {
|
||
const capabilities: DeviceCapabilities = { webnn: false, webgpu: false };
|
||
const hardware = navigator as NavigatorWithHardware;
|
||
if (preference === "auto" || preference === "webnn") {
|
||
try {
|
||
capabilities.webnn = Boolean(hardware.ml);
|
||
} catch {
|
||
// Probing threw — treat WebNN as unavailable.
|
||
}
|
||
}
|
||
if (preference === "auto" || preference === "webgpu") {
|
||
try {
|
||
capabilities.webgpu = Boolean(await hardware.gpu?.requestAdapter());
|
||
} catch {
|
||
// Probing threw — treat WebGPU as unavailable.
|
||
}
|
||
}
|
||
return capabilities;
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Pipeline lifecycle (load, switch, dispose)
|
||
// ---------------------------------------------------------------------------
|
||
|
||
interface ActivePipeline {
|
||
/** Identity of this load request; a different key means "switch models". */
|
||
key: string;
|
||
promise: Promise<AutomaticSpeechRecognitionPipeline>;
|
||
}
|
||
|
||
let active: ActivePipeline | null = null;
|
||
|
||
function loadKey(message: LoadMessage): string {
|
||
return JSON.stringify({
|
||
modelId: message.modelId,
|
||
device: message.device,
|
||
dtype: message.dtype ?? null,
|
||
});
|
||
}
|
||
|
||
async function createPipeline(
|
||
message: LoadMessage,
|
||
): Promise<AutomaticSpeechRecognitionPipeline> {
|
||
// Ordered devices to try (D-004/D-017): the first is the selected device;
|
||
// if pipeline creation throws there (driver quirks, unsupported ops,
|
||
// out-of-memory), the next one is tried automatically. For "auto" the chain
|
||
// is WebNN → WebGPU → WASM (filtered by the probe); WASM always terminates it.
|
||
const attempts: InferenceDevice[] = deviceAttempts(
|
||
message.device,
|
||
await probeCapabilities(message.device),
|
||
);
|
||
|
||
let lastError: unknown;
|
||
for (const device of attempts) {
|
||
const dtype = message.dtype ?? resolveDefaultDtype(message.modelId, device);
|
||
try {
|
||
const asr = await pipeline("automatic-speech-recognition", message.modelId, {
|
||
device,
|
||
dtype,
|
||
// WASM sessions cap graph optimization at "basic" (D-014): the level-2
|
||
// QDQ transform in the bundled onnxruntime-web build rejects the 8-bit
|
||
// decoders' tied embed_tokens weight ("Missing required scale …
|
||
// DequantizeLinear"). q4 models use native MatMulNBits ops, so the
|
||
// lower level costs them nothing. WebGPU keeps the default level.
|
||
...(device === "wasm" && {
|
||
session_options: { graphOptimizationLevel: "basic" as const },
|
||
}),
|
||
progress_callback: forwardProgress,
|
||
});
|
||
post({ type: "load-progress", status: "ready", device });
|
||
return asr;
|
||
} catch (error) {
|
||
lastError = error;
|
||
}
|
||
}
|
||
throw lastError;
|
||
}
|
||
|
||
function load(message: LoadMessage): Promise<AutomaticSpeechRecognitionPipeline> {
|
||
const key = loadKey(message);
|
||
// Idempotent: repeated "load" messages for the same model/device/dtype reuse
|
||
// the in-flight or completed pipeline.
|
||
if (active?.key === key) return active.promise;
|
||
|
||
const previous = active;
|
||
const promise = (async () => {
|
||
if (previous) {
|
||
// Dispose the old pipeline (frees WASM memory / GPU buffers) before
|
||
// creating the new one, so both never coexist.
|
||
try {
|
||
const old = await previous.promise;
|
||
await old.dispose();
|
||
} catch {
|
||
// The previous load itself failed — nothing to dispose.
|
||
}
|
||
}
|
||
return createPipeline(message);
|
||
})();
|
||
active = { key, promise };
|
||
promise.catch(() => {
|
||
// Allow a clean retry after a failed load; ignore stale failures that
|
||
// were already superseded by a newer load request.
|
||
if (active?.key === key) active = null;
|
||
});
|
||
return promise;
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Transcription
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/** Audio sample rate the whole app standardizes on (D-012). */
|
||
const SAMPLE_RATE = 16_000;
|
||
/** Chunked-pipeline window / overlap, in seconds (Whisper's usual 30/5). */
|
||
const CHUNK_LENGTH_S = 30;
|
||
const STRIDE_LENGTH_S = 5;
|
||
|
||
async function transcribe(message: TranscribeMessage): Promise<void> {
|
||
if (active === null) {
|
||
post({
|
||
type: "error",
|
||
stage: "transcribe",
|
||
message: "No model loaded — send a \"load\" message first.",
|
||
requestId: message.requestId,
|
||
});
|
||
return;
|
||
}
|
||
const asr = await active.promise;
|
||
|
||
const options: Record<string, unknown> = {
|
||
chunk_length_s: CHUNK_LENGTH_S,
|
||
stride_length_s: STRIDE_LENGTH_S,
|
||
return_timestamps: true,
|
||
// With return_timestamps and no explicit max_new_tokens, Transformers.js
|
||
// 4.2's Whisper generate() takes a seek-loop path that re-decodes each
|
||
// window's tail (1–3 generate passes per window). Under this pipeline's
|
||
// own 30 s/5 s chunking that is pure overhead: ~1.4× slower on long
|
||
// files, duplicated boundary text, on_finalize firing once per pass
|
||
// (breaking TranscribeProgressTracker's one-finalize-per-window
|
||
// assumption — the progress bar saturated early, then looked frozen),
|
||
// and a potential infinite loop when a window's last boundary timestamp
|
||
// is 0. Setting max_new_tokens selects the classic single-pass path;
|
||
// 440 ≈ Whisper's 448 decoder positions minus prompt headroom, far more
|
||
// than any real 30 s window produces.
|
||
max_new_tokens: 440,
|
||
};
|
||
if (message.stream ?? message.requestId === undefined) {
|
||
// File-flow pass (id-less, or routed through BrowserEngine with an
|
||
// explicit stream flag, D-010): stream partial tokens and report audio
|
||
// progress. The streamer's chunk timestamps are window-relative; the
|
||
// tracker maps them to monotone absolute seconds (see protocol.ts).
|
||
// Live passes skip both — they resolve wholesale via "result" (D-016).
|
||
const totalSeconds = message.audio.length / SAMPLE_RATE;
|
||
const tracker = new TranscribeProgressTracker(
|
||
totalSeconds,
|
||
CHUNK_LENGTH_S,
|
||
CHUNK_LENGTH_S - 2 * STRIDE_LENGTH_S,
|
||
);
|
||
const postProgress = (processedSeconds: number | null) => {
|
||
if (processedSeconds === null) return;
|
||
post({
|
||
type: "transcribe-progress",
|
||
processedSeconds,
|
||
totalSeconds,
|
||
requestId: message.requestId,
|
||
});
|
||
};
|
||
// Immediate 0/total so the UI shows the audio length before the first
|
||
// (possibly slow) token arrives.
|
||
post({
|
||
type: "transcribe-progress",
|
||
processedSeconds: 0,
|
||
totalSeconds,
|
||
requestId: message.requestId,
|
||
});
|
||
options.streamer = new WhisperTextStreamer(
|
||
// The ASR pipeline for Whisper models always carries a WhisperTokenizer.
|
||
asr.tokenizer as ConstructorParameters<typeof WhisperTextStreamer>[0],
|
||
{
|
||
skip_prompt: true,
|
||
callback_function: (text: string) =>
|
||
post({ type: "partial", text, requestId: message.requestId }),
|
||
on_chunk_start: (time: number) => postProgress(tracker.timestamp(time)),
|
||
on_chunk_end: (time: number) => postProgress(tracker.timestamp(time)),
|
||
on_finalize: () => postProgress(tracker.finalizeWindow()),
|
||
},
|
||
);
|
||
}
|
||
// Only meaningful for multilingual models; the UI omits these for `.en`
|
||
// variants. "auto" / undefined lets Whisper detect the language.
|
||
if (message.language !== undefined && message.language !== "auto") {
|
||
options.language = message.language;
|
||
}
|
||
if (message.task !== undefined) {
|
||
options.task = message.task;
|
||
}
|
||
|
||
const output = await asr(message.audio, options);
|
||
|
||
const single = (Array.isArray(output) ? output[0] : output) as {
|
||
text?: string;
|
||
chunks?: TranscriptChunk[];
|
||
};
|
||
const result: ResultMessage = {
|
||
type: "result",
|
||
text: (single.text ?? "").trim(),
|
||
chunks: single.chunks ?? [],
|
||
requestId: message.requestId,
|
||
};
|
||
post(result);
|
||
}
|
||
|
||
scope.addEventListener("message", (event) => {
|
||
const data = event.data;
|
||
if (!isUiToWorkerMessage(data)) {
|
||
post({
|
||
type: "error",
|
||
stage: "transcribe",
|
||
message: "Malformed message sent to inference worker.",
|
||
});
|
||
return;
|
||
}
|
||
switch (data.type) {
|
||
case "load":
|
||
load(data).catch((error: unknown) => {
|
||
post({ type: "error", stage: "load", message: errorMessage(error) });
|
||
});
|
||
break;
|
||
case "transcribe":
|
||
transcribe(data).catch((error: unknown) => {
|
||
post({
|
||
type: "error",
|
||
stage: "transcribe",
|
||
message: errorMessage(error),
|
||
requestId: data.requestId,
|
||
});
|
||
});
|
||
break;
|
||
}
|
||
});
|