330 lines
11 KiB
TypeScript
330 lines
11 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,
|
|
};
|
|
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;
|
|
}
|
|
});
|