au_whisper_web/tests/models.test.ts

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

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

/** Hub repo-name decorations that follow the Whisper size suffix. */
function stripRepoDecorations(id: string): string {
  return id.replace(/(_timestamped|-ONNX)$/, "");
}

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_timestamped",
      "onnx-community/whisper-medium-ONNX",
      "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) {
      // Hub repo names may carry decorations (e.g. "_timestamped", "-ONNX")
      // after the size suffix — judge English-ness on the undecorated name.
      expect(model.multilingual).toBe(!stripRepoDecorations(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();
  });

  test("every model maps to a faster-whisper name for the server engine (D-013)", () => {
    for (const model of MODELS) {
      // The onnx-community ids and faster-whisper's aliases share the same
      // size suffix ("tiny.en", …, "large-v3-turbo") once repo-name
      // decorations ("_timestamped", "-ONNX") are stripped; keeping them
      // aligned is what lets one catalog drive both engines.
      expect(model.fasterWhisperModel).toBe(
        stripRepoDecorations(model.id).replace("onnx-community/whisper-", ""),
      );
    }
    expect(getFasterWhisperModel(DEFAULT_MODEL_ID)).toBe("tiny.en");
    // Unknown ids stay undefined so /api/transcribe can answer 400 instead of
    // forwarding arbitrary strings to the sidecar.
    expect(getFasterWhisperModel("someone/some-model")).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-ONNX", "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("medium never defaults to fp16 on any backend (D-022: broken conversions)", () => {
    for (const id of [
      "onnx-community/whisper-medium.en_timestamped",
      "onnx-community/whisper-medium-ONNX",
    ]) {
      for (const device of ["wasm", "webgpu", "webnn"] as const) {
        expect(resolveDefaultDtype(id, device)).toBe("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);
  });
});