whisper_web/src/worker/parakeetEngine.ts

/**
 * 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?.();
    },
  };
}