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

import { afterAll, describe, test, expect } from "bun:test";
import { existsSync, readdirSync } from "node:fs";
import { join } from "node:path";
import ffmpegPath from "ffmpeg-static";
import { env, pipeline } from "@huggingface/transformers";
import { StreamingResampler } from "../src/lib/audio";
import {
  LIVE_TUNING,
  LiveTranscriber,
  type LiveTranscribeFn,
} from "../src/lib/live";
import { resolveDefaultDtype } from "../src/lib/models";
import { wordErrorRate } from "./dsp";

/**
 * Real-model, real-speech accuracy of the live pipeline, measured as WER over
 * every recording in tests/audio_files. tests/live-accuracy.test.ts proves
 * the pipeline moves audio losslessly with a deterministic oracle; this test
 * closes the remaining gap by running actual Whisper weights end to end over
 * the exact mic path: 48 kHz PCM → StreamingResampler (as src/lib/mic.ts) →
 * LiveTranscriber with production LIVE_TUNING → the same Transformers.js
 * pipeline options the inference worker sends (src/worker/inference.worker.ts).
 *
 * Adding a clip: drop `<lang>_<name>.<m4a|mp3|wav|…>` plus a reference
 * transcript `<lang>_<name>.txt` into tests/audio_files — it is discovered
 * and held to the same bar automatically. `<lang>` is the Whisper language
 * code the UI user would pick (e.g. `da`, `en`); it is pinned per clip
 * because "auto" can misdetect (the Danish clip comes back as a rough English
 * translation, in whole-file mode too — see tests/live.md). Audio is decoded
 * at test time with ffmpeg-static (trusted devDependency), so any container
 * ffmpeg reads is fine.
 *
 * Model: the catalog's Whisper Large v3 Turbo at its WASM-default dtype (q4,
 * D-006) — the model quality-sensitive users should pick. Under bun test it
 * runs on onnxruntime-node (CPU): the same weights and quantization as the
 * browser WASM engine, just a different execution provider. The first run
 * downloads ~600 MB into node_modules/@huggingface/transformers/.cache/;
 * when the model is neither cached nor reachable on the Hub the accuracy
 * tests skip, so offline development stays green (this repo is offline-first).
 */

const MODEL_ID = "onnx-community/whisper-large-v3-turbo";
const MIC_RATE = 48_000;
const RATE = 16_000;
/**
 * "Close enough" for a quantized model on real-world speech: the Danish clip
 * measures ≈ 0.11 live and the English one lower; the slack absorbs
 * onnxruntime version-to-version drift, not real regressions — one mid-word
 * cut or dropped segment on these short clips moves WER far more than this.
 */
const MAX_WER = 0.25;

/**
 * Per-clip absolute WER bars for recordings where even whole-file mode with
 * this model cannot reach MAX_WER, so the bar would measure the clip, not
 * the pipeline. Keep these rare and record the measurement that justifies
 * them; the relative "live ≤ file + 0.1" check still applies to every clip.
 *
 * - en_alien: 15-min gaming livestream (music, game voices, crosstalk) added
 *   as the D-021 long-file fixture. Measured 2026-10-29: whole-file 0.261,
 *   live 0.251 (tiny.en whole-file: 0.315, D-021).
 */
const MAX_WER_BY_CLIP: Record<string, number> = { en_alien: 0.35 };

/**
 * Time budget: headroom for the first run's ~600 MB weight download plus
 * time proportional to the clip. Large v3 Turbo on CPU measured ≈ 1.34 s per
 * second of audio for the whole-file + live passes combined (en_alien,
 * 888 s of audio → 1189 s); 3 s per second leaves ~2× margin.
 */
const DOWNLOAD_BUDGET_MS = 900_000;
const MS_PER_AUDIO_SECOND = 3_000;

// ---------------------------------------------------------------------------
// Clip discovery: every reference transcript must have audio, and vice versa
// ---------------------------------------------------------------------------

const AUDIO_DIR = join(import.meta.dir, "audio_files");

interface Clip {
  /** `<lang>_<name>`, e.g. "da_festuge". */
  base: string;
  language: string;
  audioFile: string;
  transcriptFile: string;
}

/** Clip length from ffmpeg's header probe (no decode); sizes the time budget. */
function probeDurationSeconds(path: string): number {
  const { stderr } = Bun.spawnSync(
    [ffmpegPath as unknown as string, "-hide_banner", "-i", path],
    { stdout: "ignore", stderr: "pipe" },
  );
  const match = /Duration: (\d+):(\d+):(\d+(?:\.\d+)?)/.exec(stderr.toString());
  if (!match) throw new Error(`ffmpeg could not read the duration of ${path}`);
  return Number(match[1]) * 3600 + Number(match[2]) * 60 + Number(match[3]);
}

function discoverClips(): { clips: Clip[]; problems: string[] } {
  const entries = readdirSync(AUDIO_DIR, { withFileTypes: true })
    .filter((entry) => entry.isFile())
    .map((entry) => entry.name);
  const clips: Clip[] = [];
  const problems: string[] = [];
  const claimed = new Set<string>();

  for (const transcriptFile of entries.filter((name) => name.endsWith(".txt"))) {
    const base = transcriptFile.slice(0, -".txt".length);
    const language = /^([a-z]{2,3})_./.exec(base)?.[1];
    const audioFile = entries.find(
      (name) => name !== transcriptFile && name.startsWith(`${base}.`),
    );
    if (!language) {
      problems.push(
        `${transcriptFile}: name clips "<lang>_<name>.<ext>" (Whisper language code prefix, e.g. da_, en_) so the test knows which language to pin.`,
      );
      continue;
    }
    if (!audioFile) {
      problems.push(`${transcriptFile}: no matching audio file "${base}.<ext>".`);
      continue;
    }
    claimed.add(transcriptFile);
    claimed.add(audioFile);
    clips.push({ base, language, audioFile, transcriptFile });
  }

  for (const name of entries) {
    if (!claimed.has(name) && !name.endsWith(".txt")) {
      problems.push(
        `${name}: audio without a reference transcript — add "${name.replace(/\.[^.]+$/, "")}.txt" so accuracy is regression-tested.`,
      );
    }
  }
  return { clips, problems };
}

const { clips, problems } = discoverClips();

/** Decode any audio container to 48 kHz mono float32 PCM (the mic rate). */
async function decodeAudio(path: string): Promise<Float32Array> {
  const proc = Bun.spawn(
    [
      ffmpegPath as unknown as string,
      ...["-hide_banner", "-loglevel", "error"],
      ...["-i", path],
      ...["-ac", "1", "-ar", String(MIC_RATE), "-f", "f32le", "-"],
    ],
    { stdout: "pipe", stderr: "pipe" },
  );
  const [bytes, errors, exitCode] = await Promise.all([
    new Response(proc.stdout).arrayBuffer(),
    new Response(proc.stderr).text(),
    proc.exited,
  ]);
  if (exitCode !== 0) throw new Error(`ffmpeg failed on ${path}: ${errors}`);
  return new Float32Array(bytes);
}

/**
 * WER normalization: casing, punctuation, and apostrophes are not
 * transcription accuracy. "Århus"/"Aarhus" are the same word in Danish
 * orthography, so they unify.
 */
function normalize(text: string): string {
  return text
    .toLowerCase()
    .replaceAll("århus", "aarhus")
    .replace(/['’]/g, "")
    .replace(/[.,!?;:"“”‘‚«»()\-–—…]/g, " ")
    .replace(/\s+/g, " ")
    .trim();
}

/** Cached locally, or downloadable right now — otherwise accuracy tests skip. */
async function modelAvailable(id: string): Promise<boolean> {
  if (env.cacheDir && existsSync(join(env.cacheDir, id))) return true;
  try {
    const response = await fetch(
      `https://huggingface.co/${id}/resolve/main/config.json`,
      { method: "HEAD", signal: AbortSignal.timeout(5_000) },
    );
    return response.ok;
  } catch {
    return false;
  }
}

const available = await modelAvailable(MODEL_ID);

// The pipeline is loaded lazily by the first accuracy test and shared.
type Asr = Awaited<ReturnType<typeof pipeline<"automatic-speech-recognition">>>;
let asrPromise: Promise<Asr> | null = null;
function loadAsr(): Promise<Asr> {
  asrPromise ??= pipeline("automatic-speech-recognition", MODEL_ID, {
    dtype: resolveDefaultDtype(MODEL_ID, "wasm"),
  });
  return asrPromise;
}

afterAll(async () => {
  if (asrPromise) await (await asrPromise).dispose();
});

describe("live pipeline accuracy (real model, real recordings)", () => {
  test("every file in tests/audio_files forms a <lang>_<name> audio+transcript pair", () => {
    expect(problems).toEqual([]);
    expect(clips.length).toBeGreaterThan(0);
  });

  for (const clip of clips) {
    test.skipIf(!available)(
      `${clip.base} (${clip.language}): live mic-path transcript stays close to the reference`,
      async () => {
        const reference = await Bun.file(
          join(AUDIO_DIR, clip.transcriptFile),
        ).text();
        const micAudio = await decodeAudio(join(AUDIO_DIR, clip.audioFile));
        expect(micAudio.length).toBeGreaterThan(MIC_RATE); // ≥ 1 s of speech

        const asr = await loadAsr();
        const transcribe: LiveTranscribeFn = async (audio) => {
          // Mirrors src/worker/inference.worker.ts `transcribe` options with
          // the UI language set to the clip's language.
          const output = await asr(audio, {
            chunk_length_s: 30,
            stride_length_s: 5,
            return_timestamps: true,
            // Selects the single-pass generate path, like the worker (D-021).
            max_new_tokens: 440,
            language: clip.language,
            task: "transcribe",
          });
          const single = (Array.isArray(output) ? output[0] : output) as {
            text?: string;
            chunks?: { text: string; timestamp: [number, number | null] }[];
          };
          return { text: single.text ?? "", chunks: single.chunks ?? [] };
        };

        // Resample in ~100 ms mic-sized blocks exactly like src/lib/mic.ts.
        const resampler = new StreamingResampler(MIC_RATE, RATE);
        const blocks: Float32Array[] = [];
        const chunk = MIC_RATE / 10;
        for (let i = 0; i < micAudio.length; i += chunk) {
          const out = resampler.process(
            micAudio.slice(i, Math.min(i + chunk, micAudio.length)),
          );
          if (out.length > 0) blocks.push(out);
        }
        const tail = resampler.flush();
        if (tail.length > 0) blocks.push(tail);
        const totalSamples = blocks.reduce((n, b) => n + b.length, 0);
        const durationSeconds = totalSamples / RATE;

        // Whole-clip baseline: what the file-upload flow would produce from
        // the same audio and model. The live pipeline is then held to it.
        const whole = new Float32Array(totalSamples);
        let at = 0;
        for (const block of blocks) {
          whole.set(block, at);
          at += block.length;
        }
        const fileWer = wordErrorRate(
          normalize(reference),
          normalize((await transcribe(whole)).text),
        );

        // The live session, with the injected clock advancing with the audio
        // fed so pass scheduling behaves as in real time.
        let fedSamples = 0;
        let sawTentative = false;
        const live = new LiveTranscriber({
          ...LIVE_TUNING,
          transcribe,
          onUpdate: (update) => {
            if (update.tentativeText.trim().length > 0) sawTentative = true;
          },
          now: () => fedSamples / (RATE / 1000),
        });
        for (const block of blocks) {
          live.push(block);
          fedSamples += block.length;
          await new Promise((resolve) => setTimeout(resolve, 0));
        }
        const result = await live.stop();
        const liveWer = wordErrorRate(normalize(reference), normalize(result.text));

        expect(liveWer).toBeLessThanOrEqual(MAX_WER_BY_CLIP[clip.base] ?? MAX_WER);
        // Segmentation must not cost accuracy relative to whole-file mode.
        expect(liveWer).toBeLessThanOrEqual(fileWer + 0.1);

        // Live UX invariants, with a real model: tentative text streamed, and
        // the final chunks carry absolute session timestamps spanning the clip.
        expect(sawTentative).toBe(true);
        expect(result.chunks.length).toBeGreaterThanOrEqual(1);
        const first = result.chunks[0]!;
        const last = result.chunks[result.chunks.length - 1]!;
        expect(first.timestamp[0]).toBeGreaterThanOrEqual(0);
        expect(first.timestamp[0]).toBeLessThan(durationSeconds / 2);
        expect(last.timestamp[1] ?? last.timestamp[0]).toBeGreaterThan(
          durationSeconds / 2,
        );
        expect(last.timestamp[1] ?? last.timestamp[0]).toBeLessThanOrEqual(
          durationSeconds + 1,
        );
      },
      DOWNLOAD_BUDGET_MS +
        probeDurationSeconds(join(AUDIO_DIR, clip.audioFile)) * MS_PER_AUDIO_SECOND,
    );
  }
});