jev-web 0.3.2 → 0.3.3

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/src/laya.mjs CHANGED
@@ -1,625 +1,625 @@
1
- /**
2
- * laya.mjs — Laya System-1 decision models (ModernBERT encoder + typed head).
3
- *
4
- * The Laya family is a *different graph* from open-jev:
5
- * encoder: input_ids, attention_mask -> hidden (last_hidden_state)
6
- * head: hidden, marker_pos, marker_mask, qtype -> logits, act_logits
7
- *
8
- * Sequence layout (mirrors the reference `laya` Python runtime,
9
- * `laya/common.py build_sequence`):
10
- * [CLS] <type> question: <instructions> [SEP] [MASK]opt0 [MASK]opt1 … [SEP] <state> [SEP]
11
- * `marker_pos` points at each option's [MASK]; `qtype` is 0 choice / 1 score /
12
- * 2 noul. Logits are calibrated per question type *and* option-count bucket
13
- * (rl_agent_config.json `temperature` / `temperature_by_options`), clamped to
14
- * [0.5, 5] like the Python runtime.
15
- *
16
- * The browser export used here is a two-file q8 ONNX split
17
- * (alfred361/laya-typed-decisions-web-q8, pinned by revision).
18
- */
19
-
20
- import { normalizeQuestions } from "./questions.mjs";
21
- import { softmaxWithTemperature } from "./answers.mjs";
22
-
23
- export const LAYA_DEFAULT_MODEL = "alfred361/laya-typed-decisions-web-q8";
24
- export const LAYA_DEFAULT_REVISION = "96f637c0bee046dce5c61a3174de9f7cdcf09ace";
25
- export const LAYA_DEFAULT_SUBFOLDER = "v1";
26
-
27
- /** Reference calibration from the typed-decisions rl_agent_config.json. */
28
- export const LAYA_DEFAULT_TEMPERATURES = [1.0148024559020996, 1.0374259948730469, 1.0575125217437744];
29
- export const LAYA_DEFAULT_TEMPERATURES_BY_OPTIONS = {
30
- "choice:3-5": 1.7601518630981445,
31
- "choice:6-10": 1.0000158548355103,
32
- "score:3-5": 1.2514300346374512,
33
- "noul:2": 1.983399510383606,
34
- "choice:11+": 0.10058280825614929,
35
- "choice:2": 1.9063563346862793,
36
- };
37
- export const LAYA_DEFAULT_MAX_LEN = 1024;
38
- export const LAYA_DEFAULT_HEAD_MAX_LEN = 256;
39
- export const LAYA_DEFAULT_MAX_PREFIXES = 6;
40
-
41
- export const LAYA_QTYPE = { choice: 0, score: 1, noul: 2 };
42
- const QTYPE_NAMES = { 0: "choice", 1: "score", 2: "noul" };
43
-
44
- // The Python runtime refuses temperatures outside [0.5, 5]: the shipped
45
- // `choice:11+` bucket (0.1006) would sharpen logits ~10x and publish a coin
46
- // flip as near-certainty.
47
- export const LAYA_TEMP_MIN = 0.5;
48
- export const LAYA_TEMP_MAX = 5;
49
-
50
- export function clampTemperature(t, lo = LAYA_TEMP_MIN, hi = LAYA_TEMP_MAX) {
51
- const n = Number(t);
52
- if (!Number.isFinite(n)) return 1;
53
- return Math.min(hi, Math.max(lo, n));
54
- }
55
-
56
- export function layaTempBucket(qtype, k) {
57
- const size = k <= 2 ? "2" : k <= 5 ? "3-5" : k <= 10 ? "6-10" : "11+";
58
- return `${QTYPE_NAMES[qtype]}:${size}`;
59
- }
60
-
61
- /**
62
- * Option texts exactly as the reference runtime renders them.
63
- * Our normalized questions carry no per-option criteria, so choice options
64
- * render as their labels and noul falls back to the reference wording.
65
- */
66
- export function renderLayaOptions(question) {
67
- if (question.type === "choice") return question.options.map(String);
68
- if (question.type === "score") return question.options.map((c, i) => `level ${i}: ${c}`);
69
- // noul
70
- const [no, yes] = question.options;
71
- if (no === "no" && yes === "yes") {
72
- return ["false: no, the statement does not hold", "true: yes, the statement holds"];
73
- }
74
- return [`false: ${no}`, `true: ${yes}`];
75
- }
76
-
77
- /**
78
- * Build one Laya sequence + marker positions.
79
- * @param {{encode:(text:string)=>number[], ids:{cls:number,sep:number,mask:number},
80
- * state:string, question:object, maxLen?:number, headMaxLen?:number,
81
- * maxPrefixes?:number}} spec
82
- * @returns {{ids:number[], markers:number[], qtype:number}}
83
- */
84
- export function buildLayaSequence({
85
- encode,
86
- ids,
87
- state,
88
- question,
89
- maxLen = LAYA_DEFAULT_MAX_LEN,
90
- headMaxLen = LAYA_DEFAULT_HEAD_MAX_LEN,
91
- maxPrefixes = LAYA_DEFAULT_MAX_PREFIXES,
92
- }) {
93
- const options = renderLayaOptions(question);
94
- if (options.length > maxPrefixes) {
95
- throw new RangeError(
96
- `laya: ${question.type} with ${options.length} options exceeds max_prefixes=${maxPrefixes}`);
97
- }
98
- const qtype = LAYA_QTYPE[question.type];
99
- if (qtype === undefined) throw new TypeError(`laya: unsupported question type ${question.type}`);
100
-
101
- const strip = (s) => String(s).replace(/\[MASK\]/g, " ");
102
- let headIds = encode(`${question.type} question: ${strip(question.instructions)}`);
103
- let optIds = options.map((o) => [ids.mask, ...encode(` ${strip(o)}`).slice(0, 48)]);
104
-
105
- let optBudget = headMaxLen - optIds.reduce((n, o) => n + o.length, 0);
106
- if (optBudget < 16) {
107
- const per = Math.max(4, Math.floor((headMaxLen - 16) / Math.max(1, optIds.length)));
108
- optIds = optIds.map((o) => o.slice(0, per));
109
- optBudget = headMaxLen - optIds.reduce((n, o) => n + o.length, 0);
110
- }
111
- headIds = headIds.slice(0, Math.max(8, optBudget));
112
-
113
- let seq = [ids.cls, ...headIds, ids.sep];
114
- const markers = [];
115
- for (const o of optIds) {
116
- markers.push(seq.length);
117
- seq.push(...o);
118
- }
119
- seq.push(ids.sep);
120
- const room = Math.max(0, maxLen - seq.length - 1);
121
- const stateIds = encode(String(state ?? "")).slice(0, room);
122
- seq = [...seq, ...stateIds, ids.sep];
123
-
124
- const clipped = seq.slice(0, maxLen);
125
- return {
126
- ids: clipped,
127
- markers: markers.filter((m) => m < maxLen),
128
- qtype,
129
- };
130
- }
131
-
132
- /** Pad a batch of Laya items into int64/bool tensors. */
133
- export function collateLayaItems(items, padId = 0) {
134
- const n = items.length;
135
- const L = Math.max(...items.map((it) => it.ids.length));
136
- const M = Math.max(...items.map((it) => it.markers.length));
137
- const inputIds = new BigInt64Array(n * L).fill(BigInt(padId));
138
- const attention = new BigInt64Array(n * L);
139
- const markerPos = new BigInt64Array(n * M);
140
- const markerMask = new Uint8Array(n * M);
141
- const qtype = new BigInt64Array(n);
142
- items.forEach((it, i) => {
143
- it.ids.forEach((v, j) => { inputIds[i * L + j] = BigInt(v); attention[i * L + j] = 1n; });
144
- it.markers.forEach((m, j) => { markerPos[i * M + j] = BigInt(m); markerMask[i * M + j] = 1; });
145
- qtype[i] = BigInt(it.qtype);
146
- });
147
- return { n, L, M, inputIds, attention, markerPos, markerMask, qtype };
148
- }
149
-
150
- /**
151
- * Decode head logits (row-major [question][marker]) into jev-web answers.
152
- * @param {{logits:number[]|Float32Array, questions:object[], markerCounts:number[],
153
- * temperatures?:number[], temperatureByOptions?:object}} spec
154
- */
155
- export function layaAnswersFromLogits({
156
- logits,
157
- questions,
158
- markerCounts,
159
- temperatures = LAYA_DEFAULT_TEMPERATURES,
160
- temperatureByOptions = LAYA_DEFAULT_TEMPERATURES_BY_OPTIONS,
161
- }) {
162
- const answers = [];
163
- let offset = 0;
164
- questions.forEach((q, i) => {
165
- const k = markerCounts[i];
166
- const qtype = LAYA_QTYPE[q.type];
167
- const t = clampTemperature(
168
- temperatureByOptions[layaTempBucket(qtype, k)] ?? temperatures[qtype] ?? 1,
169
- );
170
- const probs = softmaxWithTemperature(Array.from(logits.slice(offset, offset + k), Number), t);
171
- offset += k;
172
-
173
- const entropy = k > 1
174
- ? -probs.reduce((acc, p) => acc + p * Math.log(Math.max(p, 1e-12)), 0) / Math.log(k)
175
- : 0;
176
- const confidence = Math.min(1, Math.max(0, 1 - entropy));
177
- const best = probs.indexOf(Math.max(...probs));
178
-
179
- if (q.type === "choice") {
180
- answers.push({
181
- type: "choice",
182
- choice: q.options[best],
183
- index: best,
184
- probabilities: Object.fromEntries(q.options.map((o, j) => [o, probs[j]])),
185
- confidence,
186
- });
187
- } else if (q.type === "score") {
188
- const expected = probs.reduce((acc, p, j) => acc + p * j, 0);
189
- answers.push({
190
- type: "score",
191
- score: expected,
192
- level: best,
193
- probabilities: Object.fromEntries(q.options.map((o, j) => [o, probs[j]])),
194
- confidence,
195
- });
196
- } else {
197
- const yes = probs[1];
198
- answers.push({
199
- type: "noul",
200
- noul: yes,
201
- probabilities: { [q.options[0]]: probs[0], [q.options[1]]: probs[1] },
202
- confidence: Math.max(yes, 1 - yes),
203
- });
204
- }
205
- });
206
- return answers;
207
- }
208
-
209
- /**
210
- * Fetch a JSON file with optional byte progress (fetch/Response compatible).
211
- * Falls back to `res.json()` when the body isn't streamable (Node fakes/tests).
212
- */
213
- async function fetchJson(url, file, { fetchImpl, onProgress }) {
214
- const bytes = await fetchCachedBytes(url, {
215
- fetchImpl,
216
- onProgress: (p) => onProgress?.({ ...p, phase: "tokenizer", file }),
217
- });
218
- return JSON.parse(new TextDecoder().decode(bytes));
219
- }
220
-
221
- /**
222
- * Load a Laya tokenizer directly from its files.
223
- *
224
- * transformers.js `AutoTokenizer.from_pretrained` ignores `subfolder`, and the
225
- * Laya web export keeps its tokenizer under `v1/` (and its root has no
226
- * `tokenizer_config.json`), so the library path throws
227
- * "Cannot read properties of undefined (reading 'tokenizer_class')".
228
- * Fetch the two JSON files ourselves and instantiate the class the config
229
- * names (e.g. `PreTrainedTokenizerFast` -> `PreTrainedTokenizer`).
230
- *
231
- * @param {{urlBase:string, fetchImpl?:Function, transformers?:object,
232
- * onProgress?:Function}} spec `urlBase` must end with "/".
233
- * @returns {Promise<object>} tokenizer instance
234
- */
235
- export async function loadLayaTokenizer({
236
- urlBase,
237
- fetchImpl = fetch,
238
- transformers = null,
239
- onProgress = null,
240
- } = {}) {
241
- if (!urlBase) throw new TypeError("loadLayaTokenizer: urlBase is required");
242
- const base = urlBase.endsWith("/") ? urlBase : `${urlBase}/`;
243
- const [tokenizerJSON, tokenizerConfig] = await Promise.all([
244
- fetchJson(`${base}tokenizer.json`, "tokenizer.json", { fetchImpl, onProgress }),
245
- fetchJson(`${base}tokenizer_config.json`, "tokenizer_config.json", { fetchImpl, onProgress }),
246
- ]);
247
- const tf = transformers ?? (await import("@huggingface/transformers"));
248
- const className = String(tokenizerConfig?.tokenizer_class ?? "PreTrainedTokenizer")
249
- .replace(/Fast$/, "");
250
- const Cls = tf[className] ?? tf.PreTrainedTokenizer;
251
- if (typeof Cls !== "function") {
252
- throw new Error(`loadLayaTokenizer: no tokenizer class for "${className}"`);
253
- }
254
- return new Cls(tokenizerJSON, tokenizerConfig);
255
- }
256
-
257
- /** Wrap a transformers.js tokenizer into the tiny adapter laya needs. */
258
- export function layaTokenizerAdapter(tok) {
259
- const idOf = (prop, text) => {
260
- const direct = tok?.[prop];
261
- if (Number.isFinite(direct)) return Number(direct);
262
- const ids = tok.encode?.(text) ?? tok(text, { add_special_tokens: true }).input_ids?.data;
263
- return Number(Array.from(ids ?? [])[0] ?? 0);
264
- };
265
- return {
266
- encode: (text) => Array.from(tok(text, { add_special_tokens: false }).input_ids.data, Number),
267
- decode: (ids) => {
268
- try {
269
- return tok.decode(Array.from(ids), { skip_special_tokens: false });
270
- } catch {
271
- return Array.from(ids).join(" ");
272
- }
273
- },
274
- ids: {
275
- cls: idOf("cls_token_id", "[CLS]"),
276
- sep: idOf("sep_token_id", "[SEP]"),
277
- mask: idOf("mask_token_id", "[MASK]"),
278
- },
279
- pad: idOf("pad_token_id", "[PAD]"),
280
- };
281
- }
282
-
283
- const now = () => (typeof performance !== "undefined" ? performance.now() : Date.now());
284
-
285
- const hfUrl = (model, revision, file) =>
286
- `https://huggingface.co/${model}/resolve/${revision ? `${revision}/` : ""}${file}`;
287
-
288
- const MODEL_CACHE = "kf-models-v1";
289
-
290
- async function modelCacheStore() {
291
- try {
292
- return typeof caches !== "undefined" ? caches : null;
293
- } catch {
294
- return null;
295
- }
296
- }
297
-
298
- /** Ask for persistent storage once so big weights survive eviction pressure. */
299
- export async function requestPersistentStorage() {
300
- try {
301
- return (await globalThis.navigator?.storage?.persist?.()) ?? false;
302
- } catch {
303
- return false;
304
- }
305
- }
306
-
307
- /**
308
- * Fetch bytes with a Cache-API layer. Hugging Face redirects to a *signed*
309
- * CDN URL that changes per request, so the browser HTTP cache never hits for
310
- * the 468 MB encoder weights; this caches them under the stable model URL
311
- * instead (Cache API), so a page refresh is free.
312
- *
313
- * @returns {Promise<Uint8Array>}
314
- */
315
- export async function fetchCachedBytes(url, {
316
- fetchImpl = fetch,
317
- onProgress = null,
318
- cacheName = MODEL_CACHE,
319
- force = false,
320
- persist = false,
321
- } = {}) {
322
- const store = await modelCacheStore();
323
- if (persist) await requestPersistentStorage();
324
- if (store && !force) {
325
- try {
326
- const hit = await store.match(url);
327
- if (hit) {
328
- const data = new Uint8Array(await hit.arrayBuffer());
329
- onProgress?.({ phase: "cache", file: String(url).split("/").pop(), loaded: data.length, total: data.length });
330
- return data;
331
- }
332
- } catch {
333
- /* cache read failed — fall through to the network */
334
- }
335
- }
336
- const res = await fetchImpl(url);
337
- if (!res?.ok) throw new Error(`fetch ${url} failed (HTTP ${res?.status ?? "no response"})`);
338
- const data = await readBytesWithProgress(res, url, onProgress);
339
- if (store) {
340
- try {
341
- await store.put(url, new Response(data, {
342
- headers: { "content-type": "application/octet-stream", "content-length": String(data.length) },
343
- }));
344
- } catch {
345
- /* quota exceeded / private mode — caching is best-effort */
346
- }
347
- }
348
- return data;
349
- }
350
-
351
- async function readBytesWithProgress(res, url, onProgress) {
352
- const total = Number(res.headers?.get?.("content-length")) || 0;
353
- const file = String(url).split("/").pop();
354
- if (res.body?.getReader) {
355
- const reader = res.body.getReader();
356
- let received = 0;
357
- let data = total ? new Uint8Array(total) : null;
358
- const chunks = total ? null : [];
359
- for (;;) {
360
- const { done, value } = await reader.read();
361
- if (done) break;
362
- if (data) {
363
- if (received + value.length > data.length) {
364
- // content-length can describe the *encoded* size (gzip/br), so the
365
- // decoded stream may be larger — grow instead of overflowing.
366
- const grown = new Uint8Array(Math.max(data.length * 2, received + value.length));
367
- grown.set(data.subarray(0, received));
368
- data = grown;
369
- }
370
- data.set(value, received);
371
- } else {
372
- chunks.push(value);
373
- }
374
- received += value.length;
375
- onProgress?.({ phase: "download", file, loaded: received, total: total || null });
376
- }
377
- if (!data) {
378
- data = new Uint8Array(received);
379
- let off = 0;
380
- for (const c of chunks) { data.set(c, off); off += c.length; }
381
- } else {
382
- // the buffer may have grown past the decoded length (content-length is
383
- // the *encoded* size) — always trim to what was actually received
384
- data = data.subarray(0, received);
385
- }
386
- return data;
387
- }
388
- let data;
389
- if (typeof res.arrayBuffer === "function") {
390
- data = new Uint8Array(await res.arrayBuffer());
391
- } else if (typeof res.text === "function") {
392
- data = new TextEncoder().encode(await res.text());
393
- } else if (typeof res.json === "function") {
394
- data = new TextEncoder().encode(JSON.stringify(await res.json()));
395
- } else {
396
- throw new Error(`fetch ${url}: response has no readable body`);
397
- }
398
- onProgress?.({ phase: "download", file, loaded: data.length, total: data.length });
399
- return data;
400
- }
401
-
402
- /**
403
- * ONNX external-data loader.
404
- *
405
- * The pinned Laya export stores weights in `<model>.onnx.data` siblings
406
- * (encoder ~468 MB, head ~53 MB). onnxruntime-web cannot resolve those from a
407
- * URL by itself (`Module.MountedFiles is not available`); it needs the bytes
408
- * passed via `session_options.externalData`. Fetch the sibling (streaming,
409
- * with progress) and hand it over.
410
- *
411
- * @returns {Promise<{path:string, data:Uint8Array}|null>} null when the model
412
- * has no external data (404) or the fetch fails.
413
- */
414
- export async function fetchOnnxExternalData(url, { fetchImpl = fetch, onProgress = null, force = false } = {}) {
415
- const dataUrl = `${url}.data`;
416
- const path = `${String(url).split("/").pop()}.data`;
417
- try {
418
- // Probe cheaply first: a 404 means the model has no external data.
419
- const head = await fetchImpl(dataUrl, { method: "HEAD" }).catch(() => null);
420
- if (head && !head.ok) return null;
421
- const data = await fetchCachedBytes(dataUrl, { fetchImpl, onProgress, force, persist: true });
422
- return { path, data };
423
- } catch {
424
- return null;
425
- }
426
- }
427
-
428
- /**
429
- * Create a Laya decider.
430
- *
431
- * @param {object} [opts]
432
- * model/revision/subfolder — pinned web-q8 ONNX export (or your own mirror)
433
- * encoderUrl/headUrl — override the resolved ONNX URLs
434
- * ort — injected onnxruntime-web (lazy import otherwise)
435
- * transformers — injected @huggingface/transformers (tokenizer only)
436
- * wasmPaths — ORT wasm/loader directory (self-hosted builds)
437
- * device — "auto" | "wasm" | "webgpu"
438
- */
439
- export async function createLayaDecider({
440
- model = LAYA_DEFAULT_MODEL,
441
- revision = LAYA_DEFAULT_REVISION,
442
- subfolder = LAYA_DEFAULT_SUBFOLDER,
443
- encoderUrl = null,
444
- headUrl = null,
445
- tokenizerModel = null,
446
- tokenizerSubfolder = null,
447
- ort = null,
448
- transformers = null,
449
- fetchImpl = fetch,
450
- wasmPaths = null,
451
- device = "auto",
452
- // (url, sessionOptions) => Promise<InferenceSession>. Hosts that need custom
453
- // loading (local files, bundled weights, offline mirrors) override this.
454
- sessionFactory = null,
455
- onProgress = null,
456
- maxLen = LAYA_DEFAULT_MAX_LEN,
457
- headMaxLen = LAYA_DEFAULT_HEAD_MAX_LEN,
458
- maxPrefixes = LAYA_DEFAULT_MAX_PREFIXES,
459
- temperatures = LAYA_DEFAULT_TEMPERATURES,
460
- temperatureByOptions = LAYA_DEFAULT_TEMPERATURES_BY_OPTIONS,
461
- scope = globalThis,
462
- } = {}) {
463
- const ortMod = ort ?? (await import("onnxruntime-web"));
464
- const tf = transformers ?? (await import("@huggingface/transformers"));
465
- if (wasmPaths) {
466
- try { ortMod.env.wasm.wasmPaths = wasmPaths; } catch { /* older runtime */ }
467
- }
468
-
469
- const resolvedDevice = device !== "auto"
470
- ? device
471
- : (scope?.navigator?.gpu ? "webgpu" : "wasm");
472
- const executionProviders = resolvedDevice === "webgpu" ? ["webgpu", "wasm"] : ["wasm"];
473
-
474
- const report = (phase) => (p) =>
475
- onProgress?.({ phase, file: p?.file ?? null, loaded: p?.loaded ?? null, total: p?.total ?? null });
476
-
477
- let tok;
478
- if (tokenizerModel) {
479
- // Host-provided path/repo (local dir in Node, or a repo whose tokenizer
480
- // lives at the root) — use the library loader.
481
- const tokSubfolder = tokenizerSubfolder === null
482
- ? undefined
483
- : (tokenizerSubfolder ?? subfolder);
484
- tok = await tf.AutoTokenizer.from_pretrained(tokenizerModel, {
485
- revision: tokenizerSubfolder ? undefined : revision,
486
- ...(tokSubfolder ? { subfolder: tokSubfolder } : {}),
487
- progress_callback: report("tokenizer"),
488
- });
489
- } else {
490
- // Default: the pinned web export keeps the tokenizer under `v1/`, which
491
- // AutoTokenizer cannot address (no `subfolder` support) — load it directly.
492
- tok = await loadLayaTokenizer({
493
- urlBase: hfUrl(model, revision, `${subfolder}/`),
494
- transformers: tf,
495
- fetchImpl,
496
- onProgress,
497
- });
498
- }
499
- const adapter = layaTokenizerAdapter(tok);
500
-
501
- const encUrl = encoderUrl ?? hfUrl(model, revision, `${subfolder}/encoder_q8.onnx`);
502
- const headUrlResolved = headUrl ?? hfUrl(model, revision, `${subfolder}/head_q8.onnx`);
503
- const createSession = sessionFactory
504
- ?? (async (url, opts) => {
505
- const report2 = (p) => onProgress?.(p);
506
- const modelBytes = await fetchCachedBytes(url, { fetchImpl, onProgress: report2, persist: true });
507
- const external = await fetchOnnxExternalData(url, { fetchImpl, onProgress: report2 });
508
- const sessionOpts = external ? { ...opts, externalData: [external] } : opts;
509
- return ortMod.InferenceSession.create(modelBytes, sessionOpts);
510
- });
511
- const loadPair = async (providers) => {
512
- const opts = { executionProviders: providers, graphOptimizationLevel: "all" };
513
- const enc = await createSession(encUrl, { ...opts, progress_callback: report("encoder") });
514
- const hd = await createSession(headUrlResolved, { ...opts, progress_callback: report("head") });
515
- return { enc, hd };
516
- };
517
- let encoder, head, activeDevice = resolvedDevice;
518
- try {
519
- ({ enc: encoder, hd: head } = await loadPair(executionProviders));
520
- } catch (e) {
521
- // A quantized graph can be rejected by the WebGPU EP at session creation;
522
- // retry wasm-only before giving up (the caller may still fall back).
523
- if (resolvedDevice === "webgpu") {
524
- onProgress?.({ phase: `webgpu unavailable (${e?.message ?? e}) — retrying on wasm` });
525
- ({ enc: encoder, hd: head } = await loadPair(["wasm"]));
526
- activeDevice = "wasm";
527
- } else {
528
- throw e;
529
- }
530
- }
531
-
532
- const int64 = (arr, dims) => new ortMod.Tensor("int64", arr, dims);
533
- const bool = (arr, dims) => new ortMod.Tensor("bool", arr, dims);
534
-
535
- return {
536
- info: {
537
- family: "laya",
538
- model,
539
- revision,
540
- device: activeDevice,
541
- dtype: "q8",
542
- maxLen,
543
- headMaxLen,
544
- maxPrefixes,
545
- // Single representative temperature for UIs that show one value
546
- // (noul:2 bucket, else the noul type temperature); per-bucket values
547
- // below are what the decoder actually applies.
548
- temperature: clampTemperature(temperatureByOptions["noul:2"] ?? temperatures[2] ?? 1),
549
- temperatures: temperatures.map(clampTemperature),
550
- temperatureByOptions: Object.fromEntries(
551
- Object.entries(temperatureByOptions).map(([k, v]) => [k, clampTemperature(v)])),
552
- },
553
- async decide(state, questions) {
554
- const t0 = now();
555
- const normalized = normalizeQuestions(questions);
556
- const items = normalized.map((q) =>
557
- buildLayaSequence({ encode: adapter.encode, ids: adapter.ids, state, question: q, maxLen, headMaxLen, maxPrefixes }));
558
- const b = collateLayaItems(items, adapter.pad);
559
-
560
- const tEncode = now();
561
- const encOut = await encoder.run({
562
- input_ids: int64(b.inputIds, [b.n, b.L]),
563
- attention_mask: int64(b.attention, [b.n, b.L]),
564
- });
565
- const encoderMs = now() - tEncode;
566
- const hidden = encOut.hidden ?? encOut.last_hidden_state ?? encOut[Object.keys(encOut)[0]];
567
-
568
- // The q8 web head is exported batch-1 only (despite dynamic batch
569
- // axes), so run it once per question over the shared encoder output.
570
- // The encoder — the expensive half — still runs as one batched pass.
571
- const headLogits = [];
572
- const headPerQuestionMs = [];
573
- const markerCounts = items.map((it) => it.markers.length);
574
- const tHead = now();
575
- for (let i = 0; i < b.n; i++) {
576
- const k = markerCounts[i];
577
- const row = hidden.data.subarray(i * b.L * hidden.dims[2], (i + 1) * b.L * hidden.dims[2]);
578
- const rowHidden = new ortMod.Tensor(hidden.type, row, [1, b.L, hidden.dims[2]]);
579
- const rowMask = b.attention.subarray(i * b.L, (i + 1) * b.L);
580
- const rowPos = b.markerPos.subarray(i * b.M, i * b.M + k);
581
- const rowMarkerMask = b.markerMask.subarray(i * b.M, i * b.M + k);
582
- const tHeadTotal = now();
583
- const headOut = await head.run({
584
- hidden: rowHidden,
585
- attention_mask: int64(rowMask, [1, b.L]),
586
- marker_pos: int64(rowPos, [1, k]),
587
- marker_mask: bool(rowMarkerMask, [1, k]),
588
- qtype: int64(b.qtype.subarray(i, i + 1), [1]),
589
- });
590
- const rowLogits = headOut.logits ?? headOut[Object.keys(headOut)[0]];
591
- headLogits.push(...Array.from(rowLogits.data, Number));
592
- headPerQuestionMs.push(Math.round(now() - tHeadTotal));
593
- }
594
-
595
- const answers = layaAnswersFromLogits({
596
- logits: headLogits,
597
- questions: normalized,
598
- markerCounts,
599
- temperatures,
600
- temperatureByOptions,
601
- });
602
- const length = Number(b.attention.reduce((a, v) => a + v, 0n));
603
- // Exact, human-readable view of what the encoder consumed: each
604
- // question's token sequence decoded back to text (special tokens kept).
605
- const prompts = items.map((it) => adapter.decode(it.ids));
606
- const headMs = now() - tHead;
607
- return {
608
- answers,
609
- truncated: items.some((it) => it.ids.length >= maxLen),
610
- length,
611
- prompts,
612
- timings: {
613
- totalMs: Math.round(now() - t0),
614
- encoderMs: Math.round(encoderMs),
615
- headMs: Math.round(headMs),
616
- headPerQuestionMs,
617
- },
618
- };
619
- },
620
- async dispose() {
621
- try { await encoder.release?.(); } catch { /* optional */ }
622
- try { await head.release?.(); } catch { /* optional */ }
623
- },
624
- };
625
- }
1
+ /**
2
+ * laya.mjs — Laya System-1 decision models (ModernBERT encoder + typed head).
3
+ *
4
+ * The Laya family is a *different graph* from open-jev:
5
+ * encoder: input_ids, attention_mask -> hidden (last_hidden_state)
6
+ * head: hidden, marker_pos, marker_mask, qtype -> logits, act_logits
7
+ *
8
+ * Sequence layout (mirrors the reference `laya` Python runtime,
9
+ * `laya/common.py build_sequence`):
10
+ * [CLS] <type> question: <instructions> [SEP] [MASK]opt0 [MASK]opt1 … [SEP] <state> [SEP]
11
+ * `marker_pos` points at each option's [MASK]; `qtype` is 0 choice / 1 score /
12
+ * 2 noul. Logits are calibrated per question type *and* option-count bucket
13
+ * (rl_agent_config.json `temperature` / `temperature_by_options`), clamped to
14
+ * [0.5, 5] like the Python runtime.
15
+ *
16
+ * The browser export used here is a two-file q8 ONNX split
17
+ * (alfred361/laya-typed-decisions-web-q8, pinned by revision).
18
+ */
19
+
20
+ import { normalizeQuestions } from "./questions.mjs";
21
+ import { softmaxWithTemperature } from "./answers.mjs";
22
+
23
+ export const LAYA_DEFAULT_MODEL = "alfred361/laya-typed-decisions-web-q8";
24
+ export const LAYA_DEFAULT_REVISION = "96f637c0bee046dce5c61a3174de9f7cdcf09ace";
25
+ export const LAYA_DEFAULT_SUBFOLDER = "v1";
26
+
27
+ /** Reference calibration from the typed-decisions rl_agent_config.json. */
28
+ export const LAYA_DEFAULT_TEMPERATURES = [1.0148024559020996, 1.0374259948730469, 1.0575125217437744];
29
+ export const LAYA_DEFAULT_TEMPERATURES_BY_OPTIONS = {
30
+ "choice:3-5": 1.7601518630981445,
31
+ "choice:6-10": 1.0000158548355103,
32
+ "score:3-5": 1.2514300346374512,
33
+ "noul:2": 1.983399510383606,
34
+ "choice:11+": 0.10058280825614929,
35
+ "choice:2": 1.9063563346862793,
36
+ };
37
+ export const LAYA_DEFAULT_MAX_LEN = 1024;
38
+ export const LAYA_DEFAULT_HEAD_MAX_LEN = 256;
39
+ export const LAYA_DEFAULT_MAX_PREFIXES = 6;
40
+
41
+ export const LAYA_QTYPE = { choice: 0, score: 1, noul: 2 };
42
+ const QTYPE_NAMES = { 0: "choice", 1: "score", 2: "noul" };
43
+
44
+ // The Python runtime refuses temperatures outside [0.5, 5]: the shipped
45
+ // `choice:11+` bucket (0.1006) would sharpen logits ~10x and publish a coin
46
+ // flip as near-certainty.
47
+ export const LAYA_TEMP_MIN = 0.5;
48
+ export const LAYA_TEMP_MAX = 5;
49
+
50
+ export function clampTemperature(t, lo = LAYA_TEMP_MIN, hi = LAYA_TEMP_MAX) {
51
+ const n = Number(t);
52
+ if (!Number.isFinite(n)) return 1;
53
+ return Math.min(hi, Math.max(lo, n));
54
+ }
55
+
56
+ export function layaTempBucket(qtype, k) {
57
+ const size = k <= 2 ? "2" : k <= 5 ? "3-5" : k <= 10 ? "6-10" : "11+";
58
+ return `${QTYPE_NAMES[qtype]}:${size}`;
59
+ }
60
+
61
+ /**
62
+ * Option texts exactly as the reference runtime renders them.
63
+ * Our normalized questions carry no per-option criteria, so choice options
64
+ * render as their labels and noul falls back to the reference wording.
65
+ */
66
+ export function renderLayaOptions(question) {
67
+ if (question.type === "choice") return question.options.map(String);
68
+ if (question.type === "score") return question.options.map((c, i) => `level ${i}: ${c}`);
69
+ // noul
70
+ const [no, yes] = question.options;
71
+ if (no === "no" && yes === "yes") {
72
+ return ["false: no, the statement does not hold", "true: yes, the statement holds"];
73
+ }
74
+ return [`false: ${no}`, `true: ${yes}`];
75
+ }
76
+
77
+ /**
78
+ * Build one Laya sequence + marker positions.
79
+ * @param {{encode:(text:string)=>number[], ids:{cls:number,sep:number,mask:number},
80
+ * state:string, question:object, maxLen?:number, headMaxLen?:number,
81
+ * maxPrefixes?:number}} spec
82
+ * @returns {{ids:number[], markers:number[], qtype:number}}
83
+ */
84
+ export function buildLayaSequence({
85
+ encode,
86
+ ids,
87
+ state,
88
+ question,
89
+ maxLen = LAYA_DEFAULT_MAX_LEN,
90
+ headMaxLen = LAYA_DEFAULT_HEAD_MAX_LEN,
91
+ maxPrefixes = LAYA_DEFAULT_MAX_PREFIXES,
92
+ }) {
93
+ const options = renderLayaOptions(question);
94
+ if (options.length > maxPrefixes) {
95
+ throw new RangeError(
96
+ `laya: ${question.type} with ${options.length} options exceeds max_prefixes=${maxPrefixes}`);
97
+ }
98
+ const qtype = LAYA_QTYPE[question.type];
99
+ if (qtype === undefined) throw new TypeError(`laya: unsupported question type ${question.type}`);
100
+
101
+ const strip = (s) => String(s).replace(/\[MASK\]/g, " ");
102
+ let headIds = encode(`${question.type} question: ${strip(question.instructions)}`);
103
+ let optIds = options.map((o) => [ids.mask, ...encode(` ${strip(o)}`).slice(0, 48)]);
104
+
105
+ let optBudget = headMaxLen - optIds.reduce((n, o) => n + o.length, 0);
106
+ if (optBudget < 16) {
107
+ const per = Math.max(4, Math.floor((headMaxLen - 16) / Math.max(1, optIds.length)));
108
+ optIds = optIds.map((o) => o.slice(0, per));
109
+ optBudget = headMaxLen - optIds.reduce((n, o) => n + o.length, 0);
110
+ }
111
+ headIds = headIds.slice(0, Math.max(8, optBudget));
112
+
113
+ let seq = [ids.cls, ...headIds, ids.sep];
114
+ const markers = [];
115
+ for (const o of optIds) {
116
+ markers.push(seq.length);
117
+ seq.push(...o);
118
+ }
119
+ seq.push(ids.sep);
120
+ const room = Math.max(0, maxLen - seq.length - 1);
121
+ const stateIds = encode(String(state ?? "")).slice(0, room);
122
+ seq = [...seq, ...stateIds, ids.sep];
123
+
124
+ const clipped = seq.slice(0, maxLen);
125
+ return {
126
+ ids: clipped,
127
+ markers: markers.filter((m) => m < maxLen),
128
+ qtype,
129
+ };
130
+ }
131
+
132
+ /** Pad a batch of Laya items into int64/bool tensors. */
133
+ export function collateLayaItems(items, padId = 0) {
134
+ const n = items.length;
135
+ const L = Math.max(...items.map((it) => it.ids.length));
136
+ const M = Math.max(...items.map((it) => it.markers.length));
137
+ const inputIds = new BigInt64Array(n * L).fill(BigInt(padId));
138
+ const attention = new BigInt64Array(n * L);
139
+ const markerPos = new BigInt64Array(n * M);
140
+ const markerMask = new Uint8Array(n * M);
141
+ const qtype = new BigInt64Array(n);
142
+ items.forEach((it, i) => {
143
+ it.ids.forEach((v, j) => { inputIds[i * L + j] = BigInt(v); attention[i * L + j] = 1n; });
144
+ it.markers.forEach((m, j) => { markerPos[i * M + j] = BigInt(m); markerMask[i * M + j] = 1; });
145
+ qtype[i] = BigInt(it.qtype);
146
+ });
147
+ return { n, L, M, inputIds, attention, markerPos, markerMask, qtype };
148
+ }
149
+
150
+ /**
151
+ * Decode head logits (row-major [question][marker]) into jev-web answers.
152
+ * @param {{logits:number[]|Float32Array, questions:object[], markerCounts:number[],
153
+ * temperatures?:number[], temperatureByOptions?:object}} spec
154
+ */
155
+ export function layaAnswersFromLogits({
156
+ logits,
157
+ questions,
158
+ markerCounts,
159
+ temperatures = LAYA_DEFAULT_TEMPERATURES,
160
+ temperatureByOptions = LAYA_DEFAULT_TEMPERATURES_BY_OPTIONS,
161
+ }) {
162
+ const answers = [];
163
+ let offset = 0;
164
+ questions.forEach((q, i) => {
165
+ const k = markerCounts[i];
166
+ const qtype = LAYA_QTYPE[q.type];
167
+ const t = clampTemperature(
168
+ temperatureByOptions[layaTempBucket(qtype, k)] ?? temperatures[qtype] ?? 1,
169
+ );
170
+ const probs = softmaxWithTemperature(Array.from(logits.slice(offset, offset + k), Number), t);
171
+ offset += k;
172
+
173
+ const entropy = k > 1
174
+ ? -probs.reduce((acc, p) => acc + p * Math.log(Math.max(p, 1e-12)), 0) / Math.log(k)
175
+ : 0;
176
+ const confidence = Math.min(1, Math.max(0, 1 - entropy));
177
+ const best = probs.indexOf(Math.max(...probs));
178
+
179
+ if (q.type === "choice") {
180
+ answers.push({
181
+ type: "choice",
182
+ choice: q.options[best],
183
+ index: best,
184
+ probabilities: Object.fromEntries(q.options.map((o, j) => [o, probs[j]])),
185
+ confidence,
186
+ });
187
+ } else if (q.type === "score") {
188
+ const expected = probs.reduce((acc, p, j) => acc + p * j, 0);
189
+ answers.push({
190
+ type: "score",
191
+ score: expected,
192
+ level: best,
193
+ probabilities: Object.fromEntries(q.options.map((o, j) => [o, probs[j]])),
194
+ confidence,
195
+ });
196
+ } else {
197
+ const yes = probs[1];
198
+ answers.push({
199
+ type: "noul",
200
+ noul: yes,
201
+ probabilities: { [q.options[0]]: probs[0], [q.options[1]]: probs[1] },
202
+ confidence: Math.max(yes, 1 - yes),
203
+ });
204
+ }
205
+ });
206
+ return answers;
207
+ }
208
+
209
+ /**
210
+ * Fetch a JSON file with optional byte progress (fetch/Response compatible).
211
+ * Falls back to `res.json()` when the body isn't streamable (Node fakes/tests).
212
+ */
213
+ async function fetchJson(url, file, { fetchImpl, onProgress }) {
214
+ const bytes = await fetchCachedBytes(url, {
215
+ fetchImpl,
216
+ onProgress: (p) => onProgress?.({ ...p, phase: "tokenizer", file }),
217
+ });
218
+ return JSON.parse(new TextDecoder().decode(bytes));
219
+ }
220
+
221
+ /**
222
+ * Load a Laya tokenizer directly from its files.
223
+ *
224
+ * transformers.js `AutoTokenizer.from_pretrained` ignores `subfolder`, and the
225
+ * Laya web export keeps its tokenizer under `v1/` (and its root has no
226
+ * `tokenizer_config.json`), so the library path throws
227
+ * "Cannot read properties of undefined (reading 'tokenizer_class')".
228
+ * Fetch the two JSON files ourselves and instantiate the class the config
229
+ * names (e.g. `PreTrainedTokenizerFast` -> `PreTrainedTokenizer`).
230
+ *
231
+ * @param {{urlBase:string, fetchImpl?:Function, transformers?:object,
232
+ * onProgress?:Function}} spec `urlBase` must end with "/".
233
+ * @returns {Promise<object>} tokenizer instance
234
+ */
235
+ export async function loadLayaTokenizer({
236
+ urlBase,
237
+ fetchImpl = fetch,
238
+ transformers = null,
239
+ onProgress = null,
240
+ } = {}) {
241
+ if (!urlBase) throw new TypeError("loadLayaTokenizer: urlBase is required");
242
+ const base = urlBase.endsWith("/") ? urlBase : `${urlBase}/`;
243
+ const [tokenizerJSON, tokenizerConfig] = await Promise.all([
244
+ fetchJson(`${base}tokenizer.json`, "tokenizer.json", { fetchImpl, onProgress }),
245
+ fetchJson(`${base}tokenizer_config.json`, "tokenizer_config.json", { fetchImpl, onProgress }),
246
+ ]);
247
+ const tf = transformers ?? (await import("@huggingface/transformers"));
248
+ const className = String(tokenizerConfig?.tokenizer_class ?? "PreTrainedTokenizer")
249
+ .replace(/Fast$/, "");
250
+ const Cls = tf[className] ?? tf.PreTrainedTokenizer;
251
+ if (typeof Cls !== "function") {
252
+ throw new Error(`loadLayaTokenizer: no tokenizer class for "${className}"`);
253
+ }
254
+ return new Cls(tokenizerJSON, tokenizerConfig);
255
+ }
256
+
257
+ /** Wrap a transformers.js tokenizer into the tiny adapter laya needs. */
258
+ export function layaTokenizerAdapter(tok) {
259
+ const idOf = (prop, text) => {
260
+ const direct = tok?.[prop];
261
+ if (Number.isFinite(direct)) return Number(direct);
262
+ const ids = tok.encode?.(text) ?? tok(text, { add_special_tokens: true }).input_ids?.data;
263
+ return Number(Array.from(ids ?? [])[0] ?? 0);
264
+ };
265
+ return {
266
+ encode: (text) => Array.from(tok(text, { add_special_tokens: false }).input_ids.data, Number),
267
+ decode: (ids) => {
268
+ try {
269
+ return tok.decode(Array.from(ids), { skip_special_tokens: false });
270
+ } catch {
271
+ return Array.from(ids).join(" ");
272
+ }
273
+ },
274
+ ids: {
275
+ cls: idOf("cls_token_id", "[CLS]"),
276
+ sep: idOf("sep_token_id", "[SEP]"),
277
+ mask: idOf("mask_token_id", "[MASK]"),
278
+ },
279
+ pad: idOf("pad_token_id", "[PAD]"),
280
+ };
281
+ }
282
+
283
+ const now = () => (typeof performance !== "undefined" ? performance.now() : Date.now());
284
+
285
+ const hfUrl = (model, revision, file) =>
286
+ `https://huggingface.co/${model}/resolve/${revision ? `${revision}/` : ""}${file}`;
287
+
288
+ const MODEL_CACHE = "kf-models-v1";
289
+
290
+ async function modelCacheStore() {
291
+ try {
292
+ return typeof caches !== "undefined" ? caches : null;
293
+ } catch {
294
+ return null;
295
+ }
296
+ }
297
+
298
+ /** Ask for persistent storage once so big weights survive eviction pressure. */
299
+ export async function requestPersistentStorage() {
300
+ try {
301
+ return (await globalThis.navigator?.storage?.persist?.()) ?? false;
302
+ } catch {
303
+ return false;
304
+ }
305
+ }
306
+
307
+ /**
308
+ * Fetch bytes with a Cache-API layer. Hugging Face redirects to a *signed*
309
+ * CDN URL that changes per request, so the browser HTTP cache never hits for
310
+ * the 468 MB encoder weights; this caches them under the stable model URL
311
+ * instead (Cache API), so a page refresh is free.
312
+ *
313
+ * @returns {Promise<Uint8Array>}
314
+ */
315
+ export async function fetchCachedBytes(url, {
316
+ fetchImpl = fetch,
317
+ onProgress = null,
318
+ cacheName = MODEL_CACHE,
319
+ force = false,
320
+ persist = false,
321
+ } = {}) {
322
+ const store = await modelCacheStore();
323
+ if (persist) await requestPersistentStorage();
324
+ if (store && !force) {
325
+ try {
326
+ const hit = await store.match(url);
327
+ if (hit) {
328
+ const data = new Uint8Array(await hit.arrayBuffer());
329
+ onProgress?.({ phase: "cache", file: String(url).split("/").pop(), loaded: data.length, total: data.length });
330
+ return data;
331
+ }
332
+ } catch {
333
+ /* cache read failed — fall through to the network */
334
+ }
335
+ }
336
+ const res = await fetchImpl(url);
337
+ if (!res?.ok) throw new Error(`fetch ${url} failed (HTTP ${res?.status ?? "no response"})`);
338
+ const data = await readBytesWithProgress(res, url, onProgress);
339
+ if (store) {
340
+ try {
341
+ await store.put(url, new Response(data, {
342
+ headers: { "content-type": "application/octet-stream", "content-length": String(data.length) },
343
+ }));
344
+ } catch {
345
+ /* quota exceeded / private mode — caching is best-effort */
346
+ }
347
+ }
348
+ return data;
349
+ }
350
+
351
+ async function readBytesWithProgress(res, url, onProgress) {
352
+ const total = Number(res.headers?.get?.("content-length")) || 0;
353
+ const file = String(url).split("/").pop();
354
+ if (res.body?.getReader) {
355
+ const reader = res.body.getReader();
356
+ let received = 0;
357
+ let data = total ? new Uint8Array(total) : null;
358
+ const chunks = total ? null : [];
359
+ for (;;) {
360
+ const { done, value } = await reader.read();
361
+ if (done) break;
362
+ if (data) {
363
+ if (received + value.length > data.length) {
364
+ // content-length can describe the *encoded* size (gzip/br), so the
365
+ // decoded stream may be larger — grow instead of overflowing.
366
+ const grown = new Uint8Array(Math.max(data.length * 2, received + value.length));
367
+ grown.set(data.subarray(0, received));
368
+ data = grown;
369
+ }
370
+ data.set(value, received);
371
+ } else {
372
+ chunks.push(value);
373
+ }
374
+ received += value.length;
375
+ onProgress?.({ phase: "download", file, loaded: received, total: total || null });
376
+ }
377
+ if (!data) {
378
+ data = new Uint8Array(received);
379
+ let off = 0;
380
+ for (const c of chunks) { data.set(c, off); off += c.length; }
381
+ } else {
382
+ // the buffer may have grown past the decoded length (content-length is
383
+ // the *encoded* size) — always trim to what was actually received
384
+ data = data.subarray(0, received);
385
+ }
386
+ return data;
387
+ }
388
+ let data;
389
+ if (typeof res.arrayBuffer === "function") {
390
+ data = new Uint8Array(await res.arrayBuffer());
391
+ } else if (typeof res.text === "function") {
392
+ data = new TextEncoder().encode(await res.text());
393
+ } else if (typeof res.json === "function") {
394
+ data = new TextEncoder().encode(JSON.stringify(await res.json()));
395
+ } else {
396
+ throw new Error(`fetch ${url}: response has no readable body`);
397
+ }
398
+ onProgress?.({ phase: "download", file, loaded: data.length, total: data.length });
399
+ return data;
400
+ }
401
+
402
+ /**
403
+ * ONNX external-data loader.
404
+ *
405
+ * The pinned Laya export stores weights in `<model>.onnx.data` siblings
406
+ * (encoder ~468 MB, head ~53 MB). onnxruntime-web cannot resolve those from a
407
+ * URL by itself (`Module.MountedFiles is not available`); it needs the bytes
408
+ * passed via `session_options.externalData`. Fetch the sibling (streaming,
409
+ * with progress) and hand it over.
410
+ *
411
+ * @returns {Promise<{path:string, data:Uint8Array}|null>} null when the model
412
+ * has no external data (404) or the fetch fails.
413
+ */
414
+ export async function fetchOnnxExternalData(url, { fetchImpl = fetch, onProgress = null, force = false } = {}) {
415
+ const dataUrl = `${url}.data`;
416
+ const path = `${String(url).split("/").pop()}.data`;
417
+ try {
418
+ // Probe cheaply first: a 404 means the model has no external data.
419
+ const head = await fetchImpl(dataUrl, { method: "HEAD" }).catch(() => null);
420
+ if (head && !head.ok) return null;
421
+ const data = await fetchCachedBytes(dataUrl, { fetchImpl, onProgress, force, persist: true });
422
+ return { path, data };
423
+ } catch {
424
+ return null;
425
+ }
426
+ }
427
+
428
+ /**
429
+ * Create a Laya decider.
430
+ *
431
+ * @param {object} [opts]
432
+ * model/revision/subfolder — pinned web-q8 ONNX export (or your own mirror)
433
+ * encoderUrl/headUrl — override the resolved ONNX URLs
434
+ * ort — injected onnxruntime-web (lazy import otherwise)
435
+ * transformers — injected @huggingface/transformers (tokenizer only)
436
+ * wasmPaths — ORT wasm/loader directory (self-hosted builds)
437
+ * device — "auto" | "wasm" | "webgpu"
438
+ */
439
+ export async function createLayaDecider({
440
+ model = LAYA_DEFAULT_MODEL,
441
+ revision = LAYA_DEFAULT_REVISION,
442
+ subfolder = LAYA_DEFAULT_SUBFOLDER,
443
+ encoderUrl = null,
444
+ headUrl = null,
445
+ tokenizerModel = null,
446
+ tokenizerSubfolder = null,
447
+ ort = null,
448
+ transformers = null,
449
+ fetchImpl = fetch,
450
+ wasmPaths = null,
451
+ device = "auto",
452
+ // (url, sessionOptions) => Promise<InferenceSession>. Hosts that need custom
453
+ // loading (local files, bundled weights, offline mirrors) override this.
454
+ sessionFactory = null,
455
+ onProgress = null,
456
+ maxLen = LAYA_DEFAULT_MAX_LEN,
457
+ headMaxLen = LAYA_DEFAULT_HEAD_MAX_LEN,
458
+ maxPrefixes = LAYA_DEFAULT_MAX_PREFIXES,
459
+ temperatures = LAYA_DEFAULT_TEMPERATURES,
460
+ temperatureByOptions = LAYA_DEFAULT_TEMPERATURES_BY_OPTIONS,
461
+ scope = globalThis,
462
+ } = {}) {
463
+ const ortMod = ort ?? (await import("onnxruntime-web"));
464
+ const tf = transformers ?? (await import("@huggingface/transformers"));
465
+ if (wasmPaths) {
466
+ try { ortMod.env.wasm.wasmPaths = wasmPaths; } catch { /* older runtime */ }
467
+ }
468
+
469
+ const resolvedDevice = device !== "auto"
470
+ ? device
471
+ : (scope?.navigator?.gpu ? "webgpu" : "wasm");
472
+ const executionProviders = resolvedDevice === "webgpu" ? ["webgpu", "wasm"] : ["wasm"];
473
+
474
+ const report = (phase) => (p) =>
475
+ onProgress?.({ phase, file: p?.file ?? null, loaded: p?.loaded ?? null, total: p?.total ?? null });
476
+
477
+ let tok;
478
+ if (tokenizerModel) {
479
+ // Host-provided path/repo (local dir in Node, or a repo whose tokenizer
480
+ // lives at the root) — use the library loader.
481
+ const tokSubfolder = tokenizerSubfolder === null
482
+ ? undefined
483
+ : (tokenizerSubfolder ?? subfolder);
484
+ tok = await tf.AutoTokenizer.from_pretrained(tokenizerModel, {
485
+ revision: tokenizerSubfolder ? undefined : revision,
486
+ ...(tokSubfolder ? { subfolder: tokSubfolder } : {}),
487
+ progress_callback: report("tokenizer"),
488
+ });
489
+ } else {
490
+ // Default: the pinned web export keeps the tokenizer under `v1/`, which
491
+ // AutoTokenizer cannot address (no `subfolder` support) — load it directly.
492
+ tok = await loadLayaTokenizer({
493
+ urlBase: hfUrl(model, revision, `${subfolder}/`),
494
+ transformers: tf,
495
+ fetchImpl,
496
+ onProgress,
497
+ });
498
+ }
499
+ const adapter = layaTokenizerAdapter(tok);
500
+
501
+ const encUrl = encoderUrl ?? hfUrl(model, revision, `${subfolder}/encoder_q8.onnx`);
502
+ const headUrlResolved = headUrl ?? hfUrl(model, revision, `${subfolder}/head_q8.onnx`);
503
+ const createSession = sessionFactory
504
+ ?? (async (url, opts) => {
505
+ const report2 = (p) => onProgress?.(p);
506
+ const modelBytes = await fetchCachedBytes(url, { fetchImpl, onProgress: report2, persist: true });
507
+ const external = await fetchOnnxExternalData(url, { fetchImpl, onProgress: report2 });
508
+ const sessionOpts = external ? { ...opts, externalData: [external] } : opts;
509
+ return ortMod.InferenceSession.create(modelBytes, sessionOpts);
510
+ });
511
+ const loadPair = async (providers) => {
512
+ const opts = { executionProviders: providers, graphOptimizationLevel: "all" };
513
+ const enc = await createSession(encUrl, { ...opts, progress_callback: report("encoder") });
514
+ const hd = await createSession(headUrlResolved, { ...opts, progress_callback: report("head") });
515
+ return { enc, hd };
516
+ };
517
+ let encoder, head, activeDevice = resolvedDevice;
518
+ try {
519
+ ({ enc: encoder, hd: head } = await loadPair(executionProviders));
520
+ } catch (e) {
521
+ // A quantized graph can be rejected by the WebGPU EP at session creation;
522
+ // retry wasm-only before giving up (the caller may still fall back).
523
+ if (resolvedDevice === "webgpu") {
524
+ onProgress?.({ phase: `webgpu unavailable (${e?.message ?? e}) — retrying on wasm` });
525
+ ({ enc: encoder, hd: head } = await loadPair(["wasm"]));
526
+ activeDevice = "wasm";
527
+ } else {
528
+ throw e;
529
+ }
530
+ }
531
+
532
+ const int64 = (arr, dims) => new ortMod.Tensor("int64", arr, dims);
533
+ const bool = (arr, dims) => new ortMod.Tensor("bool", arr, dims);
534
+
535
+ return {
536
+ info: {
537
+ family: "laya",
538
+ model,
539
+ revision,
540
+ device: activeDevice,
541
+ dtype: "q8",
542
+ maxLen,
543
+ headMaxLen,
544
+ maxPrefixes,
545
+ // Single representative temperature for UIs that show one value
546
+ // (noul:2 bucket, else the noul type temperature); per-bucket values
547
+ // below are what the decoder actually applies.
548
+ temperature: clampTemperature(temperatureByOptions["noul:2"] ?? temperatures[2] ?? 1),
549
+ temperatures: temperatures.map(clampTemperature),
550
+ temperatureByOptions: Object.fromEntries(
551
+ Object.entries(temperatureByOptions).map(([k, v]) => [k, clampTemperature(v)])),
552
+ },
553
+ async decide(state, questions) {
554
+ const t0 = now();
555
+ const normalized = normalizeQuestions(questions);
556
+ const items = normalized.map((q) =>
557
+ buildLayaSequence({ encode: adapter.encode, ids: adapter.ids, state, question: q, maxLen, headMaxLen, maxPrefixes }));
558
+ const b = collateLayaItems(items, adapter.pad);
559
+
560
+ const tEncode = now();
561
+ const encOut = await encoder.run({
562
+ input_ids: int64(b.inputIds, [b.n, b.L]),
563
+ attention_mask: int64(b.attention, [b.n, b.L]),
564
+ });
565
+ const encoderMs = now() - tEncode;
566
+ const hidden = encOut.hidden ?? encOut.last_hidden_state ?? encOut[Object.keys(encOut)[0]];
567
+
568
+ // The q8 web head is exported batch-1 only (despite dynamic batch
569
+ // axes), so run it once per question over the shared encoder output.
570
+ // The encoder — the expensive half — still runs as one batched pass.
571
+ const headLogits = [];
572
+ const headPerQuestionMs = [];
573
+ const markerCounts = items.map((it) => it.markers.length);
574
+ const tHead = now();
575
+ for (let i = 0; i < b.n; i++) {
576
+ const k = markerCounts[i];
577
+ const row = hidden.data.subarray(i * b.L * hidden.dims[2], (i + 1) * b.L * hidden.dims[2]);
578
+ const rowHidden = new ortMod.Tensor(hidden.type, row, [1, b.L, hidden.dims[2]]);
579
+ const rowMask = b.attention.subarray(i * b.L, (i + 1) * b.L);
580
+ const rowPos = b.markerPos.subarray(i * b.M, i * b.M + k);
581
+ const rowMarkerMask = b.markerMask.subarray(i * b.M, i * b.M + k);
582
+ const tHeadTotal = now();
583
+ const headOut = await head.run({
584
+ hidden: rowHidden,
585
+ attention_mask: int64(rowMask, [1, b.L]),
586
+ marker_pos: int64(rowPos, [1, k]),
587
+ marker_mask: bool(rowMarkerMask, [1, k]),
588
+ qtype: int64(b.qtype.subarray(i, i + 1), [1]),
589
+ });
590
+ const rowLogits = headOut.logits ?? headOut[Object.keys(headOut)[0]];
591
+ headLogits.push(...Array.from(rowLogits.data, Number));
592
+ headPerQuestionMs.push(Math.round(now() - tHeadTotal));
593
+ }
594
+
595
+ const answers = layaAnswersFromLogits({
596
+ logits: headLogits,
597
+ questions: normalized,
598
+ markerCounts,
599
+ temperatures,
600
+ temperatureByOptions,
601
+ });
602
+ const length = Number(b.attention.reduce((a, v) => a + v, 0n));
603
+ // Exact, human-readable view of what the encoder consumed: each
604
+ // question's token sequence decoded back to text (special tokens kept).
605
+ const prompts = items.map((it) => adapter.decode(it.ids));
606
+ const headMs = now() - tHead;
607
+ return {
608
+ answers,
609
+ truncated: items.some((it) => it.ids.length >= maxLen),
610
+ length,
611
+ prompts,
612
+ timings: {
613
+ totalMs: Math.round(now() - t0),
614
+ encoderMs: Math.round(encoderMs),
615
+ headMs: Math.round(headMs),
616
+ headPerQuestionMs,
617
+ },
618
+ };
619
+ },
620
+ async dispose() {
621
+ try { await encoder.release?.(); } catch { /* optional */ }
622
+ try { await head.release?.(); } catch { /* optional */ }
623
+ },
624
+ };
625
+ }