Run Parakeet TDT in the inference worker via parakeet.js on the shared onnxruntime-web
Adds parakeet.js 1.4.4 (MIT) as a devDependency bundled into the worker only; a build plugin keeps its own onnxruntime-web copy out of the bundle. Includes an opt-in real-model test (PARAKEET_REAL=1). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Lucas Winther
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
c5a57cd33a
commit
708fae5778
@@ -19,6 +19,7 @@
|
||||
"@types/react": "^19",
|
||||
"@types/react-dom": "^19",
|
||||
"ffmpeg-static": "^5.3.0",
|
||||
"parakeet.js": "1.4.4",
|
||||
"typescript": "^7.0.2",
|
||||
},
|
||||
},
|
||||
@@ -254,6 +255,8 @@
|
||||
|
||||
"onnxruntime-web": ["onnxruntime-web@1.26.0-dev.20260416-b7804b056c", "", { "dependencies": { "flatbuffers": "^25.1.24", "guid-typescript": "^1.0.9", "long": "^5.2.3", "onnxruntime-common": "1.24.0-dev.20251116-b39e144322", "platform": "^1.3.6", "protobufjs": "^7.2.4" } }, "sha512-MD6Ss4GSpQBo6zqoJzyT9LRbKYs7x/JVN23FT24EcEvlqF4VuzPOeH6X38orZPKHQDbprn7K+SBpu0/mj2CQiw=="],
|
||||
|
||||
"parakeet.js": ["parakeet.js@1.4.4", "", { "dependencies": { "onnxruntime-web": "1.24.1" } }, "sha512-+tYIDgp799bvsDye3y2lFyspYHgTOJ5YMZi0mWYABXyJjp/3NZmNaMTnS9jhf4qlIQJoBk0Nmj+UwsI3zk+Aqg=="],
|
||||
|
||||
"parse-cache-control": ["parse-cache-control@1.0.1", "", {}, "sha512-60zvsJReQPX5/QP0Kzfd/VrpjScIQ7SHBW6bFCYfEP+fp0Eppr1SHhIO5nd1PjZtvclzSzES9D/p5nFJurwfWg=="],
|
||||
|
||||
"passport": ["passport@0.7.0", "", { "dependencies": { "passport-strategy": "1.x.x", "pause": "0.0.1", "utils-merge": "^1.0.1" } }, "sha512-cPLl+qZpSc+ireUvt+IzqbED1cHHkDoVYMo30jbJIdOOjQ1MQYZBPiNvmi8UM6lJuOpTPXJGZQk0DtC4y61MYQ=="],
|
||||
@@ -311,5 +314,9 @@
|
||||
"http-response-object/@types/node": ["@types/node@10.17.60", "", {}, "sha512-F0KIgDJfy2nA3zMLmWGKxcH2ZVEtCZXHHdOQs2gSaQ27+lNeEfGxzkIw90aXswATX7AZ33tahPbzy6KAfUreVw=="],
|
||||
|
||||
"onnxruntime-web/onnxruntime-common": ["onnxruntime-common@1.24.0-dev.20251116-b39e144322", "", {}, "sha512-BOoomdHYmNRL5r4iQ4bMvsl2t0/hzVQ3OM3PHD0gxeXu1PmggqBv3puZicEUVOA3AtHHYmqZtjMj9FOfGrATTw=="],
|
||||
|
||||
"parakeet.js/onnxruntime-web": ["onnxruntime-web@1.24.1", "", { "dependencies": { "flatbuffers": "^25.1.24", "guid-typescript": "^1.0.9", "long": "^5.2.3", "onnxruntime-common": "1.24.1", "platform": "^1.3.6", "protobufjs": "^7.2.4" } }, "sha512-i2u395dv+ZEQBdH+aORvlu19Bzvlg5AXJ7wjxnL350hknOP9z0UeP3pVfjkpMEWMPy2T6nCQxetKTmNia6wSzg=="],
|
||||
|
||||
"parakeet.js/onnxruntime-web/onnxruntime-common": ["onnxruntime-common@1.24.1", "", {}, "sha512-UnV15u4p4XxoIV+jFP4hXPsW93s3QrwLSpi20HUDYHoTfI4z4sjzex3L4XDOxGGZJ/M/catrwAG2go958UQq0w=="],
|
||||
}
|
||||
}
|
||||
|
||||
@@ -25,6 +25,7 @@
|
||||
"@types/react": "^19",
|
||||
"@types/react-dom": "^19",
|
||||
"ffmpeg-static": "^5.3.0",
|
||||
"parakeet.js": "1.4.4",
|
||||
"typescript": "^7.0.2"
|
||||
},
|
||||
"trustedDependencies": [
|
||||
|
||||
@@ -17,6 +17,7 @@ import { join, relative } from "node:path";
|
||||
import { injectPrecacheManifest, precachePathsFromOutputs } from "../src/pwa/precache.ts";
|
||||
import { ensureFonts, FONTS_DIR } from "./fonts.ts";
|
||||
import { ensureIcons, ICON_SIZES, ICONS_DIR } from "./generate-icons.ts";
|
||||
import { ortAliasPlugin } from "./ortAlias.ts";
|
||||
|
||||
const ROOT = join(import.meta.dir, "..");
|
||||
const DIST = join(ROOT, "dist");
|
||||
@@ -50,6 +51,8 @@ const workerBuild = await Bun.build({
|
||||
target: "browser",
|
||||
format: "esm",
|
||||
minify: true,
|
||||
// One onnxruntime in the bundle (D-WS7a).
|
||||
plugins: [ortAliasPlugin],
|
||||
});
|
||||
const workerScript = await workerBuild.outputs[0]!.text();
|
||||
await Bun.write(join(DIST, "inference.worker.js"), workerScript);
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
/**
|
||||
* Bun build plugin: parakeet.js's bare `onnxruntime-web` import resolves to
|
||||
* `src/worker/ortShim.ts` (= Transformers.js's `onnxruntime-web/webgpu`), so
|
||||
* the inference worker bundles exactly one onnxruntime (D-WS7a). Used by the
|
||||
* production build (scripts/build.ts) and the dev server's worker bundle
|
||||
* (src/index.ts). Without it the bundle carries parakeet.js's own 1.24.1
|
||||
* copy as well — never executed, but ~0.5 MB of dead weight and a CDN
|
||||
* wasmPaths default.
|
||||
*/
|
||||
|
||||
import type { BunPlugin } from "bun";
|
||||
import { join } from "node:path";
|
||||
|
||||
const SHIM = join(import.meta.dir, "../src/worker/ortShim.ts");
|
||||
|
||||
export const ortAliasPlugin: BunPlugin = {
|
||||
name: "parakeet-ort-alias",
|
||||
setup(build) {
|
||||
build.onResolve({ filter: /^onnxruntime-web$/ }, (args) =>
|
||||
/[\\/]parakeet\.js[\\/]/.test(args.importer) ? { path: SHIM } : undefined,
|
||||
);
|
||||
},
|
||||
};
|
||||
@@ -3,6 +3,7 @@ import { tmpdir } from "node:os";
|
||||
import { join, normalize, sep } from "node:path";
|
||||
import { ensureFonts } from "../scripts/fonts.ts";
|
||||
import { ensureIcons, ICONS_DIR } from "../scripts/generate-icons.ts";
|
||||
import { ortAliasPlugin } from "../scripts/ortAlias.ts";
|
||||
import { injectPrecacheManifest } from "./pwa/precache.ts";
|
||||
import { auth } from "./server/auth.ts";
|
||||
import { AUTH_ENABLED } from "./server/flags.ts";
|
||||
@@ -234,6 +235,8 @@ async function bundleEntrypoint(relativePath: string): Promise<string> {
|
||||
// Never minify: sw.ts relies on its placeholder string literals surviving
|
||||
// bundling verbatim (src/pwa/precache.ts), and dev builds favor debuggability.
|
||||
minify: false,
|
||||
// One onnxruntime in the worker bundle (D-WS7a); harmless for sw.ts.
|
||||
plugins: [ortAliasPlugin],
|
||||
});
|
||||
const artifact = build.outputs[0];
|
||||
if (!artifact) throw new Error(`Failed to bundle src/${relativePath}`);
|
||||
|
||||
@@ -16,9 +16,22 @@ import {
|
||||
type PretrainedModelOptions,
|
||||
type ProgressInfo,
|
||||
} from "@huggingface/transformers";
|
||||
import { resolveDefaultDtype } from "../lib/models";
|
||||
// The very onnxruntime-web module Transformers.js imports (its web build
|
||||
// imports "onnxruntime-web/webgpu"), so the Parakeet path shares one runtime,
|
||||
// one wasm binary from /ort/, and one env (D-WS7a).
|
||||
import * as ort from "onnxruntime-web/webgpu";
|
||||
import { getModel, modelFamily, resolveDefaultDtype } from "../lib/models";
|
||||
import { PARAKEET_VOCAB_FILE, parakeetFileName, parakeetFileUrl } from "../lib/parakeet";
|
||||
import { fetchModelFile, type ModelFileCacheStorage } from "./modelFiles";
|
||||
import {
|
||||
createParakeetEngine,
|
||||
parakeetSessionOptions,
|
||||
type OrtLike,
|
||||
type ParakeetEngine,
|
||||
} from "./parakeetEngine";
|
||||
import {
|
||||
deviceAttempts,
|
||||
parakeetDeviceAttempts,
|
||||
dtypeRequiresF16,
|
||||
isUiToWorkerMessage,
|
||||
TranscribeProgressTracker,
|
||||
@@ -199,10 +212,20 @@ async function probeCapabilities(
|
||||
// Pipeline lifecycle (load, switch, dispose)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/** A loaded model of either family (D-WS7a): the worker dispatches on it. */
|
||||
type LoadedModel =
|
||||
| { family: "whisper"; asr: AutomaticSpeechRecognitionPipeline }
|
||||
| { family: "parakeet"; engine: ParakeetEngine };
|
||||
|
||||
interface ActivePipeline {
|
||||
/** Identity of this load request; a different key means "switch models". */
|
||||
key: string;
|
||||
promise: Promise<AutomaticSpeechRecognitionPipeline>;
|
||||
promise: Promise<LoadedModel>;
|
||||
}
|
||||
|
||||
async function disposeModel(model: LoadedModel): Promise<void> {
|
||||
if (model.family === "whisper") await model.asr.dispose();
|
||||
else await model.engine.dispose();
|
||||
}
|
||||
|
||||
let active: ActivePipeline | null = null;
|
||||
@@ -292,7 +315,67 @@ async function createPipeline(
|
||||
throw lastError;
|
||||
}
|
||||
|
||||
function load(message: LoadMessage): Promise<AutomaticSpeechRecognitionPipeline> {
|
||||
/**
|
||||
* Parakeet load (D-WS7a): fetch vocab + decoder + encoder through the
|
||||
* Cache-API downloader (same cache and progress messages as Transformers.js,
|
||||
* D-005), then create the sessions on our shared onnxruntime-web. Device
|
||||
* policy is `parakeetDeviceAttempts` (WASM unless WebGPU was chosen
|
||||
* explicitly); the D-025 restart rule applies exactly as for Whisper.
|
||||
*/
|
||||
async function createParakeet(message: LoadMessage): Promise<ParakeetEngine> {
|
||||
const capabilities =
|
||||
message.device === "webgpu"
|
||||
? await probeCapabilities("webgpu")
|
||||
: { webnn: false, webgpu: false };
|
||||
const attempts = parakeetDeviceAttempts(message.device, capabilities);
|
||||
const downloads = {
|
||||
caches: typeof caches === "undefined" ? undefined : (caches as unknown as ModelFileCacheStorage),
|
||||
fetch: (url: string) => fetch(url),
|
||||
onProgress: post,
|
||||
};
|
||||
const file = async (name: string) =>
|
||||
fetchModelFile(parakeetFileUrl(name), name, downloads);
|
||||
|
||||
let lastError: unknown;
|
||||
let acceleratorAttempted = false;
|
||||
for (const device of attempts) {
|
||||
if (device === "wasm" && acceleratorAttempted && desiredWasmThreads() > 1) {
|
||||
post({ type: "load-restart", modelId: message.modelId, preference: message.device, device: "wasm" });
|
||||
throw new RestartRequested();
|
||||
}
|
||||
configureWasmThreads(device);
|
||||
const dtype = message.dtype ?? resolveDefaultDtype(message.modelId, device);
|
||||
const encoderDtype = typeof dtype === "string" ? dtype : dtype.encoder_model;
|
||||
const decoderDtype = typeof dtype === "string" ? dtype : dtype.decoder_model_merged;
|
||||
try {
|
||||
const vocab = new TextDecoder().decode(await file(PARAKEET_VOCAB_FILE));
|
||||
const decoder = await file(parakeetFileName("decoder", decoderDtype));
|
||||
const engine = await createParakeetEngine(
|
||||
ort as unknown as OrtLike,
|
||||
{ encoder: await file(parakeetFileName("encoder", encoderDtype)), decoder, vocab },
|
||||
parakeetSessionOptions(device === "webgpu" ? "webgpu" : "wasm"),
|
||||
);
|
||||
const threads = ort.env.wasm.numThreads ?? 1;
|
||||
console.info(`[whisper] parakeet ready: device=${device} wasmThreads=${threads}`);
|
||||
post({ type: "load-progress", status: "ready", device, threads });
|
||||
return engine;
|
||||
} catch (error) {
|
||||
lastError = error;
|
||||
if (device !== "wasm") acceleratorAttempted = true;
|
||||
console.warn(`[whisper] parakeet creation failed on ${device}: ${errorMessage(error)}`);
|
||||
}
|
||||
}
|
||||
throw lastError;
|
||||
}
|
||||
|
||||
async function createModel(message: LoadMessage): Promise<LoadedModel> {
|
||||
if (modelFamily(getModel(message.modelId)) === "parakeet") {
|
||||
return { family: "parakeet", engine: await createParakeet(message) };
|
||||
}
|
||||
return { family: "whisper", asr: await createPipeline(message) };
|
||||
}
|
||||
|
||||
function load(message: LoadMessage): Promise<LoadedModel> {
|
||||
const key = loadKey(message);
|
||||
// Idempotent: repeated "load" messages for the same model/device/dtype reuse
|
||||
// the in-flight or completed pipeline.
|
||||
@@ -304,13 +387,12 @@ function load(message: LoadMessage): Promise<AutomaticSpeechRecognitionPipeline>
|
||||
// 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();
|
||||
await disposeModel(await previous.promise);
|
||||
} catch {
|
||||
// The previous load itself failed — nothing to dispose.
|
||||
}
|
||||
}
|
||||
return createPipeline(message);
|
||||
return createModel(message);
|
||||
})();
|
||||
active = { key, promise };
|
||||
promise.catch(() => {
|
||||
@@ -331,6 +413,34 @@ const SAMPLE_RATE = 16_000;
|
||||
const CHUNK_LENGTH_S = 30;
|
||||
const STRIDE_LENGTH_S = 5;
|
||||
|
||||
/**
|
||||
* Parakeet pass (D-WS7a). English-only: `language`/`task` are ignored (the UI
|
||||
* never sends them for a non-multilingual model). File-flow passes stream one
|
||||
* `partial` + one `transcribe-progress` per ≤ 30 s window — absolute seconds
|
||||
* straight from the window plan, so no TranscribeProgressTracker is needed.
|
||||
* Live passes resolve wholesale via "result", exactly like Whisper's (D-016).
|
||||
*/
|
||||
async function transcribeParakeet(
|
||||
engine: ParakeetEngine,
|
||||
message: TranscribeMessage,
|
||||
): Promise<void> {
|
||||
const stream = message.stream ?? message.requestId === undefined;
|
||||
const totalSeconds = message.audio.length / SAMPLE_RATE;
|
||||
if (stream) {
|
||||
post({ type: "transcribe-progress", processedSeconds: 0, totalSeconds, requestId: message.requestId });
|
||||
}
|
||||
const output = await engine.transcribe(message.audio, {
|
||||
...(stream && {
|
||||
onWindow: (chunks: TranscriptChunk[], processedSeconds: number) => {
|
||||
const text = chunks.map((chunk) => chunk.text).join("");
|
||||
if (text) post({ type: "partial", text, requestId: message.requestId });
|
||||
post({ type: "transcribe-progress", processedSeconds, totalSeconds, requestId: message.requestId });
|
||||
},
|
||||
}),
|
||||
});
|
||||
post({ type: "result", text: output.text, chunks: output.chunks, requestId: message.requestId });
|
||||
}
|
||||
|
||||
async function transcribe(message: TranscribeMessage): Promise<void> {
|
||||
if (active === null) {
|
||||
post({
|
||||
@@ -341,7 +451,12 @@ async function transcribe(message: TranscribeMessage): Promise<void> {
|
||||
});
|
||||
return;
|
||||
}
|
||||
const asr = await active.promise;
|
||||
const loaded = await active.promise;
|
||||
if (loaded.family === "parakeet") {
|
||||
await transcribeParakeet(loaded.engine, message);
|
||||
return;
|
||||
}
|
||||
const asr = loaded.asr;
|
||||
|
||||
const options: Record<string, unknown> = {
|
||||
chunk_length_s: CHUNK_LENGTH_S,
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
/**
|
||||
* Stand-in for the bare `onnxruntime-web` import inside parakeet.js
|
||||
* (D-WS7a). parakeet.js depends on its own onnxruntime-web 1.24.1, imported
|
||||
* only by its `initOrt` helper, which this app never calls — but a bundler
|
||||
* would still pull that second runtime into the worker. `scripts/ortAlias.ts`
|
||||
* redirects the import here, to the one runtime Transformers.js uses.
|
||||
*/
|
||||
export * from "onnxruntime-web/webgpu";
|
||||
@@ -0,0 +1,130 @@
|
||||
/**
|
||||
* Parakeet TDT engine (Decisions.md D-WS7a): the non-Whisper inference path
|
||||
* of the browser engine, used by the inference worker for catalog entries
|
||||
* with `family: "parakeet"`.
|
||||
*
|
||||
* Runtime: parakeet.js's `ParakeetModel` (TDT greedy decoding over the
|
||||
* encoder + decoder_joint ONNX pair, JS mel preprocessor) on an injected
|
||||
* onnxruntime — in the worker that is the very `onnxruntime-web/webgpu`
|
||||
* instance Transformers.js uses (one runtime, its wasm self-hosted at
|
||||
* /ort/), in tests onnxruntime-web or onnxruntime-node. We deliberately
|
||||
* bypass parakeet.js's `fromUrls`/`fromHub`/`initOrt`: those import their own
|
||||
* onnxruntime-web, point it at the jsdelivr CDN (blocked by COEP, and not
|
||||
* offline), and fetch weights outside our Cache-API bookkeeping.
|
||||
*
|
||||
* Windowing and word → chunk mapping are pure (src/lib/parakeet.ts).
|
||||
*/
|
||||
|
||||
import { JsPreprocessor } from "parakeet.js/mel";
|
||||
import { ParakeetModel } from "parakeet.js/parakeet";
|
||||
import { ParakeetVocabTokenizer, parseVocab, transcribeWindows, PARAKEET_SAMPLE_RATE } from "../lib/parakeet";
|
||||
import type { TranscriptChunk } from "./protocol";
|
||||
|
||||
/** The slice of an onnxruntime module this path needs (web and node agree). */
|
||||
export interface OrtLike {
|
||||
InferenceSession: {
|
||||
create(model: Uint8Array, options?: Record<string, unknown>): Promise<{ release?(): Promise<void> }>;
|
||||
};
|
||||
Tensor: unknown;
|
||||
}
|
||||
|
||||
export interface ParakeetAssets {
|
||||
encoder: Uint8Array;
|
||||
decoder: Uint8Array;
|
||||
vocab: string;
|
||||
}
|
||||
|
||||
export interface ParakeetSessionOptions {
|
||||
encoder: Record<string, unknown>;
|
||||
decoder: Record<string, unknown>;
|
||||
}
|
||||
|
||||
export interface ParakeetTranscribeCallbacks {
|
||||
/** After each ≤ 30 s window: its new chunks and absolute progress. */
|
||||
onWindow?: (chunks: TranscriptChunk[], processedSeconds: number, totalSeconds: number) => void;
|
||||
}
|
||||
|
||||
export interface ParakeetEngine {
|
||||
transcribe(
|
||||
audio: Float32Array,
|
||||
callbacks?: ParakeetTranscribeCallbacks,
|
||||
): Promise<{ text: string; chunks: TranscriptChunk[] }>;
|
||||
dispose(): Promise<void>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Session options per device. WASM: both parts on the wasm EP at the full
|
||||
* graph optimization level — measured fine for these int8 graphs on our
|
||||
* onnxruntime-web build (D-014's "basic" cap is about Whisper's 8-bit tied
|
||||
* embeddings, which Parakeet doesn't have). WebGPU (explicit opt-in only):
|
||||
* the fp16 encoder on WebGPU, the small autoregressive decoder on WASM, the
|
||||
* split parakeet.js itself uses ("webgpu-hybrid").
|
||||
*/
|
||||
export function parakeetSessionOptions(device: "wasm" | "webgpu"): ParakeetSessionOptions {
|
||||
const wasm = { executionProviders: ["wasm"], graphOptimizationLevel: "all" };
|
||||
return device === "webgpu"
|
||||
? { encoder: { executionProviders: ["webgpu"], graphOptimizationLevel: "all" }, decoder: wasm }
|
||||
: { encoder: wasm, decoder: wasm };
|
||||
}
|
||||
|
||||
/** What we call on parakeet.js's model (its constructor isn't in its .d.ts). */
|
||||
interface ParakeetModelLike {
|
||||
transcribe(
|
||||
audio: Float32Array,
|
||||
sampleRate: number,
|
||||
options: Record<string, unknown>,
|
||||
): Promise<{ words: Array<{ text: string; start_time: number; end_time: number }> }>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Create the two ONNX sessions and wrap them. Sessions are created one after
|
||||
* the other (encoder first) so only one weight buffer is being parsed at a
|
||||
* time — the D-027 peak-memory lesson. Callers should drop their references
|
||||
* to `assets` afterwards so the 650 MB encoder buffer can be collected.
|
||||
*/
|
||||
export async function createParakeetEngine(
|
||||
ort: OrtLike,
|
||||
assets: ParakeetAssets,
|
||||
sessionOptions: ParakeetSessionOptions,
|
||||
): Promise<ParakeetEngine> {
|
||||
const tokenizer = new ParakeetVocabTokenizer(parseVocab(assets.vocab));
|
||||
const encoderSession = await ort.InferenceSession.create(assets.encoder, sessionOptions.encoder);
|
||||
let joinerSession: Awaited<ReturnType<OrtLike["InferenceSession"]["create"]>>;
|
||||
try {
|
||||
joinerSession = await ort.InferenceSession.create(assets.decoder, sessionOptions.decoder);
|
||||
} catch (error) {
|
||||
await encoderSession.release?.();
|
||||
throw error;
|
||||
}
|
||||
const Model = ParakeetModel as unknown as new (config: Record<string, unknown>) => ParakeetModelLike;
|
||||
const model = new Model({
|
||||
tokenizer,
|
||||
encoderSession,
|
||||
joinerSession,
|
||||
preprocessor: new JsPreprocessor({ nMels: 128 }),
|
||||
ort,
|
||||
nMels: 128,
|
||||
});
|
||||
|
||||
return {
|
||||
transcribe(audio, callbacks = {}) {
|
||||
return transcribeWindows(
|
||||
audio,
|
||||
async (window, offsetSeconds) => {
|
||||
const result = await model.transcribe(window, PARAKEET_SAMPLE_RATE, {
|
||||
returnTimestamps: true,
|
||||
timeOffset: offsetSeconds,
|
||||
// Profiling defaults to on and logs every call; we don't need it.
|
||||
enableProfiling: false,
|
||||
});
|
||||
return result.words;
|
||||
},
|
||||
callbacks,
|
||||
);
|
||||
},
|
||||
async dispose() {
|
||||
await encoderSession.release?.();
|
||||
await joinerSession.release?.();
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
import { afterAll, describe, expect, test } from "bun:test";
|
||||
import { existsSync, mkdirSync } from "node:fs";
|
||||
import { join } from "node:path";
|
||||
import ffmpegPath from "ffmpeg-static";
|
||||
import * as ort from "onnxruntime-web";
|
||||
import { LIVE_TUNING, LiveTranscriber, type LiveTranscribeFn } from "../src/lib/live";
|
||||
import { PARAKEET_MODEL_ID, resolveDefaultDtype } from "../src/lib/models";
|
||||
import { PARAKEET_VOCAB_FILE, parakeetFileName, parakeetFileUrl } from "../src/lib/parakeet";
|
||||
import {
|
||||
createParakeetEngine,
|
||||
parakeetSessionOptions,
|
||||
type OrtLike,
|
||||
type ParakeetEngine,
|
||||
} from "../src/worker/parakeetEngine";
|
||||
import { wordErrorRate } from "./dsp";
|
||||
|
||||
/**
|
||||
* Real-model accuracy of the Parakeet engine path (D-WS7a), opt-in:
|
||||
*
|
||||
* PARAKEET_REAL=1 bun test tests/parakeet-real-model.test.ts
|
||||
*
|
||||
* Runs the worker's own `createParakeetEngine` on onnxruntime-web's WASM
|
||||
* backend (the browser engine's baseline runtime, here under Bun) with the
|
||||
* catalog's WASM int8 weights, over the English clips in tests/audio_files:
|
||||
* whole-file mode, and live mode through the production LiveTranscriber.
|
||||
* Held to the same bars as Whisper Large v3 Turbo in live-real-model.test.ts.
|
||||
* First run downloads ~670 MB into node_modules/@huggingface/transformers/.cache/.
|
||||
* Opt-in (not in `bun test` by default) so the default suite's runtime and
|
||||
* network use don't change; run it when touching src/lib/parakeet.ts or
|
||||
* src/worker/parakeetEngine.ts.
|
||||
*/
|
||||
|
||||
const ENABLED = process.env.PARAKEET_REAL === "1";
|
||||
const RATE = 16_000;
|
||||
const AUDIO_DIR = join(import.meta.dir, "audio_files");
|
||||
const MODEL_DIR = join(
|
||||
import.meta.dir,
|
||||
"../node_modules/@huggingface/transformers/.cache",
|
||||
PARAKEET_MODEL_ID,
|
||||
);
|
||||
const CLIPS = [
|
||||
{ base: "en_aiproblem", file: "en_aiproblem.m4a", maxWer: 0.25 },
|
||||
{ base: "en_alien", file: "en_alien.opus", maxWer: 0.35 },
|
||||
];
|
||||
|
||||
function normalize(text: string): string {
|
||||
return text
|
||||
.toLowerCase()
|
||||
.replace(/['’]/g, "")
|
||||
.replace(/[.,!?;:"“”‘‚«»()\-–—…]/g, " ")
|
||||
.replace(/\s+/g, " ")
|
||||
.trim();
|
||||
}
|
||||
|
||||
async function decode16k(path: string): Promise<Float32Array> {
|
||||
const proc = Bun.spawn(
|
||||
[ffmpegPath as unknown as string, "-hide_banner", "-loglevel", "error", "-i", path,
|
||||
"-ac", "1", "-ar", String(RATE), "-f", "f32le", "-"],
|
||||
{ stdout: "pipe", stderr: "pipe" },
|
||||
);
|
||||
const [bytes, exitCode] = await Promise.all([new Response(proc.stdout).arrayBuffer(), proc.exited]);
|
||||
if (exitCode !== 0) throw new Error(`ffmpeg failed on ${path}`);
|
||||
return new Float32Array(bytes);
|
||||
}
|
||||
|
||||
async function modelFile(file: string): Promise<Uint8Array> {
|
||||
const path = join(MODEL_DIR, file);
|
||||
if (!existsSync(path)) {
|
||||
mkdirSync(MODEL_DIR, { recursive: true });
|
||||
const response = await fetch(parakeetFileUrl(file));
|
||||
if (!response.ok) throw new Error(`download ${file}: HTTP ${response.status}`);
|
||||
await Bun.write(path, response);
|
||||
}
|
||||
return new Uint8Array(await Bun.file(path).arrayBuffer());
|
||||
}
|
||||
|
||||
let enginePromise: Promise<ParakeetEngine> | null = null;
|
||||
function loadEngine(): Promise<ParakeetEngine> {
|
||||
enginePromise ??= (async () => {
|
||||
ort.env.wasm.numThreads = 8;
|
||||
const dtype = resolveDefaultDtype(PARAKEET_MODEL_ID, "wasm");
|
||||
if (typeof dtype !== "string") throw new Error("expected one WASM dtype");
|
||||
return createParakeetEngine(
|
||||
ort as unknown as OrtLike,
|
||||
{
|
||||
encoder: await modelFile(parakeetFileName("encoder", dtype)),
|
||||
decoder: await modelFile(parakeetFileName("decoder", dtype)),
|
||||
vocab: new TextDecoder().decode(await modelFile(PARAKEET_VOCAB_FILE)),
|
||||
},
|
||||
parakeetSessionOptions("wasm"),
|
||||
);
|
||||
})();
|
||||
return enginePromise;
|
||||
}
|
||||
|
||||
afterAll(async () => {
|
||||
if (enginePromise) await (await enginePromise).dispose();
|
||||
});
|
||||
|
||||
describe("Parakeet engine accuracy (real model, onnxruntime-web WASM)", () => {
|
||||
for (const clip of CLIPS) {
|
||||
test.skipIf(!ENABLED)(
|
||||
`${clip.base}: whole-file and live transcripts stay close to the reference`,
|
||||
async () => {
|
||||
const reference = await Bun.file(join(AUDIO_DIR, `${clip.base}.txt`)).text();
|
||||
const audio = await decode16k(join(AUDIO_DIR, clip.file));
|
||||
const duration = audio.length / RATE;
|
||||
const engine = await loadEngine();
|
||||
|
||||
const fileStart = performance.now();
|
||||
let lastProgress = 0;
|
||||
const whole = await engine.transcribe(audio.slice(), {
|
||||
onWindow: (_chunks, processed) => {
|
||||
expect(processed).toBeGreaterThan(lastProgress);
|
||||
lastProgress = processed;
|
||||
},
|
||||
});
|
||||
const fileMs = performance.now() - fileStart;
|
||||
const fileWer = wordErrorRate(normalize(reference), normalize(whole.text));
|
||||
expect(lastProgress).toBeCloseTo(duration, 3);
|
||||
expect(whole.chunks[0]!.timestamp[0]).toBeLessThan(duration / 2);
|
||||
expect(whole.chunks.at(-1)!.timestamp[1]!).toBeGreaterThan(duration / 2);
|
||||
expect(whole.chunks.at(-1)!.timestamp[1]!).toBeLessThanOrEqual(duration + 1);
|
||||
|
||||
// Live: production segmenter, clock advancing with the audio fed.
|
||||
const transcribe: LiveTranscribeFn = (segment) => engine.transcribe(segment);
|
||||
let fed = 0;
|
||||
const liveStart = performance.now();
|
||||
const live = new LiveTranscriber({
|
||||
...LIVE_TUNING,
|
||||
transcribe,
|
||||
onUpdate: () => {},
|
||||
now: () => fed / (RATE / 1000),
|
||||
});
|
||||
const block = RATE / 10;
|
||||
for (let i = 0; i < audio.length; i += block) {
|
||||
live.push(audio.slice(i, Math.min(i + block, audio.length)));
|
||||
fed += Math.min(block, audio.length - i);
|
||||
await new Promise((resolve) => setTimeout(resolve, 0));
|
||||
}
|
||||
const liveResult = await live.stop();
|
||||
const liveMs = performance.now() - liveStart;
|
||||
const liveWer = wordErrorRate(normalize(reference), normalize(liveResult.text));
|
||||
|
||||
console.info(
|
||||
`[parakeet-real] ${clip.base}: ${duration.toFixed(1)} s audio — file WER ` +
|
||||
`${fileWer.toFixed(3)} in ${(fileMs / 1000).toFixed(1)} s; live WER ` +
|
||||
`${liveWer.toFixed(3)} in ${(liveMs / 1000).toFixed(1)} s`,
|
||||
);
|
||||
expect(fileWer).toBeLessThanOrEqual(clip.maxWer);
|
||||
expect(liveWer).toBeLessThanOrEqual(clip.maxWer);
|
||||
expect(liveWer).toBeLessThanOrEqual(fileWer + 0.1);
|
||||
},
|
||||
1_800_000,
|
||||
);
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,31 @@
|
||||
import { describe, expect, test } from "bun:test";
|
||||
import { join } from "node:path";
|
||||
import { ortAliasPlugin } from "../scripts/ortAlias";
|
||||
|
||||
// D-WS7a: the Parakeet path must not bring a second onnxruntime into the
|
||||
// inference worker. parakeet.js depends on its own onnxruntime-web (1.24.1);
|
||||
// the app runs the one Transformers.js pins. Unminified bundles carry a
|
||||
// `// node_modules/<path>` comment per bundled module, so count runtimes.
|
||||
|
||||
const WORKER = join(import.meta.dir, "../src/worker/inference.worker.ts");
|
||||
|
||||
async function bundledRuntimes(plugins: Bun.BunPlugin[]): Promise<string[]> {
|
||||
const build = await Bun.build({ entrypoints: [WORKER], target: "browser", format: "esm", plugins });
|
||||
expect(build.success).toBe(true);
|
||||
const code = await build.outputs[0]!.text();
|
||||
return [...code.matchAll(/^\/\/ (node_modules\/\S*onnxruntime-web\/dist\/\S+)$/gm)].map((m) => m[1]!);
|
||||
}
|
||||
|
||||
describe("inference worker bundle", () => {
|
||||
test("with the alias plugin it holds only the Transformers.js onnxruntime", async () => {
|
||||
expect(await bundledRuntimes([ortAliasPlugin])).toEqual([
|
||||
"node_modules/onnxruntime-web/dist/ort.webgpu.bundle.min.mjs",
|
||||
]);
|
||||
});
|
||||
|
||||
test("without it, parakeet.js's own runtime would be bundled too (the plugin is load-bearing)", async () => {
|
||||
const runtimes = await bundledRuntimes([]);
|
||||
expect(runtimes).toHaveLength(2);
|
||||
expect(runtimes.some((path) => path.includes("parakeet.js/node_modules/"))).toBe(true);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user