whisper_web/src/lib/models.ts

/**
 * Curated Whisper model catalog (Decisions.md D-006).
 *
 * Pure, typed data — no browser globals — so it is unit-testable with
 * `bun test` and shareable with the server engine (M5). ONNX conversions come
 * from the `onnx-community` Hub org (the recommended source for
 * Transformers.js v3+).
 *
 * dtype defaults per device follow the community-known-good combinations:
 * - WASM: q8 for small models, q4 for the large ones (download size).
 * - WebGPU: fp16 encoder + q4 merged decoder for the big models; fp32 for
 *   tiny/base where fp16 has known accuracy issues.
 * `verified` is only true for combinations we have actually run in a browser;
 * the rest are pending manual testing (WorkBreakdown M2).
 */

import type {
  DevicePreference,
  DtypeSpec,
  InferenceDevice,
  WhisperDtype,
  WhisperTask,
} from "../worker/protocol";

// ---------------------------------------------------------------------------
// Catalog types
// ---------------------------------------------------------------------------

/** Dtype default for one device, with verification status (D-006). */
export interface DtypeDefault {
  dtype: DtypeSpec;
  /** True only if this exact combination has been run in a real browser. */
  verified: boolean;
  /** Why this combination was chosen / caveats. */
  note: string;
}

/**
 * Which inference path a catalog entry runs on (D-051). Whisper models use
 * the Transformers.js ASR pipeline; Parakeet uses the worker's own TDT path
 * (src/worker/parakeetEngine.ts). Absent ⇒ "whisper", so the Whisper entries
 * carry no field at all.
 */
export type ModelFamily = "whisper" | "parakeet";

/** Exact weight-file sizes per dtype, for entries whose files we know (D-051). */
export interface ExactModelBytes {
  encoder: Partial<Record<WhisperDtype, number>>;
  decoder: Partial<Record<WhisperDtype, number>>;
  /** Tokenizer/vocab/config files. */
  aux: number;
}

export interface WhisperModel {
  /** Hugging Face Hub id, e.g. "onnx-community/whisper-tiny.en". */
  id: string;
  /** Inference path (D-051); absent ⇒ "whisper". */
  family?: ModelFamily;
  /**
   * Real file sizes; when present they replace the parameter-count estimate
   * in `estimateDownloadBytes` / `estimatePeakLoadBytes`.
   */
  exactBytes?: ExactModelBytes;
  /** Human-readable name for the picker. */
  displayName: string;
  /** False for `.en` (English-only) variants. */
  multilingual: boolean;
  /**
   * Whether Whisper's own `task: "translate"` (speech → English text) is
   * offered (D-044). False for every `.en` model (no task token) and for
   * large-v3-turbo, which OpenAI did not train for translation — and for any
   * turbo-based fine-tune not tested for it. The UI hides "Translate to
   * English" where false; text translation of a finished transcript
   * (D-045) works for every model regardless.
   */
  canTranslate: boolean;
  /** Encoder parameter count, in millions (drives size estimates). */
  encoderParamsM: number;
  /** Decoder parameter count, in millions (drives size estimates). */
  decoderParamsM: number;
  /** Per-device dtype defaults (D-004 / D-006). */
  defaults: Record<InferenceDevice, DtypeDefault>;
  /**
   * Model name the server engine passes to faster-whisper (D-013). These are
   * faster-whisper's own aliases, which it resolves to the official
   * CTranslate2 conversions on the Hub (`Systran/faster-whisper-*`;
   * large-v3-turbo → `mobiuslabsgmbh/faster-whisper-large-v3-turbo`), so the
   * pinned library version stays the single authority on the exact repo.
   * Absent for browser-only families (Parakeet, D-051): the server engine
   * answers 400 for them and the picker hides them while it is selected.
   */
  fasterWhisperModel?: string;
}

/** The entry's inference path; absent family (or unknown model) ⇒ "whisper". */
export function modelFamily(model: WhisperModel | undefined): ModelFamily {
  return model?.family ?? "whisper";
}

/** Bytes of one part at one dtype: the exact file size when known, else the estimate. */
function partBytes(model: WhisperModel, part: "encoder" | "decoder", dtype: WhisperDtype): number {
  const exact = model.exactBytes?.[part][dtype];
  if (exact !== undefined) return exact;
  const params = part === "encoder" ? model.encoderParamsM : model.decoderParamsM;
  return params * 1e6 * BYTES_PER_PARAM[dtype];
}

// ---------------------------------------------------------------------------
// Size estimation (approximate, shown before download — D-006)
// ---------------------------------------------------------------------------

/** Approximate stored bytes per parameter for each quantization level. */
export const BYTES_PER_PARAM: Record<WhisperDtype, number> = {
  fp32: 4,
  fp16: 2,
  q8: 1.05, // int8 weights + scales/zero-points overhead
  q4: 0.6, // 4-bit blocks + per-block scales; some tensors stay higher precision
};

/** Tokenizer + configs + preprocessor, roughly constant per model. */
const AUX_FILES_BYTES = 3_000_000;

/** Approximate total download size for a model at a given dtype spec. */
export function estimateDownloadBytes(
  model: WhisperModel,
  dtype: DtypeSpec,
): number {
  const encDtype = typeof dtype === "string" ? dtype : dtype.encoder_model;
  const decDtype = typeof dtype === "string" ? dtype : dtype.decoder_model_merged;
  return Math.round(
    partBytes(model, "encoder", encDtype) +
      partBytes(model, "decoder", decDtype) +
      (model.exactBytes?.aux ?? AUX_FILES_BYTES),
  );
}

/** "487 MB" / "1.6 GB" style formatting for size estimates. */
export function formatModelSize(bytes: number): string {
  if (bytes >= 1e9) return `${(bytes / 1e9).toFixed(1)} GB`;
  return `${Math.round(bytes / 1e6)} MB`;
}

/**
 * Approximate peak transient memory a browser-engine load needs (D-027): for
 * whichever of the encoder/decoder weight files is larger, transformers.js
 * buffers the full download into one JS-heap `Uint8Array`, then
 * onnxruntime-web mallocs a second full-size copy inside WASM linear memory
 * to parse it at session-creation time — a per-file ~2x peak. The
 * sequential-session-load patch (`patches/`) keeps the encoder and decoder
 * from both being resident at once, so the two files' peaks don't add
 * together; only the larger one sets the peak. Approximate — real browser/
 * runtime overhead (page baseline, Cache Storage, other tabs) adds more.
 */
export function estimatePeakLoadBytes(model: WhisperModel, dtype: DtypeSpec): number {
  const encDtype = typeof dtype === "string" ? dtype : dtype.encoder_model;
  const decDtype = typeof dtype === "string" ? dtype : dtype.decoder_model_merged;
  const encoderBytes = partBytes(model, "encoder", encDtype);
  const decoderBytes = partBytes(model, "decoder", decDtype);
  return Math.round(2 * Math.max(encoderBytes, decoderBytes));
}

// ---------------------------------------------------------------------------
// dtype defaults (community-known-good; see file header)
// ---------------------------------------------------------------------------

const WASM_Q8: DtypeDefault = {
  dtype: "q8",
  verified: false,
  note: "q8 is the Transformers.js WASM default; good quality/size trade-off.",
};

const WASM_Q4: DtypeDefault = {
  dtype: "q4",
  verified: false,
  note: "q4 keeps the download manageable for larger models on WASM.",
};

/** fp16 has known accuracy issues for the small Whisper encoders on WebGPU. */
const WEBGPU_FP32: DtypeDefault = {
  dtype: { encoder_model: "fp32", decoder_model_merged: "fp32" },
  verified: false,
  note: "fp32 both parts: fp16 encoders are known-bad for tiny/base on WebGPU.",
};

const WEBGPU_FP16_Q4: DtypeDefault = {
  dtype: { encoder_model: "fp16", decoder_model_merged: "q4" },
  verified: false,
  note: "fp16 encoder + q4 merged decoder — the common working WebGPU setup.",
};

const WEBNN_FP16: DtypeDefault = {
  dtype: "fp16",
  verified: false,
  note: "fp16 is commonly used with WebNN.",
};

/**
 * The medium fp16 conversions are unusable (D-022, verified 2026-07-15): the
 * ONNX files crash onnxruntime's graph optimizer (SimplifiedLayerNormFusion
 * chokes on their InsertedPrecisionFreeCast nodes — reproduced on both
 * ort-node and ort-web 1.26.0-dev), and Whisper medium/large encoder
 * activations are the classic fp16 overflow case besides. In the browser the
 * symptom was a degenerate one-word transcript with no error. q4 is verified
 * correct for medium on the WASM runtime; a backend that cannot run it
 * throws at session creation and falls down the D-017 device chain instead
 * of silently producing garbage.
 */
const ACCEL_Q4: DtypeDefault = {
  dtype: "q4",
  verified: false,
  note: "q4: the medium fp16 conversions fail to load / overflow (D-022).",
};

// ---------------------------------------------------------------------------
// Danish fine-tunes (D-046)
// ---------------------------------------------------------------------------

/**
 * Hugging Face account that hosts our conversions of the Danish Whisper
 * fine-tunes (D-046). Both the ONNX repo (browser engine) and the
 * CTranslate2 int8 repo (server engine) hang off this one value. Only
 * whisper-small-danish is published: the turbo fine-tunes (turbo-danish,
 * edda-v0.2) showed no clear gain over stock turbo on our Danish clip.
 */
export const DANISH_MODELS_HF_OWNER = "lucasw89";

/**
 * WASM default for the Danish conversions (D-046): measured on our own files
 * (da_festuge, onnxruntime-node CPU), q4 matched or beat q8 on WER and was
 * faster for all three — small included, unlike the stock small (q8).
 */
const DANISH_WASM_Q4: DtypeDefault = {
  dtype: "q4",
  verified: false,
  note: "q4: best measured WER and fastest of q8/q4 on our Danish conversions (D-046).",
};

/** Catalog id + matching faster-whisper repo for one published conversion. */
function danishRepos(name: string): Pick<WhisperModel, "id" | "fasterWhisperModel"> {
  return {
    id: `${DANISH_MODELS_HF_OWNER}/${name}-ONNX`,
    fasterWhisperModel: `${DANISH_MODELS_HF_OWNER}/${name}-ct2-int8`,
  };
}

// ---------------------------------------------------------------------------
// The catalog (D-006: curated, not exhaustive)
// ---------------------------------------------------------------------------

/** Parakeet TDT 0.6B v3 ONNX repo and the revision its file sizes were taken from (D-051). */
export const PARAKEET_MODEL_ID = "ysdede/parakeet-tdt-0.6b-v3-onnx";
export const PARAKEET_REVISION = "e5356bf6e24bdb7a269536b0e9e7869bc034b893";

export const MODELS: readonly WhisperModel[] = [
  {
    id: "onnx-community/whisper-tiny.en",
    fasterWhisperModel: "tiny.en",
    displayName: "Whisper Tiny (English)",
    multilingual: false,
    canTranslate: false,
    encoderParamsM: 8,
    decoderParamsM: 31,
    defaults: {
      // Verified in-browser during M1 (WASM, Transformers.js default q8).
      wasm: { ...WASM_Q8, verified: true },
      webgpu: WEBGPU_FP32,
      webnn: WEBNN_FP16,
    },
  },
  {
    id: "onnx-community/whisper-tiny",
    fasterWhisperModel: "tiny",
    displayName: "Whisper Tiny",
    multilingual: true,
    canTranslate: true,
    encoderParamsM: 8,
    decoderParamsM: 31,
    defaults: { wasm: WASM_Q8, webgpu: WEBGPU_FP32, webnn: WEBNN_FP16 },
  },
  {
    id: "onnx-community/whisper-base.en",
    fasterWhisperModel: "base.en",
    displayName: "Whisper Base (English)",
    multilingual: false,
    canTranslate: false,
    encoderParamsM: 20,
    decoderParamsM: 54,
    defaults: { wasm: WASM_Q8, webgpu: WEBGPU_FP32, webnn: WEBNN_FP16 },
  },
  {
    id: "onnx-community/whisper-base",
    fasterWhisperModel: "base",
    displayName: "Whisper Base",
    multilingual: true,
    canTranslate: true,
    encoderParamsM: 20,
    decoderParamsM: 54,
    defaults: { wasm: WASM_Q8, webgpu: WEBGPU_FP32, webnn: WEBNN_FP16 },
  },
  {
    id: "onnx-community/whisper-small.en",
    fasterWhisperModel: "small.en",
    displayName: "Whisper Small (English)",
    multilingual: false,
    canTranslate: false,
    encoderParamsM: 88,
    decoderParamsM: 156,
    defaults: { wasm: WASM_Q8, webgpu: WEBGPU_FP16_Q4, webnn: WEBNN_FP16 },
  },
  {
    id: "onnx-community/whisper-small",
    fasterWhisperModel: "small",
    displayName: "Whisper Small",
    multilingual: true,
    canTranslate: true,
    encoderParamsM: 88,
    decoderParamsM: 156,
    defaults: { wasm: WASM_Q8, webgpu: WEBGPU_FP16_Q4, webnn: WEBNN_FP16 },
  },
  {
    // onnx-community has no plain "whisper-medium.en"/"whisper-medium" repos
    // (the Hub answers 401 "Unauthorized" for missing repos) — the medium
    // conversions live under these names instead. Verified to carry the full
    // q4/q8/fp16 encoder + merged-decoder variant set.
    id: "onnx-community/whisper-medium.en_timestamped",
    fasterWhisperModel: "medium.en",
    displayName: "Whisper Medium (English)",
    multilingual: false,
    canTranslate: false,
    encoderParamsM: 306,
    decoderParamsM: 463,
    defaults: { wasm: WASM_Q4, webgpu: ACCEL_Q4, webnn: ACCEL_Q4 },
  },
  {
    id: "onnx-community/whisper-medium-ONNX",
    fasterWhisperModel: "medium",
    displayName: "Whisper Medium",
    multilingual: true,
    canTranslate: true,
    encoderParamsM: 306,
    decoderParamsM: 463,
    defaults: { wasm: WASM_Q4, webgpu: ACCEL_Q4, webnn: ACCEL_Q4 },
  },
  {
    id: "onnx-community/whisper-large-v3-turbo",
    fasterWhisperModel: "large-v3-turbo",
    displayName: "Whisper Large v3 Turbo",
    multilingual: true,
    // Not trained for translation (OpenAI's turbo model card); its
    // "translate" output is unreliable (D-044).
    canTranslate: false,
    encoderParamsM: 635,
    decoderParamsM: 174,
    defaults: { wasm: WASM_Q4, webgpu: WEBGPU_FP16_Q4, webnn: WEBNN_FP16 },
  },
  // Danish fine-tune (D-046): unchanged Whisper architecture, so the
  // parameter counts are its base model's. Conversions by us, weights by
  // the original authors (credited in each published repo's model card).
  {
    // thorhojhus/whisper-small-danish (MIT), fine-tune of whisper-small.
    ...danishRepos("whisper-small-danish"),
    displayName: "Whisper Small (Danish)",
    multilingual: true,
    // Translation untested on these fine-tunes (spec 001, D-044).
    canTranslate: false,
    encoderParamsM: 88,
    decoderParamsM: 156,
    defaults: { wasm: DANISH_WASM_Q4, webgpu: WEBGPU_FP16_Q4, webnn: WEBNN_FP16 },
  },
  {
    // NVIDIA Parakeet TDT 0.6B v3 (CC-BY-4.0), NeMo→ONNX conversion by
    // ysdede, run through parakeet.js on our onnxruntime-web (D-051).
    // English-only in this app (the model itself knows 25 European
    // languages, but only English is measured here). Desktop-only: its
    // ~1.3 GB load peak is far past every phone budget (D-027–D-029), so
    // the picker hides it on mobile instead of just warning.
    id: PARAKEET_MODEL_ID,
    family: "parakeet",
    displayName: "Parakeet v3 (English)",
    multilingual: false,
    canTranslate: false,
    // Parameter counts only feed fallbacks; the sizes below are exact.
    encoderParamsM: 609,
    decoderParamsM: 17,
    exactBytes: {
      encoder: { q8: 652_183_999, fp16: 1_238_960_452 },
      decoder: { q8: 18_202_004, fp16: 36_266_140 },
      aux: 102_132, // vocab.txt
    },
    defaults: {
      wasm: {
        dtype: "q8",
        verified: false,
        note: "int8 encoder + decoder: the measured WASM path (D-051).",
      },
      // Only used when WebGPU is chosen explicitly — "auto" never picks it
      // for Parakeet (D-051): int8 on WebGPU is reported broken, and fp16
      // is a separate 1.24 GB download. The decoder stays int8 on WASM.
      webgpu: {
        dtype: { encoder_model: "fp16", decoder_model_merged: "q8" },
        verified: false,
        note: "fp16 encoder on WebGPU (needs shader-f16) + int8 decoder on WASM.",
      },
      webnn: {
        dtype: "q8",
        verified: false,
        note: "No WebNN path: runs the WASM int8 setup.",
      },
    },
  },
];

export function getModel(id: string): WhisperModel | undefined {
  return MODELS.find((model) => model.id === id);
}

/**
 * faster-whisper model name for a catalog id (server engine, D-013), or
 * undefined for ids outside the catalog — the transcribe endpoint answers 400
 * rather than forwarding arbitrary strings to the sidecar.
 */
export function getFasterWhisperModel(id: string): string | undefined {
  return getModel(id)?.fasterWhisperModel;
}

/**
 * The task a run on `modelId` should actually use (D-044): a "translate"
 * left over from a model that could translate falls back to "transcribe" on
 * one that cannot (large-v3-turbo, `.en`, unknown ids). App applies it on
 * every model switch and again when building each run's options.
 */
export function taskForModel(modelId: string, task: WhisperTask): WhisperTask {
  return task === "translate" && getModel(modelId)?.canTranslate !== true ? "transcribe" : task;
}

/** Default dtype for a model on a given device. Unknown ids fall back to q8. */
export function resolveDefaultDtype(
  modelId: string,
  device: InferenceDevice,
): DtypeSpec {
  return getModel(modelId)?.defaults[device].dtype ?? "q8";
}

/** Approximate download size using the model's default dtype for `device`. */
export function estimateDefaultDownloadBytes(
  model: WhisperModel,
  device: InferenceDevice,
): number {
  return estimateDownloadBytes(model, model.defaults[device].dtype);
}

// ---------------------------------------------------------------------------
// Defaults & language list
// ---------------------------------------------------------------------------

export const DEFAULT_MODEL_ID = "onnx-community/whisper-tiny.en";

/** UI default: worker probes WebNN → WebGPU → WASM in order (D-004, D-017). */
export const DEFAULT_DEVICE_PREFERENCE: DevicePreference = "auto";

export const DEFAULT_TASK: WhisperTask = "transcribe";

/** Sentinel for "let Whisper detect the language". */
export const AUTO_LANGUAGE = "auto";

/**
 * Curated language list for multilingual models (Whisper language codes).
 * Not exhaustive — Whisper supports ~100; extend on demand.
 */
export const WHISPER_LANGUAGES: readonly { code: string; label: string }[] = [
  { code: "en", label: "English" },
  { code: "da", label: "Danish" },
  { code: "de", label: "German" },
  { code: "es", label: "Spanish" },
  { code: "fr", label: "French" },
  { code: "it", label: "Italian" },
  { code: "nl", label: "Dutch" },
  { code: "no", label: "Norwegian" },
  { code: "sv", label: "Swedish" },
  { code: "fi", label: "Finnish" },
  { code: "is", label: "Icelandic" },
  { code: "pl", label: "Polish" },
  { code: "pt", label: "Portuguese" },
  { code: "ro", label: "Romanian" },
  { code: "cs", label: "Czech" },
  { code: "el", label: "Greek" },
  { code: "hu", label: "Hungarian" },
  { code: "uk", label: "Ukrainian" },
  { code: "ru", label: "Russian" },
  { code: "tr", label: "Turkish" },
  { code: "ar", label: "Arabic" },
  { code: "he", label: "Hebrew" },
  { code: "hi", label: "Hindi" },
  { code: "id", label: "Indonesian" },
  { code: "vi", label: "Vietnamese" },
  { code: "th", label: "Thai" },
  { code: "zh", label: "Chinese" },
  { code: "ja", label: "Japanese" },
  { code: "ko", label: "Korean" },
];