whisper_web/tests/longform-transformers.test.ts

import { afterEach, describe, expect, test } from "bun:test";
import { existsSync } from "node:fs";
import { join } from "node:path";
import { env } from "@huggingface/transformers";
import { createWhisperDeps, loadSileroVad, VAD_MODEL_ID } from "../src/worker/longform/transformers";
import { PROMPT_PREFIX_KWARG } from "../src/worker/prompt";
import { getSpeechTimestamps, SAMPLE_RATE } from "../src/worker/longform/vad";
import { decodeClip } from "./clips";

const realFetch = globalThis.fetch;
afterEach(() => {
  globalThis.fetch = realFetch;
});

describe("loadSileroVad offline (plan 001-0 Review Focus 3)", () => {
  test("localOnly load with nothing cached rejects without calling fetch", async () => {
    let fetches = 0;
    globalThis.fetch = (async () => {
      fetches += 1;
      throw new Error("network must not be touched");
    }) as unknown as typeof fetch;
    await expect(
      loadSileroVad({ localOnly: true, isCached: async () => false }),
    ).rejects.toThrow(/not downloaded/);
    expect(fetches).toBe(0);
  });
});

/** Cached locally, or downloadable right now — otherwise the real-VAD tests skip. */
async function vadAvailable(): Promise<boolean> {
  if (env.cacheDir && existsSync(join(env.cacheDir, VAD_MODEL_ID))) return true;
  try {
    const response = await fetch(
      `https://huggingface.co/${VAD_MODEL_ID}/resolve/main/onnx/model.onnx`,
      { method: "HEAD", signal: AbortSignal.timeout(5_000) },
    );
    return response.ok;
  } catch {
    return false;
  }
}
const available = await vadAvailable();

describe.skipIf(!available)("Silero VAD on the Transformers.js runtime (real model)", () => {
  const seconds = 10;
  const tone = new Float32Array(seconds * SAMPLE_RATE).map(
    (_, i) => 0.3 * Math.sin((2 * Math.PI * 440 * i) / SAMPLE_RATE),
  );
  let seed = 1;
  const noise = new Float32Array(seconds * SAMPLE_RATE).map(() => {
    seed = (seed * 1_103_515_245 + 12_345) % 2 ** 31;
    return 0.2 * (seed / 2 ** 30 - 1);
  });

  for (const [name, audio] of [
    ["silence", new Float32Array(seconds * SAMPLE_RATE)],
    ["a 440 Hz tone", tone],
    ["white noise", noise],
  ] as const) {
    test(`${name} contains no speech`, async () => {
      const vad = await loadSileroVad({ localOnly: false });
      const probs = await vad.speechProbabilities(audio);
      expect(probs.length).toBe(Math.ceil(audio.length / 512));
      expect(getSpeechTimestamps(probs, audio.length)).toEqual([]);
    });
  }

  test("en_aiproblem: speech is found and covers most of the clip", async () => {
    const audio = await decodeClip("en_aiproblem.m4a");
    const vad = await loadSileroVad({ localOnly: false });
    const spans = getSpeechTimestamps(await vad.speechProbabilities(audio), audio.length);
    expect(spans.length).toBeGreaterThan(0);
    const covered = spans.reduce((n, s) => n + s.end - s.start, 0);
    expect(covered / audio.length).toBeGreaterThan(0.8);
  });
});

describe("createWhisperDeps decode with the names-and-terms hint (D-WS2a × D-WS3a)", () => {
  /** A fake ASR pipeline that records the generate options of each call. */
  function fakeAsr() {
    const calls: Record<string, unknown>[] = [];
    const asr = Object.assign(
      async (_audio: Float32Array, options: Record<string, unknown>) => {
        calls.push(options);
        return { text: " hello", chunks: [] };
      },
      {
        tokenizer: {
          encode: (text: string) => [...text].map((char) => char.charCodeAt(0)),
          convert_tokens_to_ids: (token: string) => (token === "<|startofprev|>" ? 50361 : -1),
        },
        model: {
          forward_params: ["input_features"],
          generate: async () => ({}),
          _prepare_generation_config: () => ({}),
          _retrieve_init_tokens: () => [50258],
        },
      },
    );
    return { asr, calls };
  }

  const baseRequest = { temperature: 0, noRepeatNgramSize: 0, repetitionPenalty: 1 };

  test("passes the hint as a decoder-prompt prefix on every chunk, and nothing without one", async () => {
    const { asr, calls } = fakeAsr();
    const deps = createWhisperDeps(asr as never, { multilingual: false });
    await deps.decode(new Float32Array(16), { ...baseRequest, prompt: "Maxwell" });
    await deps.decode(new Float32Array(16), { ...baseRequest, prompt: "Maxwell" });
    await deps.decode(new Float32Array(16), baseRequest);

    const prefix = calls[0]![PROMPT_PREFIX_KWARG] as number[];
    expect(prefix[0]).toBe(50361); // <|startofprev|>
    expect(prefix.length).toBeGreaterThan(1);
    expect(calls[1]![PROMPT_PREFIX_KWARG]).toEqual(prefix);
    expect(PROMPT_PREFIX_KWARG in calls[2]!).toBe(false);
  });
});