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

import { describe, test, expect } from "bun:test";
import { existsSync } from "node:fs";
import { join } from "node:path";
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
 * a real Danish recording (tests/audio_files/Festuge.m4a, reference
 * transcript Festuge.txt). 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).
 *
 * Model: the catalog's Whisper Large v3 Turbo at its WASM-default dtype (q4,
 * D-006) — the model a Danish user should pick. Under bun test it runs on
 * onnxruntime-node (CPU), i.e. the same weights and quantization as the
 * browser WASM engine, just a different execution provider. `language` is
 * pinned to "da" exactly as the worker does when the user selects Danish in
 * the UI; "auto" misdetects this clip as English (whole-file mode too) and
 * comes back as a rough English translation — see tests/live.md.
 *
 * The first run downloads ~600 MB into @huggingface/transformers' cache
 * (node_modules/@huggingface/transformers/.cache/). When the model is neither
 * cached nor reachable on the Hub, the test skips so offline development
 * stays green (this repo is offline-first).
 *
 * Fixture: Festuge.48k.wav is Festuge.m4a decoded once, so the test needs no
 * audio decoder:  ffmpeg -i Festuge.m4a -ac 1 -ar 48000 -sample_fmt s16 \
 *   Festuge.48k.wav
 */

const MODEL_ID = "onnx-community/whisper-large-v3-turbo";
const MIC_RATE = 48_000;
const RATE = 16_000;

/**
 * WER normalization: casing and punctuation 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(/\s+/g, " ")
    .trim();
}

/** Minimal RIFF/WAVE reader for the fixture (16-bit PCM mono). */
function readWav(buffer: ArrayBuffer): { samples: Float32Array; rate: number } {
  const view = new DataView(buffer);
  const tag = (offset: number) =>
    String.fromCharCode(
      view.getUint8(offset),
      view.getUint8(offset + 1),
      view.getUint8(offset + 2),
      view.getUint8(offset + 3),
    );
  if (tag(0) !== "RIFF" || tag(8) !== "WAVE") throw new Error("not a WAV file");

  let rate = 0;
  let offset = 12;
  while (offset + 8 <= view.byteLength) {
    const id = tag(offset);
    const size = view.getUint32(offset + 4, true);
    if (id === "fmt ") {
      const format = view.getUint16(offset + 8, true);
      const channels = view.getUint16(offset + 10, true);
      const bits = view.getUint16(offset + 22, true);
      if (format !== 1 || channels !== 1 || bits !== 16) {
        throw new Error("fixture must be 16-bit PCM mono");
      }
      rate = view.getUint32(offset + 12, true);
    } else if (id === "data") {
      const pcm = new Int16Array(buffer, offset + 8, size / 2);
      const samples = new Float32Array(pcm.length);
      for (let i = 0; i < pcm.length; i++) samples[i] = pcm[i]! / 32768;
      return { samples, rate };
    }
    offset += 8 + size + (size % 2);
  }
  throw new Error("WAV data chunk not found");
}

/** Cached locally, or downloadable right now — otherwise the test skips. */
async function modelAvailable(id: string): Promise<boolean> {
  if (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);

describe("live pipeline accuracy (real model, real Danish speech)", () => {
  test.skipIf(!available)(
    "Festuge.m4a over the mic path stays close to the reference transcript",
    async () => {
      const dir = join(import.meta.dir, "audio_files");
      const reference = await Bun.file(join(dir, "Festuge.txt")).text();
      const wav = readWav(await Bun.file(join(dir, "Festuge.48k.wav")).arrayBuffer());
      expect(wav.rate).toBe(MIC_RATE);

      const asr = await pipeline("automatic-speech-recognition", MODEL_ID, {
        dtype: resolveDefaultDtype(MODEL_ID, "wasm"),
      });
      try {
        const transcribe: LiveTranscribeFn = async (audio) => {
          // Mirrors src/worker/inference.worker.ts `transcribe` options with
          // the UI set to Danish.
          const output = await asr(audio, {
            chunk_length_s: 30,
            stride_length_s: 5,
            return_timestamps: true,
            language: "da",
            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 < wav.samples.length; i += chunk) {
          const out = resampler.process(
            wav.samples.slice(i, Math.min(i + chunk, wav.samples.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));

        // "Close enough": measured 0.11 (live) / 0.14 (whole-file) on this
        // clip; the margin absorbs onnxruntime version-to-version drift, not
        // real regressions — a single mid-word cut or dropped segment on a
        // 64-word reference moves WER by far more than the slack here.
        expect(liveWer).toBeLessThanOrEqual(0.25);
        // 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(3);
        const first = result.chunks[0]!;
        const last = result.chunks[result.chunks.length - 1]!;
        expect(first.timestamp[0]).toBeGreaterThan(0.5); // speech starts ~1.8 s in
        expect(first.timestamp[0]).toBeLessThan(5);
        expect(last.timestamp[1] ?? last.timestamp[0]).toBeGreaterThan(
          durationSeconds - 5,
        );
        expect(last.timestamp[1] ?? last.timestamp[0]).toBeLessThanOrEqual(
          durationSeconds + 1,
        );
      } finally {
        await asr.dispose();
      }
    },
    900_000, // first run downloads ~600 MB of weights
  );
});