diff --git a/AGENTS.md b/AGENTS.md index b49a62b..1f07fed 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -19,6 +19,8 @@ Build: `bun run build` (rewrites `dist/`) Opt-in real sidecar test (needs `pip install -r server-py/requirements.txt`): `FW_INTEGRATION=1 bun test tests/transcribe.test.ts` Opt-in real translation test (~3 min, downloads 2 × 113 MB Opus-MT): `MT_INTEGRATION=1 bun test tests/translate-real-model.test.ts`; translation speed: `bun scripts/bench-translate.ts` +Opt-in real Parakeet test (~670 MB download on first run; run when touching `src/lib/parakeet.ts` or +`src/worker/parakeetEngine.ts`): `PARAKEET_REAL=1 bun test tests/parakeet-real-model.test.ts` Typecheck: `bun run typecheck` (tsc 7, < 1 s). No linter configured. Browser-only behaviour (offline PWA, mic, WebGPU) has no automated coverage: follow `tests/offline.md`, @@ -39,7 +41,9 @@ step with the commands above, in the same commit. - **Two engines, two libraries (D-013).** Browser inference = Transformers.js, inside `src/worker/inference.worker.ts` only (transcript translation: `src/worker/translate.worker.ts`, D-WS4b) — never on the main thread. Server inference = faster-whisper - via `server-py/sidecar.py` — never Transformers.js on the server. + via `server-py/sidecar.py` — never Transformers.js on the server. The one exception is Parakeet + (`family: "parakeet"`, D-WS7a): parakeet.js in the same worker, on the **same** onnxruntime-web + instance — never call its `fromUrls`/`fromHub`/`initOrt` (own ort copy, CDN wasm path). - **`@huggingface/transformers` is patched and a devDependency on purpose.** It is pinned by `patches/@huggingface%2Ftransformers@4.3.2.patch` (D-027/D-028, D-WS1a) and kept out of runtime deps so the Docker image stays small (D-031). Upgrading or moving it means redoing the patch: ask first. diff --git a/README.md b/README.md index 94ea7fb..b1cf3e9 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,13 @@ users can optionally run inference on the **server** instead, backed by [faster-whisper](https://github.com/SYSTRAN/faster-whisper) (CTranslate2) through a small Python sidecar (`server-py/`, Decisions.md D-013). +On desktop browsers the on-device engine also offers **NVIDIA Parakeet TDT 0.6B v3** +(English only) via [parakeet.js](https://github.com/ysdede/parakeet.js) (MIT) on the +same onnxruntime-web. Model: [nvidia/parakeet-tdt-0.6b-v3](https://huggingface.co/nvidia/parakeet-tdt-0.6b-v3), +© NVIDIA, licensed [CC-BY-4.0](https://creativecommons.org/licenses/by/4.0/); ONNX +conversion [ysdede/parakeet-tdt-0.6b-v3-onnx](https://huggingface.co/ysdede/parakeet-tdt-0.6b-v3-onnx), +downloaded unmodified by the browser on first use (Decisions.md D-WS7a). + See `docs/Architecture.md` for the design, `docs/Decisions.md` for the decision log, and `PRODUCT.md` for the product spec. diff --git a/bun.lock b/bun.lock index 5871853..270dade 100644 --- a/bun.lock +++ b/bun.lock @@ -19,6 +19,7 @@ "@types/react": "^19", "@types/react-dom": "^19", "ffmpeg-static": "^5.3.0", + "parakeet.js": "1.4.4", "typescript": "^7.0.2", }, }, @@ -250,6 +251,8 @@ "onnxruntime-web": ["onnxruntime-web@1.31.0-dev.20260914-8d85527a0", "", { "dependencies": { "flatbuffers": "^25.1.24", "guid-typescript": "^1.0.9", "long": "^5.2.3", "onnxruntime-common": "1.31.0-dev.20260911-2a43ec07e", "platform": "^1.3.6", "protobufjs": "^7.2.4" } }, "sha512-Iy7rtadoBgxS/LLvDr3QW38DB1PNXRnr0GJMcL0TAt7c9qjgVQl83UlGCVyeAnK2InpmW8Uc3PL8XIuqtDeF6g=="], + "parakeet.js": ["parakeet.js@1.4.4", "", { "dependencies": { "onnxruntime-web": "1.24.1" } }, "sha512-+tYIDgp799bvsDye3y2lFyspYHgTOJ5YMZi0mWYABXyJjp/3NZmNaMTnS9jhf4qlIQJoBk0Nmj+UwsI3zk+Aqg=="], + "parse-cache-control": ["parse-cache-control@1.0.1", "", {}, "sha512-60zvsJReQPX5/QP0Kzfd/VrpjScIQ7SHBW6bFCYfEP+fp0Eppr1SHhIO5nd1PjZtvclzSzES9D/p5nFJurwfWg=="], "passport": ["passport@0.7.0", "", { "dependencies": { "passport-strategy": "1.x.x", "pause": "0.0.1", "utils-merge": "^1.0.1" } }, "sha512-cPLl+qZpSc+ireUvt+IzqbED1cHHkDoVYMo30jbJIdOOjQ1MQYZBPiNvmi8UM6lJuOpTPXJGZQk0DtC4y61MYQ=="], @@ -301,5 +304,9 @@ "http-response-object/@types/node": ["@types/node@10.17.60", "", {}, "sha512-F0KIgDJfy2nA3zMLmWGKxcH2ZVEtCZXHHdOQs2gSaQ27+lNeEfGxzkIw90aXswATX7AZ33tahPbzy6KAfUreVw=="], "onnxruntime-web/onnxruntime-common": ["onnxruntime-common@1.31.0-dev.20260911-2a43ec07e", "", {}, "sha512-gBuF6U32YErKIAt+yD7DeGBpjRxhJ1uJwako6ygjokSZULDU6Pd+KWComX8Tumk1hYVHRxjPj7mdMnD0ryAPOw=="], + + "parakeet.js/onnxruntime-web": ["onnxruntime-web@1.24.1", "", { "dependencies": { "flatbuffers": "^25.1.24", "guid-typescript": "^1.0.9", "long": "^5.2.3", "onnxruntime-common": "1.24.1", "platform": "^1.3.6", "protobufjs": "^7.2.4" } }, "sha512-i2u395dv+ZEQBdH+aORvlu19Bzvlg5AXJ7wjxnL350hknOP9z0UeP3pVfjkpMEWMPy2T6nCQxetKTmNia6wSzg=="], + + "parakeet.js/onnxruntime-web/onnxruntime-common": ["onnxruntime-common@1.24.1", "", {}, "sha512-UnV15u4p4XxoIV+jFP4hXPsW93s3QrwLSpi20HUDYHoTfI4z4sjzex3L4XDOxGGZJ/M/catrwAG2go958UQq0w=="], } } diff --git a/docs/Architecture.md b/docs/Architecture.md index 8d1e582..b5e768d 100644 --- a/docs/Architecture.md +++ b/docs/Architecture.md @@ -72,13 +72,17 @@ src/ audio.ts # decode + resample to 16 kHz mono Float32Array (Web Audio API) live.ts # live-mode segment loop: RMS gate, LiveSegmenter, LiveTranscriber (pure, D-016) mic.ts # microphone capture: AudioWorklet chunker (Blob URL) + fallback (D-016) - models.ts # model catalog: ids, sizes, dtype defaults per device, languages + models.ts # model catalog: ids, family (whisper|parakeet), sizes, dtype defaults per device, languages + parakeet.ts # pure Parakeet logic: ≤30 s window plan, words→chunks, vocab tokenizer, pinned URLs modelCache.ts # "downloaded?" per model via Cache Storage (D-005; reused by M3 offline indicator) engine.ts # engine abstraction: BrowserEngine | ServerEngine translation.ts # transcript translation planning: Opus-MT catalog, units, sentence pieces (pure, D-WS4b) hint.ts # "Names and terms" hint: normalize, localStorage, prompt truncation rule (D-WS2a) worker/ - inference.worker.ts # Transformers.js pipeline, WebNN/WebGPU/WASM selection + inference.worker.ts # Transformers.js pipeline, WebNN/WebGPU/WASM selection; dispatches on catalog `family` + parakeetEngine.ts # Parakeet TDT path: parakeet.js ParakeetModel on the shared onnxruntime-web (D-WS7a) + modelFiles.ts # Cache-API weight download for non-Transformers.js paths (same cache, same progress msgs) + ortShim.ts # re-exports onnxruntime-web/webgpu; parakeet.js's own ort import is aliased here protocol.ts # typed postMessage messages (shared with UI) prompt.ts # decoder prompt (<|startofprev|> hint) via a generate() wrapper (D-WS2a) translate.worker.ts # transcript translation (Opus-MT), its own worker + ort instance (D-WS4b) @@ -279,11 +283,40 @@ environment). See `WorkBreakdown.md` for the itemized checklist.) - Anonymous users are never blocked from browser inference — auth only unlocks the server engine. +### 4.1c Parakeet (English-only, desktop-only browser model — D-WS7a) +Catalog entries carry `family?: "whisper" | "parakeet"` (absent ⇒ whisper). For +`family: "parakeet"` the worker skips the Transformers.js pipeline: +1. **Load:** `modelFiles.ts` fetches `vocab.txt`, the decoder_joint and the + encoder from the Hub (revision-pinned URLs, `lib/parakeet.ts`) into the same + `transformers-cache` Cache Storage cache Transformers.js uses, posting the same + `load-progress` messages — so badges, the storage panel, delete and + ModelLoadProgress work unchanged. A file is cached only once complete. + Devices: `parakeetDeviceAttempts` — WASM int8 unless WebGPU is chosen + explicitly (fp16 encoder on WebGPU + int8 decoder on WASM, needs + `shader-f16`); the D-025 restart rule applies. +2. **Runtime:** `parakeetEngine.ts` builds parakeet.js's `ParakeetModel` from our + own sessions on `onnxruntime-web/webgpu` — the module Transformers.js imports, + so there is one runtime and one `/ort/` wasm. parakeet.js's `initOrt` (own ort + copy, CDN wasm path) is never called; `scripts/ortAlias.ts` points its import + at `worker/ortShim.ts` so the bundle carries one ort (`tests/workerBundle.test.ts`). +3. **Transcribe:** audio is planned into ≤ 30 s windows cut at the quietest 100 ms + near each limit (`subarray` views, no copies); each window's word timestamps + become sentence-sized `TranscriptChunk`s with absolute times. File passes post + one `partial` + `transcribe-progress` per window; live passes resolve via + `result` like Whisper's, so `LiveTranscriber`, TranscriptView and export are + untouched. `language`/`task` are ignored (English only). +4. **UI:** the picker lists it under English-only with a one-line note, hides it on + phones/tablets (`modelsForDevice`), and the server-engine list never includes + it (no faster-whisper name). With it selected, a note replaces the language/task + controls. + ## 5. Invariants - **Privacy:** browser-engine audio and transcripts never touch the network. - **Anonymous-first:** every feature except the server engine works logged out. - **Offline-first:** app shell + previously-downloaded models must work offline. -- **One browser inference codebase:** Transformers.js in the browser; the server - engine runs faster-whisper (D-013). The model catalog stays shared across both. +- **One browser inference runtime:** onnxruntime-web, driven by Transformers.js for + Whisper and by parakeet.js for Parakeet (D-WS7a) — the same ort instance, never a + second runtime; the server engine runs faster-whisper (D-013). The model catalog + stays shared across both engines (browser-only entries have no faster-whisper name). - **Worker isolation:** inference never runs on the main thread. diff --git a/docs/Decisions.md b/docs/Decisions.md index ac96da0..dcde821 100644 --- a/docs/Decisions.md +++ b/docs/Decisions.md @@ -1703,3 +1703,93 @@ path (`choosePipeline`). Live mode keeps its energy gate (D-016). **Rejected:** bundling the model in `dist/` (2 MB on every install for an off-by-default option); running it on the main thread (D-013). **Consequence:** the first "Skip silence" run online downloads 2.2 MB. + +## D-WS7a — Parakeet TDT 0.6B v3 as an English-only, desktop-only browser model via parakeet.js on the shared onnxruntime-web +**Status:** Accepted (2026-10-29, spec 001 Goal 7) +**Runtime decision (evidence).** `@huggingface/transformers` 4.2.0 (ours) and +4.3.2 (npm `latest`) were inspected from their published tarballs: both ship +only `ParakeetForCTC` (`models/parakeet/modeling_parakeet.js`, identical in +both), routed by the ASR pipeline as `model_type: 'parakeet_ctc'` through the +wav2vec2 CTC path (argmax, no timestamps, no chunking). Neither contains a +TDT/transducer decoder. Parakeet TDT therefore runs on **parakeet.js 1.4.4** +(MIT; pre-approved by the spec for exactly this case), a devDependency +bundled into the worker only (like Transformers.js, D-031). +**How it runs on one runtime.** parakeet.js pins its own onnxruntime-web +1.24.1 and its `fromUrls`/`fromHub` → `initOrt` path imports it, sets a +jsdelivr `wasmPaths` (blocked by COEP, D-023; not offline) and fetches +weights outside our cache. We never call those: `src/worker/parakeetEngine.ts` +constructs `ParakeetModel` directly from sessions created on +`onnxruntime-web/webgpu` — the module Transformers.js imports — with our own +vocab tokenizer and parakeet.js's `JsPreprocessor` (`parakeet.js/mel`). A Bun +build plugin (`scripts/ortAlias.ts`, used by the build and the dev server) +resolves parakeet.js's bare `onnxruntime-web` import to `src/worker/ortShim.ts`, +so the bundle carries one ort (`tests/workerBundle.test.ts` asserts it, and +that without the plugin the 1.24.1 copy *would* be bundled). We rely on +`onnxruntime-web` being hoisted from Transformers.js rather than pinning it +ourselves, so the two can never drift apart (workstream 1's upgrade moves both). +**Model.** `ysdede/parakeet-tdt-0.6b-v3-onnx` (NVIDIA's CC-BY-4.0 weights, +NeMo→ONNX), URLs pinned to revision `e5356bf6…` — the one measured and +whose exact file sizes the catalog shows (int8 encoder 652.2 MB + int8 +decoder_joint 18.2 MB + vocab 0.1 MB = "~670 MB"; WebGPU fp16 encoder +1.24 GB). Attribution in README. The model knows 25 European languages; the +app offers it **English-only** (spec), so there is no language/task UI. +**Engine path.** Catalog entries gain `family?: "whisper" | "parakeet"` +(absent ⇒ whisper) and optional `exactBytes`; `fasterWhisperModel` becomes +optional (browser-only entries have none, so the server engine 400s them and +the picker's server list excludes them). The worker dispatches on the family. +Weights go through `src/worker/modelFiles.ts` into the `transformers-cache` +Cache Storage cache under their Hub URLs, posting the same `load-progress` +messages — badges, storage panel, delete, ModelLoadProgress unchanged +(D-005, D-033); a file is `put` only once complete. Download stays +user-initiated (D-020). Audio is cut into ≤ 30 s windows (spec: ≤ ~35 s per +encoder call) at the quietest 100 ms in the last 6 s before each limit, so no +overlap merging is needed; we don't use parakeet.js's `transcribeLongAudio` +(20–180 s windows, 90 s default). Word timestamps become sentence-sized +`TranscriptChunk`s (break after `.?!`, on ≥ 1 s pauses, or at 12 s). File +passes post one `partial` + absolute `transcribe-progress` per window; live +passes return `result` like Whisper, so live mode works unchanged. +**Devices.** WASM int8 is the baseline and the only path "auto" picks. An +explicit WebGPU choice uses the fp16 encoder on WebGPU + int8 decoder on WASM +(parakeet.js's "hybrid" split), only with `shader-f16`, falling back to WASM +through the D-025 restart. Not auto: WebGPU int8 is reported broken, the fp16 +encoder is a separate 1.24 GB download, and nothing WebGPU can be measured +here. WebNN has no path (runs WASM). WASM sessions use graph optimization +"all" (D-014's "basic" cap is about Whisper's 8-bit tied embeddings; measured +fine here). +**Mobile.** Load peak ≈ 2 × 652 MB = 1.3 GB (D-027 arithmetic) — above every +budget D-029 treats as safe (deviceMemory 8 ⇒ 800 MB). Instead of a confirm, +`modelsForDevice` hides it on phones/tablets (incl. iPadOS) and +`confirmRiskyMobileLoad` refuses it. +**Measurements (CPU only; this machine, 16 cores; synthetic nothing).** +`PARAKEET_REAL=1 bun test tests/parakeet-real-model.test.ts` (onnxruntime-web +WASM under Bun, 8 threads, the worker's own engine code): + +| clip | audio | file WER | file wall | live WER | live wall | +|---|---|---|---|---|---| +| en_aiproblem | 14.4 s | 0.030 | 0.9 s | 0.030 | 2.0 s | +| en_alien | 888.4 s | 0.096 | 36.4 s | 0.115 | 282 s | + +`scripts/bench-parakeet.ts` (whole-file, same normalization, 12 threads, +model load excluded, one run each under the shared lock): + +| engine (runtime) | en_aiproblem WER / wall | en_alien WER / wall | +|---|---|---| +| Parakeet int8 (onnxruntime-node CPU) | 0.030 / 0.3 s | 0.100 / 15.3 s | +| Parakeet int8 (onnxruntime-web WASM) | 0.030 / 0.9 s | 0.096 / 34.4 s | +| Large v3 Turbo q4 (onnxruntime-node, worker options) | 0.000 / 2.9 s | 0.255 / 183.7 s | + +(D-039 recorded Turbo en_alien whole-file 0.261 / 420 s on a busier run.) On +the long, noisy clip Parakeet more than halves Turbo's WER and is ~12× faster +on the same runtime; on the short clean clip Turbo is perfect and Parakeet +misses one word ("agentic" → "agency"). No WASM number for Turbo: the +Transformers.js node build always uses onnxruntime-node. +**Consequence.** For English, Parakeet is the accurate *and* fast choice on +desktop; Whisper stays the default and the only multilingual/translate +option. Live mode is supported (faster than real time on WASM here). +**Needs the owner's browser check:** real-Chrome WASM load + transcription +speed/memory (headless Chrome run here, see WorkBreakdown); explicit WebGPU +(fp16 encoder) on the owner's GPU — success, speed vs WASM, and whether the +fallback restart behaves; Firefox/Safari WASM. +**Revisit:** if Transformers.js gains a TDT decoder, move to it and drop +parakeet.js + the alias plugin; if WebGPU measures well on real hardware, +consider letting "auto" pick it. diff --git a/docs/WorkBreakdown.md b/docs/WorkBreakdown.md index 871f9d2..1b9cffd 100644 --- a/docs/WorkBreakdown.md +++ b/docs/WorkBreakdown.md @@ -372,6 +372,22 @@ cached. --- +## Spec 001 Goal 7 — Parakeet TDT 0.6B v3 (English, desktop, D-WS7a) +- [x] Runtime decision: Transformers.js 4.2.0/4.3.2 have CTC-only Parakeet → parakeet.js on the + shared onnxruntime-web (one runtime in the bundle, `tests/workerBundle.test.ts`). +- [x] Catalog `family` + exact sizes; worker engine path (≤ 30 s windows, word → chunk + timestamps, progress/partials); Cache-API downloader (`modelFiles.ts`). +- [x] File and live mode; WASM int8 baseline; explicit WebGPU = fp16 encoder (opt-in only). +- [x] Hidden on phones/tablets; English-only note instead of language/task controls. +- [x] Real-model WER/wall-clock vs Large v3 Turbo (CPU), opt-in `PARAKEET_REAL=1` test. +- [x] Headless Chrome (no GPU, fresh profile, prod build, COOP/COEP on): picking Parakeet and + dropping `en_aiproblem.m4a` downloaded the weights, cached them ("available offline"), + ran on WASM (12 threads) and showed the transcript, 2 timestamp segments and the + .txt/.srt/.txt+ts/Copy exports — 15 s end to end including the ~670 MB download. Picker + screenshots: desktop light/dark show the row + English-only note; a phone UA hides it. +- [ ] Owner's browser check: real Chrome WASM load/transcribe (speed, memory), explicit + WebGPU on a real GPU (works? faster than WASM? fallback restart), Firefox/Safari WASM. + ## Deliberately out of scope (for now) Arbitrary Hub model ids, OAuth providers, multi-user transcript storage, whisper.cpp server backend, i18n of the UI. @@ -433,6 +449,11 @@ Found while building the long-form pipeline (spec 001 Goal 3, D-WS3a), 2026-10-2 `en_alien`'s reference; `tests/meetingFixture.ts` returns audio + reference only. Add the list when the vocabulary hint (workstream 2) lands. +- [ ] Parakeet v3 also transcribes 24 other European languages, Danish included, but the app + offers it English-only (spec 001) — unmeasured on `da_festuge`. Candidate for a later spec. +- [ ] parakeet.js still logs `[Parakeet.js]` lines to the worker console on some paths; harmless, + not silenced (no logger hook in the library). + ## Recurring - [ ] Maintain `AGENTS.md`: - **Add a rule only in response to a real failure.** When an agent repeatedly diff --git a/package.json b/package.json index 70d464e..c0b66a2 100644 --- a/package.json +++ b/package.json @@ -25,6 +25,7 @@ "@types/react": "^19", "@types/react-dom": "^19", "ffmpeg-static": "^5.3.0", + "parakeet.js": "1.4.4", "typescript": "^7.0.2" }, "trustedDependencies": [ diff --git a/scripts/bench-parakeet.ts b/scripts/bench-parakeet.ts new file mode 100644 index 0000000..7e645fd --- /dev/null +++ b/scripts/bench-parakeet.ts @@ -0,0 +1,102 @@ +/** + * Parakeet vs Whisper Large v3 Turbo benchmark (D-WS7a) — CPU only. + * + * flock /tmp/whisper-heavy.lock bun scripts/bench-parakeet.ts [parakeet-node|parakeet-web|turbo ...] + * + * Whole-file WER (same normalization as tests/live-real-model.test.ts) and + * wall-clock per English clip in tests/audio_files. Parakeet runs the + * worker's own `createParakeetEngine` on onnxruntime-node (CPU EP) and on + * onnxruntime-web (WASM, the browser engine's baseline); Turbo runs the + * Transformers.js pipeline with the worker's file-flow options on + * onnxruntime-node at its WASM-default q4 dtype. Weights come from / go to + * node_modules/@huggingface/transformers/.cache (Parakeet must already be + * there — `PARAKEET_REAL=1 bun test tests/parakeet-real-model.test.ts` + * downloads it). Wall-clock excludes model load. + */ + +import { join } from "node:path"; +import ffmpegPath from "ffmpeg-static"; +import { pipeline } from "@huggingface/transformers"; +import { PARAKEET_MODEL_ID, resolveDefaultDtype } from "../src/lib/models"; +import { PARAKEET_VOCAB_FILE, parakeetFileName } from "../src/lib/parakeet"; +import { createParakeetEngine, parakeetSessionOptions, type OrtLike } from "../src/worker/parakeetEngine"; +import { wordErrorRate } from "../tests/dsp"; + +const ROOT = join(import.meta.dir, ".."); +const AUDIO = join(ROOT, "tests/audio_files"); +const PARAKEET_DIR = join(ROOT, "node_modules/@huggingface/transformers/.cache", PARAKEET_MODEL_ID); +const CLIPS = ["en_aiproblem.m4a", "en_alien.opus"]; +const THREADS = 12; + +const normalize = (text: string) => + text.toLowerCase().replace(/['’]/g, "").replace(/[.,!?;:"“”‘‚«»()\-–—…]/g, " ").replace(/\s+/g, " ").trim(); + +async function decode(file: string): Promise { + const proc = Bun.spawn( + [ffmpegPath as unknown as string, "-hide_banner", "-loglevel", "error", "-i", join(AUDIO, file), + "-ac", "1", "-ar", "16000", "-f", "f32le", "-"], + { stdout: "pipe" }, + ); + return new Float32Array(await new Response(proc.stdout).arrayBuffer()); +} + +type Transcribe = (audio: Float32Array) => Promise; + +async function parakeet(runtime: "node" | "web"): Promise { + let ort: OrtLike; + let options = parakeetSessionOptions("wasm"); + if (runtime === "web") { + const web = await import("onnxruntime-web"); + web.env.wasm.numThreads = THREADS; + ort = web as unknown as OrtLike; + } else { + ort = (await import("onnxruntime-node")) as unknown as OrtLike; + const cpu = { executionProviders: ["cpu"], intraOpNumThreads: THREADS, graphOptimizationLevel: "all" }; + options = { encoder: cpu, decoder: cpu }; + } + const read = async (file: string) => new Uint8Array(await Bun.file(join(PARAKEET_DIR, file)).arrayBuffer()); + const engine = await createParakeetEngine( + ort, + { + encoder: await read(parakeetFileName("encoder", "q8")), + decoder: await read(parakeetFileName("decoder", "q8")), + vocab: await Bun.file(join(PARAKEET_DIR, PARAKEET_VOCAB_FILE)).text(), + }, + options, + ); + return async (audio) => (await engine.transcribe(audio)).text; +} + +async function turbo(): Promise { + const id = "onnx-community/whisper-large-v3-turbo"; + const asr = await pipeline("automatic-speech-recognition", id, { dtype: resolveDefaultDtype(id, "wasm") }); + return async (audio) => { + const output = await asr(audio, { + chunk_length_s: 30, stride_length_s: 5, return_timestamps: true, max_new_tokens: 440, + language: "en", task: "transcribe", + }); + return ((Array.isArray(output) ? output[0] : output) as { text?: string }).text ?? ""; + }; +} + +const engines: Record Promise> = { + "parakeet-node": () => parakeet("node"), + "parakeet-web": () => parakeet("web"), + turbo, +}; + +for (const name of process.argv.slice(2).length ? process.argv.slice(2) : Object.keys(engines)) { + const transcribe = await engines[name]!(); + for (const clip of CLIPS) { + const audio = await decode(clip); + const reference = await Bun.file(join(AUDIO, clip.replace(/\.[^.]+$/, ".txt"))).text(); + const start = performance.now(); + const text = await transcribe(audio); + const seconds = (performance.now() - start) / 1000; + const wer = wordErrorRate(normalize(reference), normalize(text)); + console.log( + `${name}\t${clip}\taudio ${(audio.length / 16000).toFixed(1)} s\twall ${seconds.toFixed(1)} s\tWER ${wer.toFixed(3)}`, + ); + } +} +process.exit(0); diff --git a/scripts/build.ts b/scripts/build.ts index 0bd50ef..2855187 100644 --- a/scripts/build.ts +++ b/scripts/build.ts @@ -18,6 +18,7 @@ import { join, relative } from "node:path"; import { injectPrecacheManifest, precachePathsFromOutputs } from "../src/pwa/precache.ts"; import { ensureFonts, FONTS_DIR } from "./fonts.ts"; import { ensureIcons, ICON_SIZES, ICONS_DIR } from "./generate-icons.ts"; +import { ortAliasPlugin } from "./ortAlias.ts"; const ROOT = join(import.meta.dir, ".."); const DIST = join(ROOT, "dist"); @@ -51,6 +52,8 @@ const workerBuild = await Bun.build({ target: "browser", format: "esm", minify: true, + // One onnxruntime in the bundle (D-WS7a). + plugins: [ortAliasPlugin], }); const workerScript = await workerBuild.outputs[0]!.text(); await Bun.write(join(DIST, "inference.worker.js"), workerScript); diff --git a/scripts/ortAlias.ts b/scripts/ortAlias.ts new file mode 100644 index 0000000..6dbdcfb --- /dev/null +++ b/scripts/ortAlias.ts @@ -0,0 +1,23 @@ +/** + * Bun build plugin: parakeet.js's bare `onnxruntime-web` import resolves to + * `src/worker/ortShim.ts` (= Transformers.js's `onnxruntime-web/webgpu`), so + * the inference worker bundles exactly one onnxruntime (D-WS7a). Used by the + * production build (scripts/build.ts) and the dev server's worker bundle + * (src/index.ts). Without it the bundle carries parakeet.js's own 1.24.1 + * copy as well — never executed, but ~0.5 MB of dead weight and a CDN + * wasmPaths default. + */ + +import type { BunPlugin } from "bun"; +import { join } from "node:path"; + +const SHIM = join(import.meta.dir, "../src/worker/ortShim.ts"); + +export const ortAliasPlugin: BunPlugin = { + name: "parakeet-ort-alias", + setup(build) { + build.onResolve({ filter: /^onnxruntime-web$/ }, (args) => + /[\\/]parakeet\.js[\\/]/.test(args.importer) ? { path: SHIM } : undefined, + ); + }, +}; diff --git a/src/App.tsx b/src/App.tsx index b3c4781..b0fadc9 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -42,10 +42,16 @@ import { estimatePeakLoadBytes, formatModelSize, getModel, + modelFamily, resolveDefaultDtype, taskForModel, } from "./lib/models"; -import { shouldWarnBeforeLoad, type MobileMemoryHint } from "./lib/mobileMemoryGuard"; +import { + isModelAllowedOnDevice, + modelsForDevice, + shouldWarnBeforeLoad, + type MobileMemoryHint, +} from "./lib/mobileMemoryGuard"; import { applyLoadProgress, initialLoadProgress, @@ -77,6 +83,16 @@ const MODEL_IDS = MODELS.map((model) => model.id); // Downloaded badges cover the Opus-MT models too (transcript translation, D-WS4b). const CACHED_MODEL_IDS = [...MODEL_IDS, ...TRANSLATION_MODEL_IDS]; +/** Device signals for the mobile memory guard (D-027) and desktop-only models (D-WS7a). */ +function memoryHint(): MobileMemoryHint { + const nav = typeof navigator !== "undefined" ? navigator : undefined; + return { + userAgent: nav?.userAgent, + maxTouchPoints: nav?.maxTouchPoints, + deviceMemoryGB: (nav as (Navigator & { deviceMemory?: number }) | undefined)?.deviceMemory, + }; +} + export function App() { const workerRef = useRef(null); @@ -277,15 +293,7 @@ export function App() { const targetModel = getModel(forModelId); if (!targetModel) return false; const peakBytes = estimatePeakLoadBytes(targetModel, resolveDefaultDtype(forModelId, device)); - const hint: MobileMemoryHint = { - userAgent: typeof navigator !== "undefined" ? navigator.userAgent : undefined, - maxTouchPoints: typeof navigator !== "undefined" ? navigator.maxTouchPoints : undefined, - deviceMemoryGB: - typeof navigator !== "undefined" - ? (navigator as Navigator & { deviceMemory?: number }).deviceMemory - : undefined, - }; - return shouldWarnBeforeLoad(hint, peakBytes); + return shouldWarnBeforeLoad(memoryHint(), peakBytes); } /** @@ -298,6 +306,9 @@ export function App() { function confirmRiskyMobileLoad(forModelId: string, preference: DevicePreference): boolean { const targetModel = getModel(forModelId); if (!targetModel) return true; + // Desktop-only models are never offered on phones (D-WS7a); this guards + // any path that could still name one (e.g. a stale selection). + if (!isModelAllowedOnDevice(targetModel, memoryHint())) return false; const device = guessDevice(preference); const key = overrideKey(forModelId, device); if (mobileWarningAcknowledged.current.has(key)) return true; @@ -1016,9 +1027,15 @@ export function App() { const serverSelected = engineChoice === "server" && auth.user !== null && serverStatus?.available === true; - const availableModels = serverSelected - ? MODELS.filter((m) => m.fasterWhisperModel === "large-v3-turbo" || m.fasterWhisperModel === "medium.en") - : MODELS; + // Phones never see desktop-only models (Parakeet, D-WS7a); the server + // engine only offers what the sidecar runs. + const availableModels = modelsForDevice( + serverSelected + ? MODELS.filter((m) => m.fasterWhisperModel === "large-v3-turbo" || m.fasterWhisperModel === "medium.en") + : MODELS, + memoryHint(), + ); + const parakeetSelected = modelFamily(model) === "parakeet"; useEffect(() => { if (serverSelected) { @@ -1070,6 +1087,16 @@ export function App() { serverEngine={serverSelected} /> + {/* Parakeet has no language or task choice (D-WS7a) — say so where + those controls would be, rather than leaving their absence unexplained. */} + {parakeetSelected && ( +

+ Parakeet understands English only, so there is no language or translation + choice, and it doesn't use a names-and-terms list. For other languages, or + to hint names, pick a Whisper model. +

+ )} + {model?.multilingual && ( )} - { - setHint(value); - saveHint(value); - }} - disabled={busy || liveActive} - /> + {/* Hint and "Skip silence" are Whisper decode options; the Parakeet + path (D-WS7a) has neither, so they are hidden rather than ignored. */} + {!parakeetSelected && ( + <> + { + setHint(value); + saveHint(value); + }} + disabled={busy || liveActive} + /> - { - setLongForm(enabled); - saveLongFormChoice(enabled); - }} - disabled={busy || liveActive} - /> + { + setLongForm(enabled); + saveLongFormChoice(enabled); + }} + disabled={busy || liveActive} + /> + + )} {/* Engine toggle (M5, D-010): signed-in users choose where uploaded files are transcribed. Live mode always stays on-device (D-016). */} @@ -1126,7 +1159,10 @@ export function App() { models — every pass pads to a full 30 s Whisper window (D-024), so tentative text lags and the stop drain grows. Say so instead of looking frozen. */} - {liveActive && activeDevice === "wasm" && (model?.encoderParamsM ?? 0) >= 300 && ( + {liveActive && + activeDevice === "wasm" && + !parakeetSelected && + (model?.encoderParamsM ?? 0) >= 300 && (

{model?.displayName} runs slower than real time on the WASM backend, so live text will lag behind speech. A smaller model (Tiny, Base or Small) diff --git a/src/components/ModelPicker.tsx b/src/components/ModelPicker.tsx index 61834b1..d8221f0 100644 --- a/src/components/ModelPicker.tsx +++ b/src/components/ModelPicker.tsx @@ -158,7 +158,16 @@ export function ModelPicker({ disabled={disabled} onChange={() => onSelect(model.id)} /> - {model.displayName} + + {model.displayName} + {model.family === "parakeet" && ( + // Not Whisper (D-WS7a): say what it is good for and + // what it can't do, in plain words. + + Not Whisper · English only · fast on long recordings + + )} + {isDownloaded && ( { // Never minify: sw.ts relies on its placeholder string literals surviving // bundling verbatim (src/pwa/precache.ts), and dev builds favor debuggability. minify: false, + // One onnxruntime in the worker bundle (D-WS7a); harmless for sw.ts. + plugins: [ortAliasPlugin], }); const artifact = build.outputs[0]; if (!artifact) throw new Error(`Failed to bundle src/${relativePath}`); diff --git a/src/lib/mobileMemoryGuard.ts b/src/lib/mobileMemoryGuard.ts index 10095c2..ae8429f 100644 --- a/src/lib/mobileMemoryGuard.ts +++ b/src/lib/mobileMemoryGuard.ts @@ -19,6 +19,8 @@ * advance, so a benchmark probe would cost more than it saves. */ +import type { WhisperModel } from "./models"; + export interface MobileMemoryHint { /** `navigator.userAgent`, or undefined outside a browser (tests, SSR). */ userAgent: string | undefined; @@ -86,3 +88,26 @@ export function shouldWarnBeforeLoad(hint: MobileMemoryHint, peakBytes: number): } return peakBytes > MOBILE_FALLBACK_RISK_THRESHOLD_BYTES; } + +/** + * Desktop-only models (D-WS7a): Parakeet's ~1.3 GB load peak (2 × the + * 652 MB int8 encoder, D-027's arithmetic) exceeds every phone budget the + * D-029 table considers safe, even at deviceMemory 8 (800 MB). Unlike the + * Whisper models — where a confirm lets the user judge their device — it is + * not offered at all on phones and tablets: no amount of RAM a phone reports + * makes it safe, so a "Continue anyway?" would only invite the crash. + */ +export function isModelAllowedOnDevice( + model: Pick, + hint: Pick, +): boolean { + return model.family !== "parakeet" || !isLikelyMobile(hint); +} + +/** The catalog as offered on this device (order preserved). */ +export function modelsForDevice>( + models: readonly T[], + hint: Pick, +): T[] { + return models.filter((model) => isModelAllowedOnDevice(model, hint)); +} diff --git a/src/lib/models.ts b/src/lib/models.ts index ad5702e..8f34da1 100644 --- a/src/lib/models.ts +++ b/src/lib/models.ts @@ -35,9 +35,32 @@ export interface DtypeDefault { note: string; } +/** + * Which inference path a catalog entry runs on (D-WS7a). 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-WS7a). */ +export interface ExactModelBytes { + encoder: Partial>; + decoder: Partial>; + /** 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-WS7a); 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. */ @@ -63,8 +86,23 @@ export interface WhisperModel { * 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-WS7a): the server engine + * answers 400 for them and the picker hides them while it is selected. */ - fasterWhisperModel: string; + 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]; } // --------------------------------------------------------------------------- @@ -90,9 +128,9 @@ export function estimateDownloadBytes( const encDtype = typeof dtype === "string" ? dtype : dtype.encoder_model; const decDtype = typeof dtype === "string" ? dtype : dtype.decoder_model_merged; return Math.round( - model.encoderParamsM * 1e6 * BYTES_PER_PARAM[encDtype] + - model.decoderParamsM * 1e6 * BYTES_PER_PARAM[decDtype] + - AUX_FILES_BYTES, + partBytes(model, "encoder", encDtype) + + partBytes(model, "decoder", decDtype) + + (model.exactBytes?.aux ?? AUX_FILES_BYTES), ); } @@ -116,8 +154,8 @@ export function formatModelSize(bytes: number): string { 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 = model.encoderParamsM * 1e6 * BYTES_PER_PARAM[encDtype]; - const decoderBytes = model.decoderParamsM * 1e6 * BYTES_PER_PARAM[decDtype]; + const encoderBytes = partBytes(model, "encoder", encDtype); + const decoderBytes = partBytes(model, "decoder", decDtype); return Math.round(2 * Math.max(encoderBytes, decoderBytes)); } @@ -209,6 +247,10 @@ function danishRepos(name: string): Pick number, + options: PlanChunksOptions = {}, +): Array<[number, number]> { + const rate = options.sampleRate ?? PARAKEET_SAMPLE_RATE; + const maxSamples = Math.round((options.maxSeconds ?? 30) * rate); + const searchSamples = Math.round((options.searchSeconds ?? 6) * rate); + const frame = Math.max(1, Math.round((options.frameSeconds ?? 0.1) * rate)); + + const windows: Array<[number, number]> = []; + let start = 0; + while (start < length) { + if (length - start <= maxSamples) { + windows.push([start, length]); + break; + } + const limit = start + maxSamples; + let cut = limit; + let best = Number.POSITIVE_INFINITY; + for (let f = Math.max(start, limit - searchSamples); f + frame <= limit; f += frame) { + const energy = energyAt(f, f + frame); + if (energy <= best) { + best = energy; + cut = f + Math.floor(frame / 2); + } + } + windows.push([start, cut]); + start = cut; + } + return windows; +} + +/** `planChunks` over real PCM, scoring frames by mean squared amplitude. */ +export function planAudioChunks( + audio: Float32Array, + options: PlanChunksOptions = {}, +): Array<[number, number]> { + return planChunks( + audio.length, + (start, end) => { + let sum = 0; + for (let i = start; i < end; i++) sum += audio[i]! * audio[i]!; + return sum / Math.max(1, end - start); + }, + options, + ); +} + +// --------------------------------------------------------------------------- +// Words → transcript chunks (the shape TranscriptView / export / live use) +// --------------------------------------------------------------------------- + +/** Sentence end: . ? ! optionally followed by closing quotes/brackets. */ +const SENTENCE_END = /[.?!]["'”’)\]]*$/; + +/** + * Group words into caption-sized chunks: a chunk ends after a sentence-final + * word, before a pause of `gapSeconds` (default 1 s), or before a word that + * would stretch it past `maxSegmentSeconds` (default 12 s). Chunk text has + * Whisper's leading-space convention, so concatenating chunks yields text. + */ +export function wordsToChunks( + words: readonly ParakeetWord[], + options: { gapSeconds?: number; maxSegmentSeconds?: number } = {}, +): TranscriptChunk[] { + const gap = options.gapSeconds ?? 1; + const maxSegment = options.maxSegmentSeconds ?? 12; + const chunks: TranscriptChunk[] = []; + let current: ParakeetWord[] = []; + + const close = () => { + if (current.length === 0) return; + chunks.push({ + text: ` ${current.map((w) => w.text).join(" ")}`, + timestamp: [current[0]!.start_time, current[current.length - 1]!.end_time], + }); + current = []; + }; + + for (const w of words) { + if (current.length > 0) { + const last = current[current.length - 1]!; + if (w.start_time - last.end_time >= gap || w.end_time - current[0]!.start_time > maxSegment) { + close(); + } + } + current.push(w); + if (SENTENCE_END.test(w.text)) close(); + } + close(); + return chunks; +} + +/** + * Run `recognize` over each planned window (a `subarray` view — never a copy + * of the whole file) and assemble the transcript. `onWindow` fires after each + * window with its new chunks and the absolute seconds processed so far. + */ +export async function transcribeWindows( + audio: Float32Array, + recognize: (window: Float32Array, offsetSeconds: number) => Promise, + callbacks: { + onWindow?: (chunks: TranscriptChunk[], processedSeconds: number, totalSeconds: number) => void; + } = {}, +): Promise<{ text: string; chunks: TranscriptChunk[] }> { + const totalSeconds = audio.length / PARAKEET_SAMPLE_RATE; + const chunks: TranscriptChunk[] = []; + for (const [start, end] of planAudioChunks(audio)) { + const words = await recognize(audio.subarray(start, end), start / PARAKEET_SAMPLE_RATE); + const windowChunks = wordsToChunks(words); + chunks.push(...windowChunks); + callbacks.onWindow?.(windowChunks, end / PARAKEET_SAMPLE_RATE, totalSeconds); + } + return { text: chunks.map((chunk) => chunk.text).join("").trim(), chunks }; +} + +// --------------------------------------------------------------------------- +// Tokenizer (vocab.txt: " " per line, SentencePiece "▁" = space) +// --------------------------------------------------------------------------- + +export function parseVocab(text: string): string[] { + const id2token: string[] = []; + for (const line of text.split(/\r?\n/)) { + const [token, idText] = line.trim().split(/\s+/); + const id = Number(idText); + if (!token || !Number.isInteger(id)) continue; + id2token[id] = token; + } + return id2token; +} + +/** + * The tokenizer shape parakeet.js's `ParakeetModel` reads (`id2token`, + * `blankId`, `blankToken`, `decode`). Its own `ParakeetTokenizer` is not + * exported and only loads by URL, so this mirrors its decode rules. + */ +export class ParakeetVocabTokenizer { + readonly blankToken = ""; + readonly blankId: number; + readonly sanitizedTokens: string[]; + + constructor(readonly id2token: string[]) { + const blank = id2token.indexOf(this.blankToken); + this.blankId = blank === -1 ? id2token.length : blank; + this.sanitizedTokens = id2token.map((token) => token.replace(/▁/g, " ")); + } + + decode(ids: readonly number[]): string { + let text = ""; + for (const id of ids) { + if (id === this.blankId) continue; + text += this.sanitizedTokens[id] ?? ""; + } + return text + .replace(/^\s+/, "") + .replace(/\s+(?=[^\w\s])/g, "") + .replace(/\s+/g, " ") + .trim(); + } +} + +// --------------------------------------------------------------------------- +// Assets +// --------------------------------------------------------------------------- + +export const PARAKEET_VOCAB_FILE = "vocab.txt"; + +/** ONNX file for one part at one dtype. Only int8 and fp16 are offered. */ +export function parakeetFileName(part: "encoder" | "decoder", dtype: WhisperDtype): string { + const base = part === "encoder" ? "encoder-model" : "decoder_joint-model"; + if (dtype === "q8") return `${base}.int8.onnx`; + if (dtype === "fp16") return `${base}.fp16.onnx`; + throw new Error(`Parakeet has no ${dtype} ${part} in this app (int8 or fp16 only).`); +} + +/** + * Hub URL pinned to the measured revision. Same `//` shape as + * Transformers.js downloads, so modelCache.ts's badges/storage panel/delete + * recognise the cached files without changes (D-005). + */ +export function parakeetFileUrl(file: string): string { + return `https://huggingface.co/${PARAKEET_MODEL_ID}/resolve/${PARAKEET_REVISION}/${file}`; +} diff --git a/src/worker/inference.worker.ts b/src/worker/inference.worker.ts index dc56d11..554b1b7 100644 --- a/src/worker/inference.worker.ts +++ b/src/worker/inference.worker.ts @@ -16,7 +16,19 @@ import { type PretrainedModelOptions, type ProgressInfo, } from "@huggingface/transformers"; -import { resolveDefaultDtype } from "../lib/models"; +// The very onnxruntime-web module Transformers.js imports (its web build +// imports "onnxruntime-web/webgpu"), so the Parakeet path shares one runtime, +// one wasm binary from /ort/, and one env (D-WS7a). +import * as ort from "onnxruntime-web/webgpu"; +import { getModel, modelFamily, resolveDefaultDtype } from "../lib/models"; +import { PARAKEET_VOCAB_FILE, parakeetFileName, parakeetFileUrl } from "../lib/parakeet"; +import { fetchModelFile, type ModelFileCacheStorage } from "./modelFiles"; +import { + createParakeetEngine, + parakeetSessionOptions, + type OrtLike, + type ParakeetEngine, +} from "./parakeetEngine"; import { LONG_FORM_DEFAULTS, transcribeLongForm } from "./longform"; import { createWhisperDeps, loadSileroVad, type SileroVad } from "./longform/transformers"; import { @@ -29,6 +41,7 @@ import { import { choosePipeline, deviceAttempts, + parakeetDeviceAttempts, dtypeRequiresF16, isUiToWorkerMessage, TranscribeProgressTracker, @@ -209,10 +222,20 @@ async function probeCapabilities( // Pipeline lifecycle (load, switch, dispose) // --------------------------------------------------------------------------- +/** A loaded model of either family (D-WS7a): the worker dispatches on it. */ +type LoadedModel = + | { family: "whisper"; asr: AutomaticSpeechRecognitionPipeline } + | { family: "parakeet"; engine: ParakeetEngine }; + interface ActivePipeline { /** Identity of this load request; a different key means "switch models". */ key: string; - promise: Promise; + promise: Promise; +} + +async function disposeModel(model: LoadedModel): Promise { + if (model.family === "whisper") await model.asr.dispose(); + else await model.engine.dispose(); } let active: ActivePipeline | null = null; @@ -300,7 +323,67 @@ async function createPipeline( throw lastError; } -function load(message: LoadMessage): Promise { +/** + * Parakeet load (D-WS7a): fetch vocab + decoder + encoder through the + * Cache-API downloader (same cache and progress messages as Transformers.js, + * D-005), then create the sessions on our shared onnxruntime-web. Device + * policy is `parakeetDeviceAttempts` (WASM unless WebGPU was chosen + * explicitly); the D-025 restart rule applies exactly as for Whisper. + */ +async function createParakeet(message: LoadMessage): Promise { + const capabilities = + message.device === "webgpu" + ? await probeCapabilities("webgpu") + : { webnn: false, webgpu: false }; + const attempts = parakeetDeviceAttempts(message.device, capabilities); + const downloads = { + caches: typeof caches === "undefined" ? undefined : (caches as unknown as ModelFileCacheStorage), + fetch: (url: string) => fetch(url), + onProgress: post, + }; + const file = async (name: string) => + fetchModelFile(parakeetFileUrl(name), name, downloads); + + let lastError: unknown; + let acceleratorAttempted = false; + for (const device of attempts) { + if (device === "wasm" && acceleratorAttempted && desiredWasmThreads() > 1) { + post({ type: "load-restart", modelId: message.modelId, preference: message.device, device: "wasm" }); + throw new RestartRequested(); + } + configureWasmThreads(device); + const dtype = message.dtype ?? resolveDefaultDtype(message.modelId, device); + const encoderDtype = typeof dtype === "string" ? dtype : dtype.encoder_model; + const decoderDtype = typeof dtype === "string" ? dtype : dtype.decoder_model_merged; + try { + const vocab = new TextDecoder().decode(await file(PARAKEET_VOCAB_FILE)); + const decoder = await file(parakeetFileName("decoder", decoderDtype)); + const engine = await createParakeetEngine( + ort as unknown as OrtLike, + { encoder: await file(parakeetFileName("encoder", encoderDtype)), decoder, vocab }, + parakeetSessionOptions(device === "webgpu" ? "webgpu" : "wasm"), + ); + const threads = ort.env.wasm.numThreads ?? 1; + console.info(`[whisper] parakeet ready: device=${device} wasmThreads=${threads}`); + post({ type: "load-progress", status: "ready", device, threads }); + return engine; + } catch (error) { + lastError = error; + if (device !== "wasm") acceleratorAttempted = true; + console.warn(`[whisper] parakeet creation failed on ${device}: ${errorMessage(error)}`); + } + } + throw lastError; +} + +async function createModel(message: LoadMessage): Promise { + if (modelFamily(getModel(message.modelId)) === "parakeet") { + return { family: "parakeet", engine: await createParakeet(message) }; + } + return { family: "whisper", asr: await createPipeline(message) }; +} + +function load(message: LoadMessage): Promise { const key = loadKey(message); // Idempotent: repeated "load" messages for the same model/device/dtype reuse // the in-flight or completed pipeline. @@ -312,13 +395,12 @@ function load(message: LoadMessage): Promise // Dispose the old pipeline (frees WASM memory / GPU buffers) before // creating the new one, so both never coexist. try { - const old = await previous.promise; - await old.dispose(); + await disposeModel(await previous.promise); } catch { // The previous load itself failed — nothing to dispose. } } - return createPipeline(message); + return createModel(message); })(); active = { key, promise }; promise.catch(() => { @@ -369,6 +451,34 @@ const SAMPLE_RATE = 16_000; const CHUNK_LENGTH_S = 30; const STRIDE_LENGTH_S = 5; +/** + * Parakeet pass (D-WS7a). English-only: `language`/`task` are ignored (the UI + * never sends them for a non-multilingual model). File-flow passes stream one + * `partial` + one `transcribe-progress` per ≤ 30 s window — absolute seconds + * straight from the window plan, so no TranscribeProgressTracker is needed. + * Live passes resolve wholesale via "result", exactly like Whisper's (D-016). + */ +async function transcribeParakeet( + engine: ParakeetEngine, + message: TranscribeMessage, +): Promise { + const stream = message.stream ?? message.requestId === undefined; + const totalSeconds = message.audio.length / SAMPLE_RATE; + if (stream) { + post({ type: "transcribe-progress", processedSeconds: 0, totalSeconds, requestId: message.requestId }); + } + const output = await engine.transcribe(message.audio, { + ...(stream && { + onWindow: (chunks: TranscriptChunk[], processedSeconds: number) => { + const text = chunks.map((chunk) => chunk.text).join(""); + if (text) post({ type: "partial", text, requestId: message.requestId }); + post({ type: "transcribe-progress", processedSeconds, totalSeconds, requestId: message.requestId }); + }, + }), + }); + post({ type: "result", text: output.text, chunks: output.chunks, requestId: message.requestId }); +} + async function transcribe(message: TranscribeMessage): Promise { if (active === null) { post({ @@ -379,7 +489,12 @@ async function transcribe(message: TranscribeMessage): Promise { }); return; } - const asr = await active.promise; + const loaded = await active.promise; + if (loaded.family === "parakeet") { + await transcribeParakeet(loaded.engine, message); + return; + } + const asr = loaded.asr; const streaming = message.stream ?? message.requestId === undefined; // Live passes (not streaming) always keep the chunked path (D-016/D-024). diff --git a/src/worker/modelFiles.ts b/src/worker/modelFiles.ts new file mode 100644 index 0000000..72a1109 --- /dev/null +++ b/src/worker/modelFiles.ts @@ -0,0 +1,121 @@ +/** + * Weight-file download + cache for engine paths that don't go through + * Transformers.js's own loader (Parakeet, D-WS7a). + * + * Files land in the same Cache Storage cache Transformers.js uses + * (`transformers-cache`, D-005), keyed by their Hub URL, so the downloaded + * badge, the storage panel, and delete (src/lib/modelCache.ts) work for them + * unchanged. Progress is reported as the same `load-progress` messages the + * Whisper path forwards, so ModelLoadProgress needs no second code path. + * + * A file is `put` only once its whole body has arrived: a cancelled or + * failed download (worker terminated, network drop) leaves nothing behind + * that could later show as "available offline". + * + * Dependencies are injected so `bun test` runs this without Cache Storage. + */ + +import { TRANSFORMERS_CACHE_NAME } from "../lib/modelCache"; +import type { LoadProgressMessage } from "./protocol"; + +/** The slice of `Cache` used here. */ +export interface ModelFileCache { + match(url: string): Promise; + put(url: string, response: Response): Promise; +} + +/** The slice of `CacheStorage` used here. */ +export interface ModelFileCacheStorage { + open(name: string): Promise; +} + +export interface FetchModelFileDeps { + /** Cache Storage, or undefined where it doesn't exist (download uncached). */ + caches: ModelFileCacheStorage | undefined; + fetch: (url: string) => Promise; + onProgress: (message: LoadProgressMessage) => void; +} + +/** Fetch one model file, from cache when present, else download and cache it. */ +export async function fetchModelFile( + url: string, + file: string, + deps: FetchModelFileDeps, +): Promise { + deps.onProgress({ type: "load-progress", status: "initiate", file }); + const cache = deps.caches ? await deps.caches.open(TRANSFORMERS_CACHE_NAME) : undefined; + + const cached = await cache?.match(url); + if (cached) { + const bytes = new Uint8Array(await cached.arrayBuffer()); + deps.onProgress({ type: "load-progress", status: "done", file }); + return bytes; + } + + deps.onProgress({ type: "load-progress", status: "download", file }); + const response = await deps.fetch(url); + if (!response.ok || !response.body) { + throw new Error(`Could not download ${file} (HTTP ${response.status})`); + } + const total = Number(response.headers.get("Content-Length")) || 0; + const bytes = await readWithProgress(response.body, total, (loaded) => { + const size = total || loaded; + deps.onProgress({ + type: "load-progress", + status: "progress", + file, + loaded, + total: size, + progress: (loaded / size) * 100, + }); + }); + + if (cache) { + try { + await cache.put( + url, + new Response(bytes as Uint8Array, { + headers: { "Content-Length": String(bytes.length) }, + }), + ); + } catch { + // Quota exceeded / private mode: the model still runs this session, it + // just won't be available offline. + } + } + deps.onProgress({ type: "load-progress", status: "done", file }); + return bytes; +} + +/** + * Read a body into one buffer — preallocated when the length is known, so a + * 650 MB encoder isn't also held as a list of chunks. Reports progress at + * most once per whole percent (a large file arrives in thousands of reads). + */ +async function readWithProgress( + body: ReadableStream, + total: number, + onLoaded: (loaded: number) => void, +): Promise { + let buffer = new Uint8Array(total); + let loaded = 0; + let lastPercent = -1; + const reader = body.getReader(); + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + if (loaded + value.length > buffer.length) { + const grown = new Uint8Array(Math.max(buffer.length * 2, loaded + value.length)); + grown.set(buffer.subarray(0, loaded)); + buffer = grown; + } + buffer.set(value, loaded); + loaded += value.length; + const percent = total ? Math.floor((loaded / total) * 100) : loaded; + if (percent !== lastPercent) { + lastPercent = percent; + onLoaded(loaded); + } + } + return loaded === buffer.length ? buffer : buffer.slice(0, loaded); +} diff --git a/src/worker/ortShim.ts b/src/worker/ortShim.ts new file mode 100644 index 0000000..824e043 --- /dev/null +++ b/src/worker/ortShim.ts @@ -0,0 +1,8 @@ +/** + * Stand-in for the bare `onnxruntime-web` import inside parakeet.js + * (D-WS7a). parakeet.js depends on its own onnxruntime-web 1.24.1, imported + * only by its `initOrt` helper, which this app never calls — but a bundler + * would still pull that second runtime into the worker. `scripts/ortAlias.ts` + * redirects the import here, to the one runtime Transformers.js uses. + */ +export * from "onnxruntime-web/webgpu"; diff --git a/src/worker/parakeetEngine.ts b/src/worker/parakeetEngine.ts new file mode 100644 index 0000000..c3cf196 --- /dev/null +++ b/src/worker/parakeetEngine.ts @@ -0,0 +1,130 @@ +/** + * Parakeet TDT engine (Decisions.md D-WS7a): the non-Whisper inference path + * of the browser engine, used by the inference worker for catalog entries + * with `family: "parakeet"`. + * + * Runtime: parakeet.js's `ParakeetModel` (TDT greedy decoding over the + * encoder + decoder_joint ONNX pair, JS mel preprocessor) on an injected + * onnxruntime — in the worker that is the very `onnxruntime-web/webgpu` + * instance Transformers.js uses (one runtime, its wasm self-hosted at + * /ort/), in tests onnxruntime-web or onnxruntime-node. We deliberately + * bypass parakeet.js's `fromUrls`/`fromHub`/`initOrt`: those import their own + * onnxruntime-web, point it at the jsdelivr CDN (blocked by COEP, and not + * offline), and fetch weights outside our Cache-API bookkeeping. + * + * Windowing and word → chunk mapping are pure (src/lib/parakeet.ts). + */ + +import { JsPreprocessor } from "parakeet.js/mel"; +import { ParakeetModel } from "parakeet.js/parakeet"; +import { ParakeetVocabTokenizer, parseVocab, transcribeWindows, PARAKEET_SAMPLE_RATE } from "../lib/parakeet"; +import type { TranscriptChunk } from "./protocol"; + +/** The slice of an onnxruntime module this path needs (web and node agree). */ +export interface OrtLike { + InferenceSession: { + create(model: Uint8Array, options?: Record): Promise<{ release?(): Promise }>; + }; + Tensor: unknown; +} + +export interface ParakeetAssets { + encoder: Uint8Array; + decoder: Uint8Array; + vocab: string; +} + +export interface ParakeetSessionOptions { + encoder: Record; + decoder: Record; +} + +export interface ParakeetTranscribeCallbacks { + /** After each ≤ 30 s window: its new chunks and absolute progress. */ + onWindow?: (chunks: TranscriptChunk[], processedSeconds: number, totalSeconds: number) => void; +} + +export interface ParakeetEngine { + transcribe( + audio: Float32Array, + callbacks?: ParakeetTranscribeCallbacks, + ): Promise<{ text: string; chunks: TranscriptChunk[] }>; + dispose(): Promise; +} + +/** + * Session options per device. WASM: both parts on the wasm EP at the full + * graph optimization level — measured fine for these int8 graphs on our + * onnxruntime-web build (D-014's "basic" cap is about Whisper's 8-bit tied + * embeddings, which Parakeet doesn't have). WebGPU (explicit opt-in only): + * the fp16 encoder on WebGPU, the small autoregressive decoder on WASM, the + * split parakeet.js itself uses ("webgpu-hybrid"). + */ +export function parakeetSessionOptions(device: "wasm" | "webgpu"): ParakeetSessionOptions { + const wasm = { executionProviders: ["wasm"], graphOptimizationLevel: "all" }; + return device === "webgpu" + ? { encoder: { executionProviders: ["webgpu"], graphOptimizationLevel: "all" }, decoder: wasm } + : { encoder: wasm, decoder: wasm }; +} + +/** What we call on parakeet.js's model (its constructor isn't in its .d.ts). */ +interface ParakeetModelLike { + transcribe( + audio: Float32Array, + sampleRate: number, + options: Record, + ): Promise<{ words: Array<{ text: string; start_time: number; end_time: number }> }>; +} + +/** + * Create the two ONNX sessions and wrap them. Sessions are created one after + * the other (encoder first) so only one weight buffer is being parsed at a + * time — the D-027 peak-memory lesson. Callers should drop their references + * to `assets` afterwards so the 650 MB encoder buffer can be collected. + */ +export async function createParakeetEngine( + ort: OrtLike, + assets: ParakeetAssets, + sessionOptions: ParakeetSessionOptions, +): Promise { + const tokenizer = new ParakeetVocabTokenizer(parseVocab(assets.vocab)); + const encoderSession = await ort.InferenceSession.create(assets.encoder, sessionOptions.encoder); + let joinerSession: Awaited>; + try { + joinerSession = await ort.InferenceSession.create(assets.decoder, sessionOptions.decoder); + } catch (error) { + await encoderSession.release?.(); + throw error; + } + const Model = ParakeetModel as unknown as new (config: Record) => ParakeetModelLike; + const model = new Model({ + tokenizer, + encoderSession, + joinerSession, + preprocessor: new JsPreprocessor({ nMels: 128 }), + ort, + nMels: 128, + }); + + return { + transcribe(audio, callbacks = {}) { + return transcribeWindows( + audio, + async (window, offsetSeconds) => { + const result = await model.transcribe(window, PARAKEET_SAMPLE_RATE, { + returnTimestamps: true, + timeOffset: offsetSeconds, + // Profiling defaults to on and logs every call; we don't need it. + enableProfiling: false, + }); + return result.words; + }, + callbacks, + ); + }, + async dispose() { + await encoderSession.release?.(); + await joinerSession.release?.(); + }, + }; +} diff --git a/src/worker/protocol.ts b/src/worker/protocol.ts index b6d9feb..955698b 100644 --- a/src/worker/protocol.ts +++ b/src/worker/protocol.ts @@ -300,6 +300,25 @@ export function choosePipeline(input: { return input.longForm === true && input.vadReady ? "longform" : "chunked"; } +/** + * Device attempts for the Parakeet engine path (D-WS7a). Narrower than + * `deviceAttempts`: "auto" always means WASM int8 — the only measured path — + * because WebGPU int8 is reported broken and WebGPU needs the separate + * 1.24 GB fp16 encoder, so auto-selecting it would silently double the + * download on an unverified backend. An explicit WebGPU choice is honoured + * only where shader-f16 exists (fp16 encoder), still falling back to WASM; + * there is no WebNN path. + */ +export function parakeetDeviceAttempts( + preference: DevicePreference, + capabilities: DeviceCapabilities, +): InferenceDevice[] { + if (preference === "webgpu" && capabilities.webgpu && capabilities.webgpuF16) { + return ["webgpu", "wasm"]; + } + return ["wasm"]; +} + // --------------------------------------------------------------------------- // Runtime guards (postMessage data arrives untyped) // --------------------------------------------------------------------------- diff --git a/tests/mobileMemoryGuard.test.ts b/tests/mobileMemoryGuard.test.ts index 2772467..859cbd4 100644 --- a/tests/mobileMemoryGuard.test.ts +++ b/tests/mobileMemoryGuard.test.ts @@ -1,5 +1,11 @@ import { describe, test, expect } from "bun:test"; -import { isLikelyMobile, shouldWarnBeforeLoad } from "../src/lib/mobileMemoryGuard"; +import { + isLikelyMobile, + isModelAllowedOnDevice, + modelsForDevice, + shouldWarnBeforeLoad, +} from "../src/lib/mobileMemoryGuard"; +import { MODELS, PARAKEET_MODEL_ID } from "../src/lib/models"; // Pure UA/heuristic logic (D-027) — no browser, no navigator. @@ -99,3 +105,27 @@ describe("shouldWarnBeforeLoad", () => { ).toBe(false); }); }); + +describe("desktop-only models (D-WS7a)", () => { + const parakeet = MODELS.find((model) => model.id === PARAKEET_MODEL_ID)!; + const phone = { userAgent: ANDROID_UA, maxTouchPoints: 5, deviceMemoryGB: 8 }; + const ipad = { userAgent: IPAD_DESKTOP_UA, maxTouchPoints: 5, deviceMemoryGB: undefined }; + const desktop = { userAgent: DESKTOP_WINDOWS_UA, maxTouchPoints: 0, deviceMemoryGB: 8 }; + + test("Parakeet is never offered on phones or tablets, even high-memory ones", () => { + expect(isModelAllowedOnDevice(parakeet, phone)).toBe(false); + expect(isModelAllowedOnDevice(parakeet, ipad)).toBe(false); + expect(isModelAllowedOnDevice(parakeet, desktop)).toBe(true); + }); + + test("Whisper models stay allowed everywhere (they get the D-027 confirm instead)", () => { + for (const model of MODELS.filter((m) => m.id !== PARAKEET_MODEL_ID)) { + expect(isModelAllowedOnDevice(model, phone)).toBe(true); + } + }); + + test("modelsForDevice filters the list, keeping order", () => { + expect(modelsForDevice(MODELS, phone)).toEqual(MODELS.filter((m) => m !== parakeet)); + expect(modelsForDevice(MODELS, desktop)).toEqual([...MODELS]); + }); +}); diff --git a/tests/modelFiles.test.ts b/tests/modelFiles.test.ts new file mode 100644 index 0000000..1b3bfcb --- /dev/null +++ b/tests/modelFiles.test.ts @@ -0,0 +1,121 @@ +import { describe, expect, test } from "bun:test"; +import { fetchModelFile, type ModelFileCacheStorage } from "../src/worker/modelFiles"; +import type { LoadProgressMessage } from "../src/worker/protocol"; + +// Cache-API-backed weight download for the Parakeet path (D-WS7a, D-005). + +const URL_A = "https://huggingface.co/org/model/resolve/rev/encoder.onnx"; + +function fakeCaches(initial: Record = {}) { + const stored = new Map(Object.entries(initial)); + const opened: string[] = []; + const storage: ModelFileCacheStorage = { + async open(name) { + opened.push(name); + return { + async match(url) { + const bytes = stored.get(url); + return bytes ? new Response(bytes as Uint8Array) : undefined; + }, + async put(url, response) { + stored.set(url, new Uint8Array(await response.arrayBuffer())); + }, + }; + }, + }; + return { storage, stored, opened }; +} + +function streamingFetch(bytes: Uint8Array, pieces: number, status = 200) { + const calls: string[] = []; + const fetchFn = async (url: string) => { + calls.push(url); + if (status !== 200) return new Response("nope", { status }); + const size = Math.ceil(bytes.length / pieces); + const body = new ReadableStream({ + start(controller) { + for (let i = 0; i < bytes.length; i += size) controller.enqueue(bytes.slice(i, i + size)); + controller.close(); + }, + }); + return new Response(body, { headers: { "Content-Length": String(bytes.length) } }); + }; + return { fetchFn, calls }; +} + +describe("fetchModelFile", () => { + test("a cached file is returned without touching the network", async () => { + const cached = new Uint8Array([1, 2, 3]); + const { storage, opened } = fakeCaches({ [URL_A]: cached }); + const { fetchFn, calls } = streamingFetch(new Uint8Array(), 1); + const events: LoadProgressMessage[] = []; + const bytes = await fetchModelFile(URL_A, "encoder.onnx", { + caches: storage, + fetch: fetchFn, + onProgress: (event) => events.push(event), + }); + expect([...bytes]).toEqual([1, 2, 3]); + expect(calls).toEqual([]); + expect(opened).toEqual(["transformers-cache"]); + expect(events.map((event) => event.status)).toEqual(["initiate", "done"]); + }); + + test("a miss downloads with byte progress and stores the complete file", async () => { + const payload = new Uint8Array(1000).map((_, i) => i % 251); + const { storage, stored } = fakeCaches(); + const { fetchFn, calls } = streamingFetch(payload, 4); + const events: LoadProgressMessage[] = []; + const bytes = await fetchModelFile(URL_A, "encoder.onnx", { + caches: storage, + fetch: fetchFn, + onProgress: (event) => events.push(event), + }); + expect(calls).toEqual([URL_A]); + expect([...bytes]).toEqual([...payload]); + expect([...stored.get(URL_A)!]).toEqual([...payload]); + const progress = events.filter((event) => event.status === "progress"); + expect(progress.length).toBe(4); + const last = progress[progress.length - 1]!; + expect(last).toMatchObject({ file: "encoder.onnx", loaded: 1000, total: 1000, progress: 100 }); + expect(events[0]!.status).toBe("initiate"); + expect(events[events.length - 1]!.status).toBe("done"); + }); + + test("an HTTP error throws a plain message and stores nothing", async () => { + const { storage, stored } = fakeCaches(); + const { fetchFn } = streamingFetch(new Uint8Array(), 1, 404); + await expect( + fetchModelFile(URL_A, "encoder.onnx", { caches: storage, fetch: fetchFn, onProgress: () => {} }), + ).rejects.toThrow("Could not download encoder.onnx (HTTP 404)"); + expect(stored.size).toBe(0); + }); + + test("an interrupted download stores nothing (no false 'available offline')", async () => { + const { storage, stored } = fakeCaches(); + const fetchFn = async () => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(10)); + controller.error(new Error("network down")); + }, + }), + { headers: { "Content-Length": "100" } }, + ); + await expect( + fetchModelFile(URL_A, "encoder.onnx", { caches: storage, fetch: fetchFn, onProgress: () => {} }), + ).rejects.toThrow(); + expect(stored.size).toBe(0); + }); + + test("without Cache Storage it still downloads", async () => { + const { fetchFn, calls } = streamingFetch(new Uint8Array([9]), 1); + const bytes = await fetchModelFile(URL_A, "encoder.onnx", { + caches: undefined, + fetch: fetchFn, + onProgress: () => {}, + }); + expect([...bytes]).toEqual([9]); + expect(calls).toEqual([URL_A]); + }); +}); diff --git a/tests/models.test.ts b/tests/models.test.ts index 4110785..ac3df68 100644 --- a/tests/models.test.ts +++ b/tests/models.test.ts @@ -10,6 +10,8 @@ import { formatModelSize, getFasterWhisperModel, getModel, + modelFamily, + PARAKEET_MODEL_ID, resolveDefaultDtype, taskForModel, } from "../src/lib/models"; @@ -21,6 +23,8 @@ 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); @@ -35,12 +39,15 @@ describe("MODELS catalog", () => { "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 MODELS) { + 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")); @@ -72,10 +79,10 @@ describe("MODELS catalog", () => { expect(getModel("onnx-community/whisper-nonexistent")).toBeUndefined(); }); - test("every model maps to a faster-whisper name for the server engine (D-013)", () => { + 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 MODELS.filter((m) => !m.id.startsWith(`${DANISH_MODELS_HF_OWNER}/`))) { + 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 @@ -298,3 +305,41 @@ describe("taskForModel (D-WS4a)", () => { 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); + }); +}); diff --git a/tests/parakeet-real-model.test.ts b/tests/parakeet-real-model.test.ts new file mode 100644 index 0000000..0b14314 --- /dev/null +++ b/tests/parakeet-real-model.test.ts @@ -0,0 +1,157 @@ +import { afterAll, describe, expect, test } from "bun:test"; +import { existsSync, mkdirSync } from "node:fs"; +import { join } from "node:path"; +import ffmpegPath from "ffmpeg-static"; +import * as ort from "onnxruntime-web"; +import { LIVE_TUNING, LiveTranscriber, type LiveTranscribeFn } from "../src/lib/live"; +import { PARAKEET_MODEL_ID, resolveDefaultDtype } from "../src/lib/models"; +import { PARAKEET_VOCAB_FILE, parakeetFileName, parakeetFileUrl } from "../src/lib/parakeet"; +import { + createParakeetEngine, + parakeetSessionOptions, + type OrtLike, + type ParakeetEngine, +} from "../src/worker/parakeetEngine"; +import { wordErrorRate } from "./dsp"; + +/** + * Real-model accuracy of the Parakeet engine path (D-WS7a), opt-in: + * + * PARAKEET_REAL=1 bun test tests/parakeet-real-model.test.ts + * + * Runs the worker's own `createParakeetEngine` on onnxruntime-web's WASM + * backend (the browser engine's baseline runtime, here under Bun) with the + * catalog's WASM int8 weights, over the English clips in tests/audio_files: + * whole-file mode, and live mode through the production LiveTranscriber. + * Held to the same bars as Whisper Large v3 Turbo in live-real-model.test.ts. + * First run downloads ~670 MB into node_modules/@huggingface/transformers/.cache/. + * Opt-in (not in `bun test` by default) so the default suite's runtime and + * network use don't change; run it when touching src/lib/parakeet.ts or + * src/worker/parakeetEngine.ts. + */ + +const ENABLED = process.env.PARAKEET_REAL === "1"; +const RATE = 16_000; +const AUDIO_DIR = join(import.meta.dir, "audio_files"); +const MODEL_DIR = join( + import.meta.dir, + "../node_modules/@huggingface/transformers/.cache", + PARAKEET_MODEL_ID, +); +const CLIPS = [ + { base: "en_aiproblem", file: "en_aiproblem.m4a", maxWer: 0.25 }, + { base: "en_alien", file: "en_alien.opus", maxWer: 0.35 }, +]; + +function normalize(text: string): string { + return text + .toLowerCase() + .replace(/['’]/g, "") + .replace(/[.,!?;:"“”‘‚«»()\-–—…]/g, " ") + .replace(/\s+/g, " ") + .trim(); +} + +async function decode16k(path: string): Promise { + const proc = Bun.spawn( + [ffmpegPath as unknown as string, "-hide_banner", "-loglevel", "error", "-i", path, + "-ac", "1", "-ar", String(RATE), "-f", "f32le", "-"], + { stdout: "pipe", stderr: "pipe" }, + ); + const [bytes, exitCode] = await Promise.all([new Response(proc.stdout).arrayBuffer(), proc.exited]); + if (exitCode !== 0) throw new Error(`ffmpeg failed on ${path}`); + return new Float32Array(bytes); +} + +async function modelFile(file: string): Promise { + const path = join(MODEL_DIR, file); + if (!existsSync(path)) { + mkdirSync(MODEL_DIR, { recursive: true }); + const response = await fetch(parakeetFileUrl(file)); + if (!response.ok) throw new Error(`download ${file}: HTTP ${response.status}`); + await Bun.write(path, response); + } + return new Uint8Array(await Bun.file(path).arrayBuffer()); +} + +let enginePromise: Promise | null = null; +function loadEngine(): Promise { + enginePromise ??= (async () => { + ort.env.wasm.numThreads = 8; + const dtype = resolveDefaultDtype(PARAKEET_MODEL_ID, "wasm"); + if (typeof dtype !== "string") throw new Error("expected one WASM dtype"); + return createParakeetEngine( + ort as unknown as OrtLike, + { + encoder: await modelFile(parakeetFileName("encoder", dtype)), + decoder: await modelFile(parakeetFileName("decoder", dtype)), + vocab: new TextDecoder().decode(await modelFile(PARAKEET_VOCAB_FILE)), + }, + parakeetSessionOptions("wasm"), + ); + })(); + return enginePromise; +} + +afterAll(async () => { + if (enginePromise) await (await enginePromise).dispose(); +}); + +describe("Parakeet engine accuracy (real model, onnxruntime-web WASM)", () => { + for (const clip of CLIPS) { + test.skipIf(!ENABLED)( + `${clip.base}: whole-file and live transcripts stay close to the reference`, + async () => { + const reference = await Bun.file(join(AUDIO_DIR, `${clip.base}.txt`)).text(); + const audio = await decode16k(join(AUDIO_DIR, clip.file)); + const duration = audio.length / RATE; + const engine = await loadEngine(); + + const fileStart = performance.now(); + let lastProgress = 0; + const whole = await engine.transcribe(audio.slice(), { + onWindow: (_chunks, processed) => { + expect(processed).toBeGreaterThan(lastProgress); + lastProgress = processed; + }, + }); + const fileMs = performance.now() - fileStart; + const fileWer = wordErrorRate(normalize(reference), normalize(whole.text)); + expect(lastProgress).toBeCloseTo(duration, 3); + expect(whole.chunks[0]!.timestamp[0]).toBeLessThan(duration / 2); + expect(whole.chunks.at(-1)!.timestamp[1]!).toBeGreaterThan(duration / 2); + expect(whole.chunks.at(-1)!.timestamp[1]!).toBeLessThanOrEqual(duration + 1); + + // Live: production segmenter, clock advancing with the audio fed. + const transcribe: LiveTranscribeFn = (segment) => engine.transcribe(segment); + let fed = 0; + const liveStart = performance.now(); + const live = new LiveTranscriber({ + ...LIVE_TUNING, + transcribe, + onUpdate: () => {}, + now: () => fed / (RATE / 1000), + }); + const block = RATE / 10; + for (let i = 0; i < audio.length; i += block) { + live.push(audio.slice(i, Math.min(i + block, audio.length))); + fed += Math.min(block, audio.length - i); + await new Promise((resolve) => setTimeout(resolve, 0)); + } + const liveResult = await live.stop(); + const liveMs = performance.now() - liveStart; + const liveWer = wordErrorRate(normalize(reference), normalize(liveResult.text)); + + console.info( + `[parakeet-real] ${clip.base}: ${duration.toFixed(1)} s audio — file WER ` + + `${fileWer.toFixed(3)} in ${(fileMs / 1000).toFixed(1)} s; live WER ` + + `${liveWer.toFixed(3)} in ${(liveMs / 1000).toFixed(1)} s`, + ); + expect(fileWer).toBeLessThanOrEqual(clip.maxWer); + expect(liveWer).toBeLessThanOrEqual(clip.maxWer); + expect(liveWer).toBeLessThanOrEqual(fileWer + 0.1); + }, + 1_800_000, + ); + } +}); diff --git a/tests/parakeet.test.ts b/tests/parakeet.test.ts new file mode 100644 index 0000000..7711c83 --- /dev/null +++ b/tests/parakeet.test.ts @@ -0,0 +1,201 @@ +import { describe, expect, test } from "bun:test"; +import { + ParakeetVocabTokenizer, + parakeetFileName, + parakeetFileUrl, + parseVocab, + planAudioChunks, + planChunks, + transcribeWindows, + wordsToChunks, + type ParakeetWord, +} from "../src/lib/parakeet"; +import { transcriptToSrt, transcriptToTimestampedTxt } from "../src/lib/export"; + +// Pure Parakeet engine logic (D-WS7a) — no model, no runtime. + +const RATE = 16_000; + +function word(text: string, start: number, end: number): ParakeetWord { + return { text, start_time: start, end_time: end }; +} + +function expectContiguous(windows: Array<[number, number]>, length: number) { + expect(windows[0]![0]).toBe(0); + expect(windows[windows.length - 1]![1]).toBe(length); + for (let i = 1; i < windows.length; i++) { + expect(windows[i]![0]).toBe(windows[i - 1]![1]); + } +} + +describe("planChunks", () => { + const flat = () => 1; + + test("audio up to the window length is a single window", () => { + expect(planChunks(10 * RATE, flat)).toEqual([[0, 10 * RATE]]); + expect(planChunks(30 * RATE, flat)).toEqual([[0, 30 * RATE]]); + expect(planChunks(0, flat)).toEqual([]); + }); + + test("windows never exceed 30 s, are contiguous, and cover everything", () => { + const length = 95 * RATE + 123; + const windows = planChunks(length, flat); + expectContiguous(windows, length); + for (const [start, end] of windows) { + expect(end - start).toBeLessThanOrEqual(30 * RATE); + expect(end).toBeGreaterThan(start); + } + }); + + test("a 90-minute file plans ≥ 180 bounded windows (no full-length copies needed)", () => { + const length = 90 * 60 * RATE; + const windows = planChunks(length, flat); + expect(windows.length).toBeGreaterThanOrEqual(180); + expectContiguous(windows, length); + for (const [start, end] of windows) expect(end - start).toBeLessThanOrEqual(30 * RATE); + }); + + test("cuts in the quietest 100 ms of the last 6 s before the limit", () => { + // Loud everywhere except a pause at 27.0–27.1 s. + const quietStart = 27 * RATE; + const energy = (start: number, end: number) => + start >= quietStart && end <= quietStart + 0.1 * RATE ? 0 : 1; + const windows = planChunks(60 * RATE, energy); + const firstCut = windows[0]![1]; + expect(firstCut).toBeGreaterThanOrEqual(quietStart); + expect(firstCut).toBeLessThanOrEqual(quietStart + 0.1 * RATE); + }); + + test("planAudioChunks finds a real silent gap in PCM", () => { + const audio = new Float32Array(50 * RATE); + for (let i = 0; i < audio.length; i++) audio[i] = Math.sin(i / 7) * 0.5; + audio.fill(0, 26 * RATE, 26.5 * RATE); + const windows = planAudioChunks(audio); + expect(windows[0]![1]).toBeGreaterThanOrEqual(26 * RATE); + expect(windows[0]![1]).toBeLessThanOrEqual(26.5 * RATE); + expectContiguous(windows, audio.length); + }); +}); + +describe("wordsToChunks", () => { + test("splits after sentence-ending punctuation, keeping absolute timestamps", () => { + const chunks = wordsToChunks([ + word("Hello", 0.5, 0.8), + word("there.", 0.9, 1.2), + word("How", 1.5, 1.7), + word("are", 1.7, 1.8), + word("you?", 1.8, 2.1), + ]); + expect(chunks).toEqual([ + { text: " Hello there.", timestamp: [0.5, 1.2] }, + { text: " How are you?", timestamp: [1.5, 2.1] }, + ]); + }); + + test("chunks export to .srt and timestamped .txt like Whisper's", () => { + const chunks = wordsToChunks([word("Hello", 0.5, 0.8), word("there.", 0.9, 1.2), word("Bye.", 61.5, 62.25)]); + expect(transcriptToSrt(chunks)).toBe( + "1\n00:00:00,500 --> 00:00:01,200\nHello there.\n\n2\n00:01:01,500 --> 00:01:02,250\nBye.\n", + ); + expect(transcriptToTimestampedTxt(chunks)).toContain("Bye."); + }); + + test("closing quotes after the full stop still end the sentence", () => { + const chunks = wordsToChunks([word('"Yes."', 0, 0.4), word("Then", 0.5, 0.7)]); + expect(chunks.map((chunk) => chunk.text)).toEqual([' "Yes."', " Then"]); + }); + + test("a pause of at least 1 s starts a new chunk", () => { + const chunks = wordsToChunks([word("one", 0, 0.3), word("two", 1.4, 1.6)]); + expect(chunks).toHaveLength(2); + }); + + test("a run-on sentence is split once it reaches 12 s", () => { + const words = Array.from({ length: 40 }, (_, i) => word(`w${i}`, i * 0.5, i * 0.5 + 0.4)); + const chunks = wordsToChunks(words); + expect(chunks.length).toBeGreaterThan(1); + for (const chunk of chunks) { + expect((chunk.timestamp[1] ?? 0) - chunk.timestamp[0]).toBeLessThanOrEqual(12.5); + } + expect(chunks.map((c) => c.text.trim()).join(" ").split(" ")).toHaveLength(40); + }); + + test("no words → no chunks", () => { + expect(wordsToChunks([])).toEqual([]); + }); +}); + +describe("vocab tokenizer", () => { + const vocab = " 0\n▁the 1\n▁cat 2\ns 3\n. 4\n 5\n"; + + test("parseVocab indexes tokens by id", () => { + expect(parseVocab(vocab)).toEqual(["", "▁the", "▁cat", "s", ".", ""]); + }); + + test("decode joins pieces, skips blank, and attaches punctuation", () => { + const tokenizer = new ParakeetVocabTokenizer(parseVocab(vocab)); + expect(tokenizer.blankId).toBe(5); + expect(tokenizer.blankToken).toBe(""); + expect(tokenizer.decode([1, 5, 2, 3, 4])).toBe("the cats."); + }); +}); + +describe("Parakeet files", () => { + test("int8 and fp16 file names; other dtypes are not offered", () => { + expect(parakeetFileName("encoder", "q8")).toBe("encoder-model.int8.onnx"); + expect(parakeetFileName("decoder", "q8")).toBe("decoder_joint-model.int8.onnx"); + expect(parakeetFileName("encoder", "fp16")).toBe("encoder-model.fp16.onnx"); + expect(() => parakeetFileName("encoder", "q4")).toThrow(); + }); + + test("URLs are pinned to the measured revision under the model id", () => { + expect(parakeetFileUrl("vocab.txt")).toBe( + "https://huggingface.co/ysdede/parakeet-tdt-0.6b-v3-onnx/resolve/" + + "e5356bf6e24bdb7a269536b0e9e7869bc034b893/vocab.txt", + ); + }); +}); + +describe("transcribeWindows", () => { + test("recognizes each window with its offset and reports progress + new chunks", async () => { + const audio = new Float32Array(70 * RATE).fill(0.1); + const offsets: number[] = []; + const sizes: number[] = []; + const progress: number[] = []; + const streamed: string[] = []; + const result = await transcribeWindows( + audio, + async (window, offset) => { + offsets.push(offset); + sizes.push(window.length); + return [word(`at${Math.round(offset)}.`, offset + 0.1, offset + 0.5)]; + }, + { + onWindow: (chunks, processed, total) => { + progress.push(processed); + expect(total).toBe(70); + streamed.push(...chunks.map((chunk) => chunk.text)); + }, + }, + ); + expect(offsets[0]).toBe(0); + expect(offsets.length).toBe(3); + for (const size of sizes) expect(size).toBeLessThanOrEqual(30 * RATE); + expect(progress[progress.length - 1]).toBe(70); + expect(progress).toEqual([...progress].sort((a, b) => a - b)); + expect(result.chunks.map((c) => c.text)).toEqual(streamed); + expect(result.text).toBe(streamed.join("").trim()); + expect(result.chunks[1]!.timestamp[0]).toBeCloseTo(offsets[1]! + 0.1, 5); + }); + + test("silence-only audio: empty transcript, progress still reaches the end", async () => { + const progress: number[] = []; + const result = await transcribeWindows( + new Float32Array(40 * RATE), + async () => [], + { onWindow: (_chunks, processed) => progress.push(processed) }, + ); + expect(result).toEqual({ text: "", chunks: [] }); + expect(progress[progress.length - 1]).toBe(40); + }); +}); diff --git a/tests/protocol.test.ts b/tests/protocol.test.ts index a08bdaf..c96f07f 100644 --- a/tests/protocol.test.ts +++ b/tests/protocol.test.ts @@ -2,6 +2,7 @@ import { describe, test, expect } from "bun:test"; import { applyLoadProgress, deviceAttempts, + parakeetDeviceAttempts, dtypeRequiresF16, initialLoadProgress, choosePipeline, @@ -247,6 +248,29 @@ describe("deviceAttempts (D-004, D-017)", () => { }); }); +describe("parakeetDeviceAttempts (D-WS7a)", () => { + const everything = { webnn: true, webgpu: true, webgpuF16: true }; + + test("auto never picks an accelerator for Parakeet", () => { + expect(parakeetDeviceAttempts("auto", everything)).toEqual(["wasm"]); + }); + + test("explicit WebGPU is tried (then WASM) only with shader-f16", () => { + expect(parakeetDeviceAttempts("webgpu", everything)).toEqual(["webgpu", "wasm"]); + expect(parakeetDeviceAttempts("webgpu", { ...everything, webgpuF16: false })).toEqual([ + "wasm", + ]); + expect(parakeetDeviceAttempts("webgpu", { ...everything, webgpu: false })).toEqual([ + "wasm", + ]); + }); + + test("WebNN and WASM preferences run WASM", () => { + expect(parakeetDeviceAttempts("webnn", everything)).toEqual(["wasm"]); + expect(parakeetDeviceAttempts("wasm", everything)).toEqual(["wasm"]); + }); +}); + describe("TranscribeProgressTracker", () => { // Production shape: 30 s windows advancing by 20 s (30 − 2·5, the worker's // chunk/stride settings), like the chunked ASR pipeline. diff --git a/tests/workerBundle.test.ts b/tests/workerBundle.test.ts new file mode 100644 index 0000000..eda8384 --- /dev/null +++ b/tests/workerBundle.test.ts @@ -0,0 +1,31 @@ +import { describe, expect, test } from "bun:test"; +import { join } from "node:path"; +import { ortAliasPlugin } from "../scripts/ortAlias"; + +// D-WS7a: the Parakeet path must not bring a second onnxruntime into the +// inference worker. parakeet.js depends on its own onnxruntime-web (1.24.1); +// the app runs the one Transformers.js pins. Unminified bundles carry a +// `// node_modules/` comment per bundled module, so count runtimes. + +const WORKER = join(import.meta.dir, "../src/worker/inference.worker.ts"); + +async function bundledRuntimes(plugins: Bun.BunPlugin[]): Promise { + const build = await Bun.build({ entrypoints: [WORKER], target: "browser", format: "esm", plugins }); + expect(build.success).toBe(true); + const code = await build.outputs[0]!.text(); + return [...code.matchAll(/^\/\/ (node_modules\/\S*onnxruntime-web\/dist\/\S+)$/gm)].map((m) => m[1]!); +} + +describe("inference worker bundle", () => { + test("with the alias plugin it holds only the Transformers.js onnxruntime", async () => { + expect(await bundledRuntimes([ortAliasPlugin])).toEqual([ + "node_modules/onnxruntime-web/dist/ort.webgpu.bundle.min.mjs", + ]); + }); + + test("without it, parakeet.js's own runtime would be bundled too (the plugin is load-bearing)", async () => { + const runtimes = await bundledRuntimes([]); + expect(runtimes).toHaveLength(2); + expect(runtimes.some((path) => path.includes("parakeet.js/node_modules/"))).toBe(true); + }); +});