whisper_web/tests/models.test.ts

import { describe, test, expect } from "bun:test";
import {
  DANISH_MODELS_HF_OWNER,
  DEFAULT_MODEL_ID,
  MODELS,
  WHISPER_LANGUAGES,
  estimateDefaultDownloadBytes,
  estimateDownloadBytes,
  estimatePeakLoadBytes,
  formatModelSize,
  getFasterWhisperModel,
  getModel,
  modelFamily,
  PARAKEET_MODEL_ID,
  resolveDefaultDtype,
  taskForModel,
} 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)$/, "");
}

const WHISPER_MODELS = MODELS.filter((model) => modelFamily(model) === "whisper");

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",
      `${DANISH_MODELS_HF_OWNER}/whisper-small-danish-ONNX`,
      // Not Whisper: the D-WS7a Parakeet engine path, last so it never
      // becomes anyone's accidental default.
      "ysdede/parakeet-tdt-0.6b-v3-onnx",
    ]);
    expect(new Set(ids).size).toBe(ids.length);
  });

  test(".en variants are English-only; the rest are multilingual", () => {
    for (const model of WHISPER_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 Whisper model maps to a faster-whisper name for the server engine (D-013)", () => {
    // The Danish conversions (D-WS6a) have no faster-whisper alias; their
    // server-engine name is a Hub repo id — covered by the Danish suite below.
    for (const model of WHISPER_MODELS.filter((m) => !m.id.startsWith(`${DANISH_MODELS_HF_OWNER}/`))) {
      // 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("Danish fine-tune conversions (D-WS6a)", () => {
  const owner = `${DANISH_MODELS_HF_OWNER}/`;
  const danish = MODELS.filter((model) => model.id.startsWith(owner));

  test("only whisper-small-danish is published (the turbo fine-tunes showed no clear gain), under lucasw89", () => {
    expect(DANISH_MODELS_HF_OWNER).toBe("lucasw89");
    expect(danish.map((model) => model.id)).toEqual([`${owner}whisper-small-danish-ONNX`]);
  });

  test("point the server engine at the matching CTranslate2 int8 conversion under the same owner", () => {
    for (const model of danish) {
      const ct2 = model.id.replace(/-ONNX$/, "-ct2-int8");
      expect(model.fasterWhisperModel).toBe(ct2);
      expect(getFasterWhisperModel(model.id)).toBe(ct2);
    }
  });

  test("are multilingual (language/task are sent) and unverified on every device", () => {
    for (const model of danish) {
      expect(model.multilingual).toBe(true);
      for (const device of ["wasm", "webgpu", "webnn"] as const) {
        expect(model.defaults[device].verified).toBe(false);
      }
    }
  });

  test("share their base model's parameter counts (unchanged architecture)", () => {
    const small = getModel("onnx-community/whisper-small")!;
    const [smallDa] = danish;
    expect(smallDa!.encoderParamsM).toBe(small.encoderParamsM);
    expect(smallDa!.decoderParamsM).toBe(small.decoderParamsM);
  });

  test("default to the dtypes measured on the converted files (D-WS6a)", () => {
    // da_festuge on onnxruntime-node: q4 matched or beat q8 and was faster for
    // all three, small included; fp16 encoders did not overflow (D-022 check).
    for (const model of danish) {
      expect(resolveDefaultDtype(model.id, "wasm")).toBe("q4");
      expect(resolveDefaultDtype(model.id, "webgpu")).toEqual({
        encoder_model: "fp16",
        decoder_model_merged: "q4",
      });
      expect(resolveDefaultDtype(model.id, "webnn")).toBe("fp16");
      expect(model.defaults.wasm.note).toContain("D-WS6a");
    }
  });
});

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("estimatePeakLoadBytes (D-027)", () => {
  const medium = getModel("onnx-community/whisper-medium-ONNX")!;
  const turbo = getModel("onnx-community/whisper-large-v3-turbo")!;

  test("is ~2x the larger of the two weight files, not their sum", () => {
    const dtype = resolveDefaultDtype(medium.id, "wasm"); // q4
    const peak = estimatePeakLoadBytes(medium, dtype);
    const download = estimateDownloadBytes(medium, dtype);
    // The combined download sums both files; the peak only doubles the
    // larger one — smaller than 2x the sum, but bigger than the sum alone
    // once the sequential-load patch keeps only one file resident at a time.
    expect(peak).toBeLessThan(2 * download);
    expect(peak).toBeGreaterThan(download);
  });

  test("mixed per-part dtypes (fp16 encoder + q4 decoder) peak on the larger encoder", () => {
    const dtype = resolveDefaultDtype(turbo.id, "webgpu"); // fp16 encoder, q4 decoder
    const peak = estimatePeakLoadBytes(turbo, dtype);
    // large-v3-turbo's encoder (635M) dwarfs its decoder (174M) even before
    // fp16 vs q4 is factored in, so the encoder alone should set the peak.
    const encoderOnlyPeak = 2 * turbo.encoderParamsM * 1e6 * 2; // fp16 = 2 bytes/param
    expect(peak).toBeCloseTo(encoderOnlyPeak, -5);
  });

  test("scales up with a larger model at the same dtype", () => {
    const tiny = getModel("onnx-community/whisper-tiny.en")!;
    expect(estimatePeakLoadBytes(medium, "q4")).toBeGreaterThan(
      estimatePeakLoadBytes(tiny, "q4"),
    );
  });
});

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);
  });
});

describe("canTranslate (D-WS4a)", () => {
  test("Whisper translate is offered only where OpenAI trained it", () => {
    const translating = MODELS.filter((model) => model.canTranslate).map((model) => model.id);
    // large-v3-turbo was not trained for translation (OpenAI model card);
    // `.en` models are English-only and have no task token.
    expect(translating).toEqual([
      "onnx-community/whisper-tiny",
      "onnx-community/whisper-base",
      "onnx-community/whisper-small",
      "onnx-community/whisper-medium-ONNX",
    ]);
  });

  test("canTranslate implies multilingual", () => {
    for (const model of MODELS) {
      if (model.canTranslate) expect(model.multilingual).toBe(true);
    }
  });
});

describe("taskForModel (D-WS4a)", () => {
  test("a stale translate resets to transcribe on a model that cannot translate", () => {
    expect(taskForModel("onnx-community/whisper-large-v3-turbo", "translate")).toBe("transcribe");
    expect(taskForModel("onnx-community/whisper-small.en", "translate")).toBe("transcribe");
  });

  test("translate survives on a model that can translate", () => {
    expect(taskForModel("onnx-community/whisper-small", "translate")).toBe("translate");
    expect(taskForModel("onnx-community/whisper-small", "transcribe")).toBe("transcribe");
  });

  test("unknown ids never translate", () => {
    expect(taskForModel("someone/unknown", "translate")).toBe("transcribe");
  });
});

describe("Parakeet catalog entry (D-WS7a)", () => {
  const parakeet = getModel(PARAKEET_MODEL_ID)!;

  test("family defaults to whisper when absent", () => {
    expect(modelFamily(getModel(DEFAULT_MODEL_ID))).toBe("whisper");
    expect(modelFamily(undefined)).toBe("whisper");
    expect(modelFamily(parakeet)).toBe("parakeet");
  });

  test("is English-only and browser-only (no faster-whisper name)", () => {
    expect(parakeet.multilingual).toBe(false);
    expect(parakeet.fasterWhisperModel).toBeUndefined();
    expect(getFasterWhisperModel(PARAKEET_MODEL_ID)).toBeUndefined();
  });

  test("WASM runs int8 for both parts; WebGPU is fp16 encoder + int8 decoder", () => {
    expect(resolveDefaultDtype(PARAKEET_MODEL_ID, "wasm")).toBe("q8");
    expect(resolveDefaultDtype(PARAKEET_MODEL_ID, "webgpu")).toEqual({
      encoder_model: "fp16",
      decoder_model_merged: "q8",
    });
  });

  test("download size is the real file sizes, not a parameter estimate", () => {
    expect(estimateDefaultDownloadBytes(parakeet, "wasm")).toBe(
      652_183_999 + 18_202_004 + 102_132,
    );
    expect(estimateDefaultDownloadBytes(parakeet, "webgpu")).toBe(
      1_238_960_452 + 18_202_004 + 102_132,
    );
    expect(formatModelSize(estimateDefaultDownloadBytes(parakeet, "wasm"))).toBe("670 MB");
  });

  test("peak load memory is twice the larger weight file", () => {
    expect(estimatePeakLoadBytes(parakeet, "q8")).toBe(2 * 652_183_999);
  });
});