Hugging Face's logo

Spaces:
SiliconLanguage
/
open-asr-model-explorer
Sleeping

App Files Community
open-asr-model-explorer / frontend /src /workers /webgpu.worker.js
Ping Long
feat(ui): stabilize metrics state and expand language samples
d4f1e7c
Raw History Blame Contribute Delete
12.1 kB
/**
* webgpu.worker.js
* ─────────────────────────────────────────────────────────────────────────────
* Runs entirely inside a WebWorker context.
* Responsibilities:
* 1. Load a transformers.js Whisper pipeline (WebGPU backend).
* 2. Cache ONNX model weights in the browser via the Cache API for fast
* subsequent cold-starts.
* 3. Implement a "LocalAgreement-2" style streaming stabilisation algorithm
* that withholds partial hypotheses until at least two consecutive
* decoding passes agree on a prefix, preventing jittery live output.
*
* Message protocol (worker ← main thread):
* { type: 'load', model: string }
* { type: 'transcribe', audio: Float32Array, id: string }
* { type: 'abort', id: string }
*
* Message protocol (worker → main thread):
* { type: 'ready' }
* { type: 'progress', id, text: string, stable: bool }
* { type: 'result', id, transcript: string, ttft_ms, itl_ms, rtfx }
* { type: 'error', id, message: string }
*/
import { pipeline, env } from '@huggingface/transformers';
// ── Cache configuration ───────────────────────────────────────────────────────
// Point transformers.js at the Cache Storage API so ONNX weights survive
// page reloads without re-downloading from the CDN.
// useBrowserCache relies on the Cache Storage API which is only available in
// secure contexts (HTTPS or localhost); it is a no-op if the API is absent.
env.useBrowserCache = typeof caches !== 'undefined';
env.allowLocalModels = false;
// ── State ─────────────────────────────────────────────────────────────────────
/** @type {import('@huggingface/transformers').Pipeline | null} */
let asr = null;
let loadedModelId = null;
const abortedIds = new Set();
const LANGUAGE_ISO_MAP = {
english: 'en',
chinese: 'zh',
spanish: 'es',
french: 'fr',
german: 'de',
italian: 'it',
japanese: 'ja',
hindi: 'hi',
};
// ── LocalAgreement-2 helpers ──────────────────────────────────────────────────
/**
* Returns the longest common prefix (token-level) shared by *a* and *b*.
* @param {number[]} a
* @param {number[]} b
* @returns {number[]}
*/
function longestCommonPrefix(a, b) {
const len = Math.min(a.length, b.length);
let i = 0;
while (i < len && a[i] === b[i]) i++;
return a.slice(0, i);
}
/**
* LocalAgreement-2 stabiliser.
*
* Keeps a sliding window of the last two hypotheses. Only the prefix
* agreed upon by *both* is emitted as "stable" text. The remainder is
* surfaced as an unstable suffix so the UI can render it in a dimmed style.
*/
class LocalAgreement2 {
constructor() {
this.prev = /** @type {number[]} */ ([]);
this.stableTokenIds = /** @type {number[]} */ ([]);
}
/**
* Feed in the latest hypothesis token IDs.
* @param {number[]} current Current hypothesis token IDs
* @returns {{ stableTokenIds: number[], unstableTokenIds: number[] }}
*/
update(current) {
const agreed = longestCommonPrefix(this.prev, current);
// Extend the stable set with newly confirmed tokens, but never claim more
// stable tokens than actually exist in the current hypothesis. Capping at
// current.length prevents stale over-commitment when the model produces a
// shorter hypothesis (e.g. at the start of a new audio chunk), which would
// otherwise show confirmed text that is no longer in the active hypothesis.
const newStableLen = Math.min(
Math.max(agreed.length, this.stableTokenIds.length),
current.length,
);
this.stableTokenIds = current.slice(0, newStableLen);
this.prev = current;
const stableTokenIds = this.stableTokenIds;
const unstableTokenIds = current.slice(this.stableTokenIds.length);
return { stableTokenIds, unstableTokenIds };
}
reset() {
this.prev = [];
this.stableTokenIds = [];
}
}
// ── Model loading ─────────────────────────────────────────────────────────────
async function loadModel(modelId) {
if (asr && loadedModelId === modelId) {
self.postMessage({ type: 'ready', modelId });
return;
}
try {
// Dispose of the previous pipeline's ONNX sessions before loading a new
// model. Without this, ONNX Runtime tries to .destroy() GPU buffers that
// were already garbage-collected, throwing "Cannot read properties of
// undefined (reading 'destroy')".
if (asr) {
try { await asr.dispose(); } catch { /* best-effort cleanup */ }
asr = null;
loadedModelId = null;
}
self.postMessage({ type: 'loading', modelId });
// Cohere 2B encoder fp16 has 2 external data shards (~2.5 GB) which
// exceeds browser memory limits. Use q4f16 (1 shard, ~600 MB) instead.
const isCohere = modelId.toLowerCase().includes('cohere');
const gpuDtype = isCohere ? 'q4f16' : 'fp16';
try {
asr = await pipeline('automatic-speech-recognition', modelId, {
device: 'webgpu',
dtype: gpuDtype,
});
} catch {
// WebGPU unavailable (non-secure context, old browser, no GPU) –
// fall back to WASM fp32 so the app still works.
self.postMessage({ type: 'loading', modelId, fallback: 'wasm' });
asr = await pipeline('automatic-speech-recognition', modelId, {
device: 'wasm',
dtype: isCohere ? 'q4' : 'fp32',
});
}
loadedModelId = modelId;
self.postMessage({ type: 'ready', modelId });
} catch (err) {
// Reset state so a subsequent load attempt can retry cleanly.
asr = null;
loadedModelId = null;
self.postMessage({ type: 'error', id: null, message: String(err) });
}
}
// ── Transcription ─────────────────────────────────────────────────────────────
async function transcribe(audio, id, language) {
if (!asr) {
self.postMessage({ type: 'error', id, message: 'Model not loaded. Send a load message first.' });
return;
}
const isoLang = language ? (LANGUAGE_ISO_MAP[language.toLowerCase()] || 'en') : undefined;
const la2 = new LocalAgreement2();
const t0 = performance.now();
let firstTokenTime = null;
const tokenTimestamps = [];
try {
// transformers.js callback_function fires after each generated token
// Whisper models support chunked long-form decoding; Cohere uses its own
// encoder and these params cause early truncation.
const isWhisper = (loadedModelId ?? '').toLowerCase().includes('whisper');
// Non-Whisper models (e.g. Cohere ONNX) default to max_length=20 in
// GenerationConfig, truncating transcripts to ~10 new tokens. Force a
// generous limit via BOTH max_new_tokens AND generation_config to ensure
// at least one path overrides the default.
const cohereGenerationOverrides = !isWhisper
? { max_new_tokens: 512, max_length: 1024, generation_config: { max_new_tokens: 512, max_length: 1024 } }
: {};
const result = await asr(audio, {
task: 'transcribe',
...(isoLang ? { language: isoLang } : {}),
// Whisper: chunked long-form decoding with timestamp grounding.
// return_timestamps forces the attention mechanism to anchor to audio
// frames, naturally cutting off silence loops without draconian n-gram
// penalties that truncate valid speech.
...(isWhisper ? {
return_timestamps: true,
chunk_length_s: 30,
stride_length_s: 5,
repetition_penalty: 1.05,
condition_on_prev_tokens: false,
} : {}),
...cohereGenerationOverrides,
callback_function: (beams) => {
if (abortedIds.has(id)) throw new Error('ABORTED');
const now = performance.now();
if (firstTokenTime === null) firstTokenTime = now;
tokenTimestamps.push(now);
// Wrap in try/catch — the WhisperTextStreamer in transformers.js can
// throw "token_ids must be a non-empty array of integers" when
// timestamp tokens are emitted. Swallowing the error here lets the
// pipeline complete; the final bulk result is still accurate.
try {
const beamText = beams[0]?.text ?? '';
/** @type {number[]} */
let tokenIds = [];
if (asr.tokenizer?.encode) {
const encoded = asr.tokenizer.encode(beamText);
if (Array.isArray(encoded)) {
tokenIds = encoded.filter((v) => Number.isFinite(v));
} else if (Array.isArray(encoded?.ids)) {
tokenIds = encoded.ids.filter((v) => Number.isFinite(v));
} else if (Array.isArray(encoded?.input_ids)) {
tokenIds = encoded.input_ids.filter((v) => Number.isFinite(v));
} else if (Array.isArray(encoded?.[0])) {
tokenIds = encoded[0].filter((v) => Number.isFinite(v));
}
}
// Fallback to beam token IDs when encode() is unavailable or returns empty.
if (tokenIds.length === 0) {
tokenIds = (beams[0]?.output_token_ids ?? []).filter((v) => Number.isFinite(v));
}
if (tokenIds.length === 0) return; // nothing to consensus on yet
const { stableTokenIds, unstableTokenIds } = la2.update(tokenIds);
const stable = asr.tokenizer?.decode
? asr.tokenizer.decode(stableTokenIds, { skip_special_tokens: true })
: '';
const unstable = asr.tokenizer?.decode
? asr.tokenizer.decode(unstableTokenIds, { skip_special_tokens: true })
: '';
self.postMessage({ type: 'progress', id, stable, unstable });
} catch (cbErr) {
// Silently skip this streaming tick; the final result is unaffected.
}
},
});
if (abortedIds.has(id)) {
abortedIds.delete(id);
return;
}
const t1 = performance.now();
const totalMs = t1 - t0;
const audioDurationS = audio.length / 16000;
const ttftMs = firstTokenTime !== null ? firstTokenTime - t0 : 0;
let itlMs = 0;
if (tokenTimestamps.length > 1) {
const gaps = tokenTimestamps
.slice(1)
.map((ts, i) => ts - tokenTimestamps[i]);
itlMs = gaps.reduce((s, g) => s + g, 0) / gaps.length;
}
const rtfx = audioDurationS / (totalMs / 1000);
self.postMessage({
type: 'result',
id,
transcript: typeof result.text === 'string' ? result.text : result.text ?? '',
ttft_ms: Math.round(ttftMs * 100) / 100,
itl_ms: Math.round(itlMs * 100) / 100,
rtfx: Math.round(rtfx * 10000) / 10000,
});
} catch (err) {
if (String(err).includes('ABORTED')) {
abortedIds.delete(id);
} else {
self.postMessage({ type: 'error', id, message: String(err) });
}
}
}
// ── Message handler ───────────────────────────────────────────────────────────
self.addEventListener('message', async (event) => {
const { type, model, audio, id, language } = event.data;
switch (type) {
case 'load':
await loadModel(model);
break;
case 'transcribe':
await transcribe(audio, id, language);
break;
case 'abort':
abortedIds.add(id);
break;
default:
self.postMessage({ type: 'error', id: null, message: `Unknown message type: ${type}` });
}
});