whisper_web/tests/parakeet-real-model.test.ts

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,
    );
  }
});