Agent completed work jubii
This commit is contained in:
Lucas Winther
@@ -0,0 +1,142 @@
|
||||
/**
|
||||
* Inference Web Worker (Decisions.md D-003 — inference never runs on the main
|
||||
* thread). Hosts the Transformers.js `automatic-speech-recognition` pipeline.
|
||||
*
|
||||
* M1: hardcoded-model, WASM-only. WebGPU detection + fallback lands in M2
|
||||
* (D-004); model switching/disposal also lands in M2.
|
||||
*/
|
||||
|
||||
import {
|
||||
pipeline,
|
||||
WhisperTextStreamer,
|
||||
type AutomaticSpeechRecognitionPipeline,
|
||||
type ProgressInfo,
|
||||
} from "@huggingface/transformers";
|
||||
import {
|
||||
isUiToWorkerMessage,
|
||||
type LoadMessage,
|
||||
type ResultMessage,
|
||||
type TranscriptChunk,
|
||||
type WorkerToUiMessage,
|
||||
} from "./protocol";
|
||||
|
||||
// `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. */
|
||||
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;
|
||||
case "ready":
|
||||
post({ type: "load-progress", status: "ready" });
|
||||
break;
|
||||
default:
|
||||
// e.g. "progress_total" — the UI derives totals per-file itself.
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let asrPromise: Promise<AutomaticSpeechRecognitionPipeline> | null = null;
|
||||
|
||||
function load(message: LoadMessage): Promise<AutomaticSpeechRecognitionPipeline> {
|
||||
// Idempotent for M1's single hardcoded model: repeated "load" messages
|
||||
// reuse the in-flight/completed pipeline.
|
||||
asrPromise ??= pipeline("automatic-speech-recognition", message.modelId, {
|
||||
device: message.device,
|
||||
progress_callback: forwardProgress,
|
||||
});
|
||||
return asrPromise;
|
||||
}
|
||||
|
||||
async function transcribe(audio: Float32Array): Promise<void> {
|
||||
if (asrPromise === null) {
|
||||
post({
|
||||
type: "error",
|
||||
stage: "transcribe",
|
||||
message: "No model loaded — send a \"load\" message first.",
|
||||
});
|
||||
return;
|
||||
}
|
||||
const asr = await asrPromise;
|
||||
|
||||
const 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 }),
|
||||
},
|
||||
);
|
||||
|
||||
const output = await asr(audio, {
|
||||
chunk_length_s: 30,
|
||||
stride_length_s: 5,
|
||||
return_timestamps: true,
|
||||
streamer,
|
||||
});
|
||||
|
||||
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 ?? [],
|
||||
};
|
||||
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) => {
|
||||
asrPromise = null; // allow retry after a failed download
|
||||
post({ type: "error", stage: "load", message: errorMessage(error) });
|
||||
});
|
||||
break;
|
||||
case "transcribe":
|
||||
transcribe(data.audio).catch((error: unknown) => {
|
||||
post({ type: "error", stage: "transcribe", message: errorMessage(error) });
|
||||
});
|
||||
break;
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,189 @@
|
||||
/**
|
||||
* Typed postMessage protocol between the UI and the inference Web Worker
|
||||
* (Decisions.md D-003). Shared by `src/worker/inference.worker.ts` and the UI.
|
||||
*
|
||||
* Everything in this file must stay free of browser/worker globals so it can
|
||||
* be unit-tested with `bun test`.
|
||||
*/
|
||||
|
||||
/** Devices supported by the worker. M1 is WASM-only; M2 adds "webgpu" (D-004). */
|
||||
export type InferenceDevice = "wasm";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// UI → worker
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export interface LoadMessage {
|
||||
type: "load";
|
||||
/** Hugging Face Hub model id, e.g. "onnx-community/whisper-tiny.en". */
|
||||
modelId: string;
|
||||
device: InferenceDevice;
|
||||
}
|
||||
|
||||
export interface TranscribeMessage {
|
||||
type: "transcribe";
|
||||
/**
|
||||
* Mono PCM at 16 kHz (see src/lib/audio.ts). The underlying ArrayBuffer
|
||||
* should be passed as a transferable for zero-copy handoff.
|
||||
*/
|
||||
audio: Float32Array;
|
||||
}
|
||||
|
||||
export type UiToWorkerMessage = LoadMessage | TranscribeMessage;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Worker → UI
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Model download / initialization progress. Mirrors the Transformers.js
|
||||
* `progress_callback` events we care about; `status: "ready"` marks the end
|
||||
* of a successful load.
|
||||
*/
|
||||
export type LoadProgressMessage =
|
||||
| { type: "load-progress"; status: "initiate"; file: string }
|
||||
| { type: "load-progress"; status: "download"; file: string }
|
||||
| {
|
||||
type: "load-progress";
|
||||
status: "progress";
|
||||
file: string;
|
||||
/** Bytes downloaded so far. */
|
||||
loaded: number;
|
||||
/** Total bytes for this file. */
|
||||
total: number;
|
||||
/** Percentage in [0, 100]. */
|
||||
progress: number;
|
||||
}
|
||||
| { type: "load-progress"; status: "done"; file: string }
|
||||
| { type: "load-progress"; status: "ready" };
|
||||
|
||||
/** Incremental transcript fragment — append to the text streamed so far. */
|
||||
export interface PartialMessage {
|
||||
type: "partial";
|
||||
text: string;
|
||||
}
|
||||
|
||||
/** One timestamped segment of the final transcript. */
|
||||
export interface TranscriptChunk {
|
||||
text: string;
|
||||
/** [start, end] in seconds; end may be null if Whisper never closed the segment. */
|
||||
timestamp: [number, number | null];
|
||||
}
|
||||
|
||||
export interface ResultMessage {
|
||||
type: "result";
|
||||
text: string;
|
||||
chunks: TranscriptChunk[];
|
||||
}
|
||||
|
||||
export interface ErrorMessage {
|
||||
type: "error";
|
||||
stage: "load" | "transcribe";
|
||||
message: string;
|
||||
}
|
||||
|
||||
export type WorkerToUiMessage =
|
||||
| LoadProgressMessage
|
||||
| PartialMessage
|
||||
| ResultMessage
|
||||
| ErrorMessage;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Runtime guards (postMessage data arrives untyped)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === "object" && value !== null;
|
||||
}
|
||||
|
||||
export function isUiToWorkerMessage(value: unknown): value is UiToWorkerMessage {
|
||||
if (!isRecord(value)) return false;
|
||||
switch (value.type) {
|
||||
case "load":
|
||||
return typeof value.modelId === "string" && value.device === "wasm";
|
||||
case "transcribe":
|
||||
return value.audio instanceof Float32Array;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
export function isWorkerToUiMessage(value: unknown): value is WorkerToUiMessage {
|
||||
if (!isRecord(value)) return false;
|
||||
switch (value.type) {
|
||||
case "load-progress":
|
||||
switch (value.status) {
|
||||
case "initiate":
|
||||
case "download":
|
||||
case "done":
|
||||
return typeof value.file === "string";
|
||||
case "progress":
|
||||
return (
|
||||
typeof value.file === "string" &&
|
||||
typeof value.loaded === "number" &&
|
||||
typeof value.total === "number" &&
|
||||
typeof value.progress === "number"
|
||||
);
|
||||
case "ready":
|
||||
return true;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
case "partial":
|
||||
return typeof value.text === "string";
|
||||
case "result":
|
||||
return typeof value.text === "string" && Array.isArray(value.chunks);
|
||||
case "error":
|
||||
return (
|
||||
typeof value.message === "string" &&
|
||||
(value.stage === "load" || value.stage === "transcribe")
|
||||
);
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Load-progress accumulation (pure; drives the download progress UI)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export interface FileProgress {
|
||||
file: string;
|
||||
loaded: number;
|
||||
total: number;
|
||||
done: boolean;
|
||||
}
|
||||
|
||||
export interface LoadProgressState {
|
||||
/** Per-file progress, keyed by file name, in first-seen order. */
|
||||
files: Record<string, FileProgress>;
|
||||
/** True once the pipeline reported ready. */
|
||||
ready: boolean;
|
||||
}
|
||||
|
||||
export const initialLoadProgress: LoadProgressState = { files: {}, ready: false };
|
||||
|
||||
/** Fold one load-progress message into the accumulated state. Pure. */
|
||||
export function applyLoadProgress(
|
||||
state: LoadProgressState,
|
||||
message: LoadProgressMessage,
|
||||
): LoadProgressState {
|
||||
if (message.status === "ready") {
|
||||
return { ...state, ready: true };
|
||||
}
|
||||
const previous = state.files[message.file];
|
||||
const next: FileProgress = {
|
||||
file: message.file,
|
||||
loaded: previous?.loaded ?? 0,
|
||||
total: previous?.total ?? 0,
|
||||
done: previous?.done ?? false,
|
||||
};
|
||||
if (message.status === "progress") {
|
||||
next.loaded = message.loaded;
|
||||
next.total = message.total;
|
||||
} else if (message.status === "done") {
|
||||
next.done = true;
|
||||
if (next.total > 0) next.loaded = next.total;
|
||||
}
|
||||
return { ...state, files: { ...state.files, [message.file]: next } };
|
||||
}
|
||||
Reference in New Issue
Block a user