753a4b2960
What review found:
- A worker let go (unload) finishes a load under way before it releases its
model and device, but it counted as gone at once: the listener could load
beside it (two models at once on an iPad), and a new download's check
could start a second copy. abbrev.js now keeps every worker let go until
it says it has given its memory back ("stopping"); gpu_busy() counts it;
a new worker starts only after it; make_room() waits for it before the
helper loads; and the listener (and first letters, for the helper and the
listener) lets the other go only once a load under way has finished.
- The listener's and first letters' first-run checks each waited for the
other's CHECKING: two downloads finishing together waited for ever. A
check now waits only for a check that holds the chip (other_check, the
rule the helper already followed); a request still says the other is busy.
- The worker watches device.lost and tells the page, which lets it go; any
error but "superseded" lets the worker go; the first-run check lets the
model go whatever its result.
- Removing one feature no longer deletes the shared runtime while the other
is downloading it (each counts the other's download as using it).
- What is typed is folded as a keyboard types it (NFKC in the browser,
curly quotes straight), in letters() and in the worker.
- Credits carry a link to where each piece is published.
Co-Authored-By: Claude Opus 5.5 <[email protected]>
568 lines
25 KiB
JavaScript
568 lines
25 KiB
JavaScript
// The first-letters worker for the AAC reference app (abbreviation
|
||
// expansion): the person types only the first letter of each word ("iwtgh"),
|
||
// and a small language model taught to do this proposes whole sentences ("I
|
||
// want to go home"), fitting what their partner said when that is given.
|
||
// Qwen3 0.6B, fine-tuned for it and stored in 4 bits, runs on onnxruntime-
|
||
// web's WebGPU execution provider; the sentences come from a beam search that
|
||
// keeps only words beginning with the letters typed, in order (the same
|
||
// algorithm as the model's reference decoder, ae_decode.py, and the
|
||
// prototype's ae-engine.mjs, both in the model library: tools/ae/README.md).
|
||
//
|
||
// Nothing leaves the device, and nothing is kept: what is typed and what the
|
||
// partner said come in one message and go once the sentences are found. The
|
||
// worker never fetches: the runtime (onnxruntime-web's bundle and its
|
||
// WebAssembly), the tokenizer library (Hugging Face's @huggingface/
|
||
// tokenizers, beside this file) and the model's files come from the model
|
||
// store (Cache API "aac-models-v1", each file SHA-256 checked when it was
|
||
// stored), so it works offline.
|
||
//
|
||
// Graphics chip only. The session is created with the processor fallback
|
||
// switched off (session.disable_cpu_ep_fallback): a device that cannot run
|
||
// the whole graph on its graphics chip fails to load it, rather than running
|
||
// it on the processor. All 371 nodes run on WebGPU (checked in Chrome and in
|
||
// WebKit).
|
||
//
|
||
// The worker has a GPU device of its own (every limit and feature the
|
||
// adapter offers, as the runtime's own would), given to the runtime for the
|
||
// session. Unloading releases the session and destroys the device, which
|
||
// gives all of its graphics memory back; only then may the page end the
|
||
// worker. (In WebKit, ending a worker whose WebGPU device the runtime's JSPI
|
||
// build still holds ends the page with it: seen with Playwright's WebKit,
|
||
// 2026-10-01.)
|
||
//
|
||
// A module worker, started by src/bindings/abbrev.js from a Blob URL. It
|
||
// imports onnxruntime-web's JSPI build (shared with the listener: it is
|
||
// offered only in browsers with JavaScript Promise Integration) and the
|
||
// tokenizer library from Blob URLs too, and hands the runtime its
|
||
// WebAssembly's bytes, so nothing is looked for by name. JSPI lets the runtime
|
||
// read the weights one tensor at a time, straight into graphics memory, so
|
||
// the model is never held whole in memory. One thread.
|
||
//
|
||
// Messages in (each with an `id`, answered in order with {id, ok, ...} or
|
||
// {id, ok: false, error}; and, unasked, {event: "lost", text} when the
|
||
// graphics chip takes the device back):
|
||
// {type: "init", bundle, wasm, tokenizers} model-store URLs of the runtime
|
||
// and the tokenizer library
|
||
// {type: "load", key, model} model: {graph, data: [{name,
|
||
// parts: [{url, bytes}]}],
|
||
// tokenizer, tokenizer_config,
|
||
// config, vocab, vocab_table}
|
||
// (URLs); loads it, and finds
|
||
// one set of sentences so the
|
||
// graphics programs are compiled
|
||
// {type: "check"} one set of sentences, which
|
||
// must not be empty
|
||
// {type: "expand", letters, partner, near_keys}
|
||
// -> {sentences: [text, ...], ms}
|
||
// {type: "cancel"} the expansion running now, and
|
||
// any asked for before, end
|
||
// "superseded"
|
||
// {type: "unload"} releases the model and destroys
|
||
// the device; the worker can then
|
||
// be ended, or load again
|
||
// An expansion asked for after another (or an unload) supersedes it: the
|
||
// earlier one stops at its next model run and ends "superseded".
|
||
//
|
||
// Vendored into vendor/hf-tokenizers-<version>-<revision>/ as ae-worker.js by
|
||
// vendor.sh, unchanged. Its bytes are pinned in crates/aac-web/src/abbrev.rs
|
||
// (RUNTIME_FILES): change this file only together with a new
|
||
// TOKENIZERS_VENDOR_REV and its new size and hash there.
|
||
|
||
const CACHE_NAME = "aac-models-v1";
|
||
|
||
// The search, as tuned on the development set (tools/ae/README.md): eight
|
||
// hypotheses, five sentences offered.
|
||
const BEAM = 8;
|
||
const OFFERED = 5;
|
||
// A slip to a key next to the one meant costs this much (natural log of the
|
||
// probability), when the person asks for slips to be forgiven.
|
||
const NEAR_KEY_COST = 1.0;
|
||
|
||
let ort = null;
|
||
let Tokenizer = null;
|
||
let engine = null; // the loaded model
|
||
let key = "";
|
||
// The latest expansion (or cancel) asked for: any other is superseded.
|
||
let latest = 0;
|
||
|
||
function describe(err) {
|
||
return String(err && err.message ? err.message : err);
|
||
}
|
||
|
||
async function stored(cache, url) {
|
||
const res = await cache.match(url);
|
||
if (!res) throw new Error(`first letters is not on this device (${url.split("/").pop()})`);
|
||
return res;
|
||
}
|
||
|
||
async function importStored(cache, url) {
|
||
const code = await (await stored(cache, url)).blob();
|
||
const blobUrl = URL.createObjectURL(new Blob([code], { type: "text/javascript" }));
|
||
try {
|
||
return await import(blobUrl);
|
||
} finally {
|
||
URL.revokeObjectURL(blobUrl);
|
||
}
|
||
}
|
||
|
||
async function init(msg) {
|
||
if (ort && Tokenizer) return {};
|
||
if (typeof caches === "undefined") throw new Error("this browser keeps no model store in workers");
|
||
const cache = await caches.open(CACHE_NAME);
|
||
const [runtime, tokenizers, wasm] = await Promise.all([
|
||
importStored(cache, msg.bundle),
|
||
importStored(cache, msg.tokenizers),
|
||
stored(cache, msg.wasm).then((r) => r.arrayBuffer()),
|
||
]);
|
||
runtime.env.wasm.wasmBinary = wasm; // never fetched by name
|
||
runtime.env.wasm.numThreads = 1;
|
||
runtime.env.wasm.proxy = false;
|
||
runtime.env.logLevel = "error";
|
||
ort = runtime;
|
||
Tokenizer = tokenizers.Tokenizer;
|
||
return {};
|
||
}
|
||
|
||
// ─── Loading ───────────────────────────────────────────────────────────────
|
||
|
||
// The part of the model the runtime is reading now. It reads the weights one
|
||
// tensor after another, in file order, so it is mostly the same part again:
|
||
// one part (64 MiB at most) is held at a time, never the whole model.
|
||
let held = { url: "", blob: null };
|
||
|
||
function partBlob(cache, url) {
|
||
if (held.url !== url) held = { url, blob: stored(cache, url).then((r) => r.blob()) };
|
||
return held.blob;
|
||
}
|
||
|
||
// A weights file as the runtime reads it: a Blob it slices one tensor at a
|
||
// time, whose bytes stay in the model store until then. A file larger than a
|
||
// part (one tensor, the 74 MiB table of word embeddings, cannot be split) is
|
||
// stored as several parts, read here as one.
|
||
class StoredFile extends Blob {
|
||
constructor(cache, parts) {
|
||
super([]);
|
||
this.cache = cache;
|
||
this.parts = parts;
|
||
this.bytes = parts.reduce((n, p) => n + p.bytes, 0);
|
||
}
|
||
|
||
get size() {
|
||
return this.bytes;
|
||
}
|
||
|
||
slice(start, end) {
|
||
const { cache, parts } = this;
|
||
return {
|
||
stream: () => new ReadableStream({
|
||
async start(controller) {
|
||
let at = 0;
|
||
for (const p of parts) {
|
||
const from = Math.max(start, at);
|
||
const to = Math.min(end, at + p.bytes);
|
||
if (from < to) {
|
||
const blob = await partBlob(cache, p.url);
|
||
controller.enqueue(new Uint8Array(await blob.slice(from - at, to - at).arrayBuffer()));
|
||
}
|
||
at += p.bytes;
|
||
}
|
||
controller.close();
|
||
},
|
||
}),
|
||
};
|
||
}
|
||
}
|
||
|
||
async function load(msg) {
|
||
if (engine && key === msg.key) return {};
|
||
await unload();
|
||
const cache = await caches.open(CACHE_NAME);
|
||
const m = msg.model;
|
||
const json = async (url) => (await stored(cache, url)).json();
|
||
const [graph, tj, tc, cfg, vj, vb] = await Promise.all([
|
||
stored(cache, m.graph).then((r) => r.arrayBuffer()),
|
||
json(m.tokenizer),
|
||
json(m.tokenizer_config),
|
||
json(m.config),
|
||
json(m.vocab),
|
||
stored(cache, m.vocab_table).then((r) => r.arrayBuffer()),
|
||
]);
|
||
const e = new Engine(new Tokenizer(tj, tc), cfg.model.decoder, vj, vb);
|
||
try {
|
||
e.device = await ownDevice();
|
||
// The graphics chip may take the device back (a driver reset, memory
|
||
// short): nothing more runs on it, and the page is told, so it lets
|
||
// this worker go and starts afresh.
|
||
e.device.lost.then((info) => {
|
||
if (info.reason === "destroyed") return; // let go on purpose
|
||
e.lost = true;
|
||
self.postMessage({ event: "lost", text: "the graphics chip stopped it" });
|
||
});
|
||
await e.open(new Uint8Array(graph), m.data.map((d) => ({ path: d.name, data: new StoredFile(cache, d.parts) })));
|
||
} catch (err) {
|
||
await e.release();
|
||
throw err;
|
||
} finally {
|
||
held = { url: "", blob: null };
|
||
}
|
||
engine = e;
|
||
key = msg.key;
|
||
// The first run compiles the graphics chip's programs for the prompt
|
||
// and for the batched steps (a fraction of a second, once).
|
||
await engine.expand("ih", null, { nearKeys: false, live: () => true });
|
||
return {};
|
||
}
|
||
|
||
// A GPU device for this worker alone: every limit and feature the adapter
|
||
// offers (the runtime asks for many of them: storage buffers per shader,
|
||
// workgroup sizes, half precision).
|
||
async function ownDevice() {
|
||
const adapter = navigator.gpu && (await navigator.gpu.requestAdapter());
|
||
if (!adapter) throw new Error("this browser offers no graphics chip to first letters");
|
||
const requiredLimits = {};
|
||
for (const k in adapter.limits) if (typeof adapter.limits[k] === "number") requiredLimits[k] = adapter.limits[k];
|
||
return adapter.requestDevice({ requiredFeatures: [...adapter.features], requiredLimits });
|
||
}
|
||
|
||
async function unload() {
|
||
latest++;
|
||
const old = engine;
|
||
engine = null;
|
||
key = "";
|
||
if (old) await old.release();
|
||
return {};
|
||
}
|
||
|
||
function loaded() {
|
||
if (!engine) throw new Error("first letters is not loaded");
|
||
}
|
||
|
||
async function check() {
|
||
loaded();
|
||
const r = await engine.expand("iwtgh", "What do you want to do?", { nearKeys: false, live: () => true });
|
||
if (r.sentences.length === 0) throw new Error("it found no sentence");
|
||
return {};
|
||
}
|
||
|
||
async function expand(msg) {
|
||
loaded();
|
||
const mine = msg.id;
|
||
const live = () => latest === mine;
|
||
if (!live()) throw new Error("superseded");
|
||
return engine.expand(String(msg.letters || ""), msg.partner ? String(msg.partner) : null, { nearKeys: !!msg.near_keys, live });
|
||
}
|
||
|
||
// ─── The letters ───────────────────────────────────────────────────────────
|
||
|
||
// The keyboard the slips are measured on (QWERTY, the paper's 3 x 10 grid).
|
||
const ROWS = ["qwertyuiop", "asdfghjkl;", "zxcvbnm,./"];
|
||
const POS = new Map();
|
||
ROWS.forEach((row, r) => [...row].forEach((ch, c) => POS.set(ch, [c, r])));
|
||
const isAlnum = (ch) => /[\p{L}\p{N}]/u.test(ch);
|
||
|
||
function normText(s) {
|
||
s = s.normalize("NFKC").replace(/[’‘ʼ]/g, "'").replace(/[“”]/g, '"')
|
||
.replace(/—/g, " - ").replace(/–/g, "-");
|
||
return s.split(/\s+/).filter(Boolean).join(" ");
|
||
}
|
||
|
||
function stripFinal(s) {
|
||
s = s.trimEnd();
|
||
let i = s.length;
|
||
while (i > 0 && !isAlnum(s[i - 1])) i--;
|
||
return s.slice(0, i);
|
||
}
|
||
|
||
const normMatch = (s) => stripFinal(normText(s)).toLowerCase().split(" ").filter(Boolean).join(" ");
|
||
|
||
// The prompt the model was taught with: "Partner: …\nLetters: i w t g h\nReply:".
|
||
function makePrompt(letters, partner) {
|
||
const head = partner ? `Partner: ${normText(partner)}\n` : "";
|
||
return `${head}Letters: ${[...letters].join(" ")}\nReply:`;
|
||
}
|
||
|
||
function neighbours(ch) {
|
||
const out = new Set();
|
||
const p = POS.get(ch);
|
||
if (!p) return out;
|
||
for (let dr = -1; dr <= 1; dr++) {
|
||
for (let dc = -1; dc <= 1; dc++) {
|
||
if (!dr && !dc) continue;
|
||
const r = p[1] + dr;
|
||
const c = p[0] + dc;
|
||
if (r >= 0 && r < 3 && c >= 0 && c < 10 && /[a-z]/.test(ROWS[r][c])) out.add(ROWS[r][c]);
|
||
}
|
||
}
|
||
return out;
|
||
}
|
||
|
||
// Can the text `e` follow a hypothesis that has used `p` of the letters `Q`?
|
||
// -> [letters used after it, cost] or null.
|
||
function advance(e, Q, p, punct, nearCost) {
|
||
let cost = 0;
|
||
for (const ch of e) {
|
||
if (isAlnum(ch)) {
|
||
if (p < Q.length && Q[p] === ch) p++;
|
||
else if (nearCost != null && p < Q.length && neighbours(Q[p]).has(ch)) {
|
||
p++;
|
||
cost += nearCost;
|
||
} else return null;
|
||
} else {
|
||
if (punct === "free") continue;
|
||
if (p < Q.length && Q[p] === ch) p++;
|
||
else if (p >= Q.length) continue;
|
||
else return null;
|
||
}
|
||
}
|
||
return [p, cost];
|
||
}
|
||
|
||
let F16 = null;
|
||
function f16table() {
|
||
if (F16) return F16;
|
||
F16 = new Float32Array(65536);
|
||
for (let h = 0; h < 65536; h++) {
|
||
const s = h & 0x8000 ? -1 : 1;
|
||
const e = (h >> 10) & 31;
|
||
const m = h & 1023;
|
||
F16[h] = e === 0 ? s * m * 2 ** -24 : e === 31 ? (m ? NaN : s * Infinity) : s * (1 + m / 1024) * 2 ** (e - 15);
|
||
}
|
||
return F16;
|
||
}
|
||
|
||
// ─── The model ─────────────────────────────────────────────────────────────
|
||
|
||
// The graph's contract (the model library's onnx_patch.py): inputs input_ids
|
||
// int64 [B,S], seqlens_k int32 [B], total_sequence_length int32 [],
|
||
// beam_idx int32 [B], past_key_values.N.{key,value} fp16 [B_prev,kvh,P,hd];
|
||
// outputs logits fp16 [B,S,V], logsumexp f32 [B,S], present.N.* (kept on the
|
||
// graphics chip, and reordered there by beam_idx at the next step).
|
||
class Engine {
|
||
constructor(tokenizer, dec, vj, vb) {
|
||
this.tok = tokenizer;
|
||
this.kvh = dec.num_key_value_heads;
|
||
this.hd = dec.head_size;
|
||
this.L = dec.num_hidden_layers;
|
||
this.V = vj.V;
|
||
this.eos = vj.eos[0];
|
||
// Which tokens may come next, by what they start with: one table for
|
||
// after a letter, one for after anything else.
|
||
const gN = new Uint16Array(vb, 0, this.V);
|
||
const gL = new Uint16Array(vb, 2 * this.V, this.V);
|
||
this.last = new Int8Array(vb, 4 * this.V, this.V);
|
||
const build = (g) => {
|
||
const counts = new Map();
|
||
for (let i = 0; i < this.V; i++) if (g[i] !== 0xffff) counts.set(g[i], (counts.get(g[i]) || 0) + 1);
|
||
const arrs = new Map();
|
||
const fill = new Map();
|
||
for (const [k, n] of counts) {
|
||
arrs.set(k, new Int32Array(n));
|
||
fill.set(k, 0);
|
||
}
|
||
for (let i = 0; i < this.V; i++) {
|
||
if (g[i] === 0xffff) continue;
|
||
const k = g[i];
|
||
const f = fill.get(k);
|
||
arrs.get(k)[f] = i;
|
||
fill.set(k, f + 1);
|
||
}
|
||
const byFirst = new Map();
|
||
for (const k of arrs.keys()) {
|
||
const h = vj.groups[k].slice(0, 1);
|
||
if (!byFirst.has(h)) byFirst.set(h, []);
|
||
byFirst.get(h).push(k);
|
||
}
|
||
return { arrs, byFirst };
|
||
};
|
||
this.tables = { N: build(gN), L: build(gL) };
|
||
this.groups = vj.groups;
|
||
const io = dec.outputs || {};
|
||
const kn = io.present_key_names || "present.%d.key";
|
||
const vn = io.present_value_names || "present.%d.value";
|
||
this.pres = [];
|
||
for (let i = 0; i < this.L; i++) this.pres.push(kn.replace("%d", i), vn.replace("%d", i));
|
||
this.pasts = this.pres.map((n) => n.replace("present", "past_key_values"));
|
||
this.sess = null;
|
||
this.device = null;
|
||
}
|
||
|
||
async open(graph, externalData) {
|
||
this.sess = await ort.InferenceSession.create(graph, {
|
||
// Options reach the WebGPU provider only when it is named as an
|
||
// object; the device is this worker's own.
|
||
executionProviders: [{ name: "webgpu", device: this.device }],
|
||
externalData,
|
||
graphOptimizationLevel: "all",
|
||
logSeverityLevel: 3,
|
||
// The caches stay on the graphics chip from one step to the next.
|
||
preferredOutputLocation: Object.fromEntries(this.pres.map((n) => [n, "gpu-buffer"])),
|
||
// Graphics chip or nothing: no node may fall back to the processor.
|
||
extra: { session: { disable_cpu_ep_fallback: "1" } },
|
||
});
|
||
}
|
||
|
||
// The session, then the device: all of its graphics memory goes back.
|
||
async release() {
|
||
if (this.sess) await this.sess.release().catch(() => {});
|
||
this.sess = null;
|
||
if (this.device) this.device.destroy();
|
||
this.device = null;
|
||
}
|
||
|
||
async run(rows, parents, state) {
|
||
const S = rows[0].length;
|
||
const B = rows.length;
|
||
state.T += S;
|
||
const flat = new BigInt64Array(B * S);
|
||
rows.forEach((r, i) => r.forEach((t, j) => { flat[i * S + j] = BigInt(t); }));
|
||
const feeds = {
|
||
input_ids: new ort.Tensor("int64", flat, [B, S]),
|
||
seqlens_k: new ort.Tensor("int32", new Int32Array(B).fill(state.T - 1), [B]),
|
||
total_sequence_length: new ort.Tensor("int32", Int32Array.of(state.T), []),
|
||
beam_idx: new ort.Tensor("int32", Int32Array.from(parents), [B]),
|
||
...state.past,
|
||
};
|
||
const o = await this.sess.run(feeds);
|
||
for (const t of Object.values(state.past)) t.dispose?.();
|
||
state.past = {};
|
||
this.pres.forEach((n, k) => { state.past[this.pasts[k]] = o[n]; });
|
||
const Vo = o.logits.dims[2];
|
||
const So = o.logits.dims[1];
|
||
const raw = o.logits.data;
|
||
const u16 = raw instanceof Uint16Array ? raw : new Uint16Array(raw.buffer, raw.byteOffset, raw.length);
|
||
return { u16, lse: o.logsumexp.data, Vo, So, B };
|
||
}
|
||
|
||
// Up to OFFERED sentences whose words begin with `typed`'s letters, in
|
||
// order (punctuation typed must be where it was typed; none typed, and
|
||
// it may go anywhere). `live()` false stops it ("superseded").
|
||
async expand(typed, partner, { nearKeys, live }) {
|
||
const t0 = performance.now();
|
||
if (this.lost) throw new Error("the graphics chip stopped it");
|
||
// What a keyboard may type for the same letter or mark: compatibility
|
||
// forms folded (NFKC), curly quotes straight.
|
||
typed = typed.normalize("NFKC").replace(/[\u2019\u2018\u02bc]/g, "'").replace(/[\u201c\u201d]/g, '"');
|
||
let Q = typed.toLowerCase().replace(/\s+/g, "");
|
||
const punct = [...Q].some((c) => !isAlnum(c)) ? "strict" : "free";
|
||
if (punct === "free") Q = [...Q].filter(isAlnum).join("");
|
||
if (!Q) return { sentences: [], ms: 0 };
|
||
const nearCost = nearKeys ? NEAR_KEY_COST : null;
|
||
const maxSteps = Math.min(48, 6 + 3 * Q.length);
|
||
const ids = this.tok.encode(makePrompt(typed.toLowerCase().replace(/\s+/g, ""), partner), { add_special_tokens: false }).ids;
|
||
const state = { T: 0, past: {} };
|
||
for (const n of this.pasts) state.past[n] = new ort.Tensor("float16", new Uint16Array(0), [1, this.kvh, 0, this.hd]);
|
||
const f16 = f16table();
|
||
const fin = new Map();
|
||
try {
|
||
let out = await this.run([ids], [0], state);
|
||
let beams = [{ score: 0, ids: [], p: 0, prev: 0 }];
|
||
for (let step = 0; step < maxSteps; step++) {
|
||
if (!live()) throw new Error("superseded");
|
||
const cands = [];
|
||
for (let bi = 0; bi < beams.length; bi++) {
|
||
const b = beams[bi];
|
||
const off = (bi * out.So + (out.So - 1)) * out.Vo;
|
||
const lse = out.lse[bi * out.So + out.So - 1];
|
||
if (b.p === Q.length && b.ids.length) {
|
||
const k = b.ids.join(",");
|
||
if (!fin.has(k)) fin.set(k, b.score + f16[out.u16[off + this.eos]] - lse);
|
||
}
|
||
const T = b.prev === 1 ? this.tables.L : this.tables.N;
|
||
const keys = [...(T.byFirst.get("") || [])];
|
||
if (b.p < Q.length) {
|
||
const heads = new Set([Q[b.p]]);
|
||
if (nearCost != null) for (const n of neighbours(Q[b.p])) heads.add(n);
|
||
for (const h of heads) keys.push(...(T.byFirst.get(h) || []));
|
||
}
|
||
if (punct === "free" || b.p >= Q.length) {
|
||
const qh = b.p < Q.length ? Q[b.p] : "";
|
||
for (const [h, ks] of T.byFirst) if (h && !isAlnum(h) && h !== qh) keys.push(...ks);
|
||
}
|
||
// The best BEAM continuations of this hypothesis.
|
||
const top = [];
|
||
let minI = -1;
|
||
let minV = Infinity;
|
||
for (const g of keys) {
|
||
const r = advance(this.groups[g], Q, b.p, punct, nearCost);
|
||
if (!r) continue;
|
||
const arr = T.arrs.get(g);
|
||
const base = b.score - r[1] - lse;
|
||
for (let j = 0; j < arr.length; j++) {
|
||
const t = arr[j];
|
||
const s = base + f16[out.u16[off + t]];
|
||
if (top.length < BEAM) {
|
||
top.push([s, t, r[0]]);
|
||
if (top.length === BEAM) {
|
||
minV = Infinity;
|
||
for (let q = 0; q < BEAM; q++) if (top[q][0] < minV) { minV = top[q][0]; minI = q; }
|
||
}
|
||
} else if (s > minV) {
|
||
top[minI] = [s, t, r[0]];
|
||
minV = Infinity;
|
||
for (let q = 0; q < BEAM; q++) if (top[q][0] < minV) { minV = top[q][0]; minI = q; }
|
||
}
|
||
}
|
||
}
|
||
for (const [s, t, np] of top) cands.push([s, bi, t, np]);
|
||
}
|
||
if (!cands.length) break;
|
||
cands.sort((a, b) => b[0] - a[0]);
|
||
cands.length = Math.min(cands.length, BEAM);
|
||
if (fin.size >= OFFERED) {
|
||
const kth = [...fin.values()].sort((a, b) => b - a)[OFFERED - 1];
|
||
if (cands[0][0] < kth) break;
|
||
}
|
||
beams = cands.map(([s, bi, t, np]) => ({ score: s, ids: [...beams[bi].ids, t], p: np, prev: this.last[t] }));
|
||
out = await this.run(cands.map((c) => [c[2]]), cands.map((c) => c[1]), state);
|
||
}
|
||
} finally {
|
||
for (const t of Object.values(state.past)) t.dispose?.();
|
||
}
|
||
// Hypotheses that read the same (case, spacing, final punctuation
|
||
// aside) are one sentence, with their probabilities added.
|
||
const merged = new Map();
|
||
for (const [k, s] of fin) {
|
||
const text = this.tok.decode(k.split(",").map(Number), { skip_special_tokens: true }).trim();
|
||
const nk = normMatch(text);
|
||
if (!nk) continue;
|
||
const m = merged.get(nk);
|
||
if (m) {
|
||
const hi = Math.max(m.lse, s);
|
||
const lo = Math.min(m.lse, s);
|
||
m.lse = hi + Math.log1p(Math.exp(lo - hi));
|
||
if (s > m.best) {
|
||
m.best = s;
|
||
m.text = text;
|
||
}
|
||
} else merged.set(nk, { lse: s, best: s, text });
|
||
}
|
||
const sentences = [...merged.values()].sort((a, b) => b.lse - a.lse).slice(0, OFFERED).map((m) => m.text);
|
||
return { sentences, ms: Math.round(performance.now() - t0) };
|
||
}
|
||
}
|
||
|
||
// ─── Messages ──────────────────────────────────────────────────────────────
|
||
|
||
const HANDLERS = { init, load, check, expand, unload, cancel: async () => ({}) };
|
||
|
||
// One thing at a time, in the order it came; an expansion (or a cancel, or
|
||
// an unload) supersedes those before it as soon as it arrives.
|
||
let queue = Promise.resolve();
|
||
|
||
self.onmessage = (event) => {
|
||
const msg = event.data || {};
|
||
if (msg.type === "expand" || msg.type === "cancel" || msg.type === "unload") latest = msg.id;
|
||
const handler = HANDLERS[msg.type];
|
||
queue = queue.then(async () => {
|
||
if (!handler) {
|
||
self.postMessage({ id: msg.id, ok: false, error: `unknown message ${msg.type}` });
|
||
return;
|
||
}
|
||
try {
|
||
const reply = await handler(msg);
|
||
self.postMessage({ id: msg.id, ok: true, ...reply });
|
||
} catch (e) {
|
||
self.postMessage({ id: msg.id, ok: false, error: describe(e) });
|
||
}
|
||
});
|
||
};
|