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:
1 parent
c5a57cd33a
commit
708fae5778
10 files changed
+485
-7
No files matched your search
@@ -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?.();
|
||||
},
|
||||
};
|
||||
}
|
||||
Reference in new issue
Block a user