whisper_web/tests/models.test.ts

import { describe, test, expect } from "bun:test";
import {
  DEFAULT_MODEL_ID,
  MODELS,
  WHISPER_LANGUAGES,
  estimateDefaultDownloadBytes,
  estimateDownloadBytes,
  formatModelSize,
  getModel,
  resolveDefaultDtype,
} from "../src/lib/models";

// Pure catalog data (D-006) — no browser, no network.

describe("MODELS catalog", () => {
  test("contains the curated D-006 lineup from the onnx-community org", () => {
    const ids = MODELS.map((model) => model.id);
    expect(ids).toEqual([
      "onnx-community/whisper-tiny.en",
      "onnx-community/whisper-tiny",
      "onnx-community/whisper-base.en",
      "onnx-community/whisper-base",
      "onnx-community/whisper-small.en",
      "onnx-community/whisper-small",
      "onnx-community/whisper-medium.en",
      "onnx-community/whisper-medium",
      "onnx-community/whisper-large-v3-turbo",
    ]);
    expect(new Set(ids).size).toBe(ids.length);
  });

  test(".en variants are English-only; the rest are multilingual", () => {
    for (const model of MODELS) {
      expect(model.multilingual).toBe(!model.id.endsWith(".en"));
    }
  });

  test("every model has dtype defaults with notes for both devices", () => {
    for (const model of MODELS) {
      for (const device of ["wasm", "webgpu"] as const) {
        const entry = model.defaults[device];
        expect(entry.dtype).toBeDefined();
        expect(entry.note.length).toBeGreaterThan(0);
        expect(typeof entry.verified).toBe("boolean");
      }
    }
  });

  test("only the M1-tested combination is marked verified", () => {
    const verified = MODELS.flatMap((model) =>
      (["wasm", "webgpu"] as const)
        .filter((device) => model.defaults[device].verified)
        .map((device) => `${model.id}@${device}`),
    );
    expect(verified).toEqual(["onnx-community/whisper-tiny.en@wasm"]);
  });

  test("the default model exists in the catalog", () => {
    expect(getModel(DEFAULT_MODEL_ID)).toBeDefined();
    expect(getModel("onnx-community/whisper-nonexistent")).toBeUndefined();
  });
});

describe("resolveDefaultDtype", () => {
  test("WASM defaults: q8 for small models, q4 for large ones", () => {
    expect(resolveDefaultDtype("onnx-community/whisper-tiny.en", "wasm")).toBe("q8");
    expect(resolveDefaultDtype("onnx-community/whisper-small", "wasm")).toBe("q8");
    expect(resolveDefaultDtype("onnx-community/whisper-medium", "wasm")).toBe("q4");
    expect(resolveDefaultDtype("onnx-community/whisper-large-v3-turbo", "wasm")).toBe("q4");
  });

  test("WebGPU defaults: fp32 for tiny/base (fp16 known-bad), fp16+q4 above", () => {
    expect(resolveDefaultDtype("onnx-community/whisper-tiny", "webgpu")).toEqual({
      encoder_model: "fp32",
      decoder_model_merged: "fp32",
    });
    expect(resolveDefaultDtype("onnx-community/whisper-base.en", "webgpu")).toEqual({
      encoder_model: "fp32",
      decoder_model_merged: "fp32",
    });
    expect(
      resolveDefaultDtype("onnx-community/whisper-large-v3-turbo", "webgpu"),
    ).toEqual({ encoder_model: "fp16", decoder_model_merged: "q4" });
  });

  test("unknown model ids fall back to q8", () => {
    expect(resolveDefaultDtype("someone/some-model", "wasm")).toBe("q8");
  });
});

describe("size estimates", () => {
  const tiny = getModel("onnx-community/whisper-tiny.en")!;
  const turbo = getModel("onnx-community/whisper-large-v3-turbo")!;

  test("scale with quantization level", () => {
    const fp32 = estimateDownloadBytes(tiny, "fp32");
    const fp16 = estimateDownloadBytes(tiny, "fp16");
    const q8 = estimateDownloadBytes(tiny, "q8");
    const q4 = estimateDownloadBytes(tiny, "q4");
    expect(fp32).toBeGreaterThan(fp16);
    expect(fp16).toBeGreaterThan(q8);
    expect(q8).toBeGreaterThan(q4);
  });

  test("per-part dtypes are priced per part", () => {
    const mixed = estimateDownloadBytes(turbo, {
      encoder_model: "fp16",
      decoder_model_merged: "q4",
    });
    expect(mixed).toBeLessThan(estimateDownloadBytes(turbo, "fp16"));
    expect(mixed).toBeGreaterThan(estimateDownloadBytes(turbo, "q4"));
  });

  test("default-dtype estimates are in a plausible range", () => {
    // tiny.en on WASM (q8, ~39M params) — tens of MB.
    const tinyWasm = estimateDefaultDownloadBytes(tiny, "wasm");
    expect(tinyWasm).toBeGreaterThan(20e6);
    expect(tinyWasm).toBeLessThan(100e6);
    // large-v3-turbo on WebGPU (fp16 encoder + q4 decoder) — around 1.4 GB.
    const turboWebgpu = estimateDefaultDownloadBytes(turbo, "webgpu");
    expect(turboWebgpu).toBeGreaterThan(1e9);
    expect(turboWebgpu).toBeLessThan(2e9);
  });

  test("formatModelSize renders MB below 1 GB and GB above", () => {
    expect(formatModelSize(42_000_000)).toBe("42 MB");
    expect(formatModelSize(1_400_000_000)).toBe("1.4 GB");
  });
});

describe("WHISPER_LANGUAGES", () => {
  test("codes are unique, lowercase, and include common languages", () => {
    const codes = WHISPER_LANGUAGES.map((entry) => entry.code);
    expect(new Set(codes).size).toBe(codes.length);
    for (const code of codes) expect(code).toMatch(/^[a-z]{2,3}$/);
    expect(codes).toContain("en");
    expect(codes).toContain("da");
  });

  test("does not contain the auto-detect sentinel (UI adds it)", () => {
    expect(WHISPER_LANGUAGES.some((entry) => entry.code === "auto")).toBe(false);
  });
});