jev-web 0.2.1 → 0.3.1

package/README.md

@@ -91,6 +91,36 @@ take `input_ids, attention_mask, seg, pair_q, pair_opt`, and return per-pair
 `logits`. `typedMarkerIds()` fails loudly when a repo is missing them.
 
 
+## Multiple families (registry)
+
+`jev-web` is a multi-family runtime: `createDecisionRuntime({ family, fallback })`
+tries the requested family and then any fallbacks, returning the shared contract
+`{ info, decide(state, questions), dispose? }`:
+
+```js
+import { createDecisionRuntime, listDecisionFamilies } from "jev-web";
+
+listDecisionFamilies(); // ["open-jev", "laya"]
+const runtime = await createDecisionRuntime({
+  family: "laya",
+  fallback: ["open-jev"],
+  onFallback: ({ family, error }) => console.warn(family, error),
+});
+const { answers, truncated, length, prompts, timings } = await runtime.decide(state, questions);
+```
+
+Third-party families plug in without touching core:
+
+```js
+registerDecisionFamily("my-model", {
+  defaults: { model: "org/my-typed-decision-ONNX", revision: "main" },
+  create: async (opts) => ({ info, decide, dispose }),
+});
+```
+
+`decide()` returns `prompts` (decoded input sequence(s), for audit/debug) and
+`timings` (per-stage ms) alongside the answers.
+
 ## Laya family (ModernBERT encoder + typed head)
 
 `createLayaDecider(options)` runs the [Laya](https://huggingface.co/convaiinnovations/laya)

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "jev-web",
-  "version": "0.2.1",
+  "version": "0.3.1",
   "description": "Run open-jev- and Laya-shaped typed-decision models in the browser: one state plus typed questions in, calibrated probabilities out.",
   "type": "module",
   "license": "MIT",

package/src/index.mjs

@@ -1,5 +1,9 @@
 // jev-web — public API (typed decisions in the browser).
 export { createDecider } from "./session.mjs";
+export {
+  registerDecisionFamily, createDecisionRuntime,
+  listDecisionFamilies, getDecisionFamily,
+} from "./registry.mjs";
 export { normalizeQuestions, QUESTION_TYPES, NOUL_OPTIONS } from "./questions.mjs";
 export { buildDecisionInput, typedMarkerIds, DecisionInputTooLongError } from "./encode.mjs";
 export { answersFromScores, softmaxWithTemperature } from "./answers.mjs";

package/src/laya.mjs

@@ -206,6 +206,54 @@ export function layaAnswersFromLogits({
   return answers;
 }
 
+/**
+ * Fetch a JSON file with optional byte progress (fetch/Response compatible).
+ * Falls back to `res.json()` when the body isn't streamable (Node fakes/tests).
+ */
+async function fetchJson(url, file, { fetchImpl, onProgress }) {
+  const bytes = await fetchCachedBytes(url, {
+    fetchImpl,
+    onProgress: (p) => onProgress?.({ ...p, phase: "tokenizer", file }),
+  });
+  return JSON.parse(new TextDecoder().decode(bytes));
+}
+
+/**
+ * Load a Laya tokenizer directly from its files.
+ *
+ * transformers.js `AutoTokenizer.from_pretrained` ignores `subfolder`, and the
+ * Laya web export keeps its tokenizer under `v1/` (and its root has no
+ * `tokenizer_config.json`), so the library path throws
+ * "Cannot read properties of undefined (reading 'tokenizer_class')".
+ * Fetch the two JSON files ourselves and instantiate the class the config
+ * names (e.g. `PreTrainedTokenizerFast` -> `PreTrainedTokenizer`).
+ *
+ * @param {{urlBase:string, fetchImpl?:Function, transformers?:object,
+ *   onProgress?:Function}} spec `urlBase` must end with "/".
+ * @returns {Promise<object>} tokenizer instance
+ */
+export async function loadLayaTokenizer({
+  urlBase,
+  fetchImpl = fetch,
+  transformers = null,
+  onProgress = null,
+} = {}) {
+  if (!urlBase) throw new TypeError("loadLayaTokenizer: urlBase is required");
+  const base = urlBase.endsWith("/") ? urlBase : `${urlBase}/`;
+  const [tokenizerJSON, tokenizerConfig] = await Promise.all([
+    fetchJson(`${base}tokenizer.json`, "tokenizer.json", { fetchImpl, onProgress }),
+    fetchJson(`${base}tokenizer_config.json`, "tokenizer_config.json", { fetchImpl, onProgress }),
+  ]);
+  const tf = transformers ?? (await import("@huggingface/transformers"));
+  const className = String(tokenizerConfig?.tokenizer_class ?? "PreTrainedTokenizer")
+    .replace(/Fast$/, "");
+  const Cls = tf[className] ?? tf.PreTrainedTokenizer;
+  if (typeof Cls !== "function") {
+    throw new Error(`loadLayaTokenizer: no tokenizer class for "${className}"`);
+  }
+  return new Cls(tokenizerJSON, tokenizerConfig);
+}
+
 /** Wrap a transformers.js tokenizer into the tiny adapter laya needs. */
 export function layaTokenizerAdapter(tok) {
   const idOf = (prop, text) => {
@@ -232,9 +280,151 @@ export function layaTokenizerAdapter(tok) {
   };
 }
 
+const now = () => (typeof performance !== "undefined" ? performance.now() : Date.now());
+
 const hfUrl = (model, revision, file) =>
   `https://huggingface.co/${model}/resolve/${revision ? `${revision}/` : ""}${file}`;
 
+const MODEL_CACHE = "kf-models-v1";
+
+async function modelCacheStore() {
+  try {
+    return typeof caches !== "undefined" ? caches : null;
+  } catch {
+    return null;
+  }
+}
+
+/** Ask for persistent storage once so big weights survive eviction pressure. */
+export async function requestPersistentStorage() {
+  try {
+    return (await globalThis.navigator?.storage?.persist?.()) ?? false;
+  } catch {
+    return false;
+  }
+}
+
+/**
+ * Fetch bytes with a Cache-API layer. Hugging Face redirects to a *signed*
+ * CDN URL that changes per request, so the browser HTTP cache never hits for
+ * the 468 MB encoder weights; this caches them under the stable model URL
+ * instead (Cache API), so a page refresh is free.
+ *
+ * @returns {Promise<Uint8Array>}
+ */
+export async function fetchCachedBytes(url, {
+  fetchImpl = fetch,
+  onProgress = null,
+  cacheName = MODEL_CACHE,
+  force = false,
+  persist = false,
+} = {}) {
+  const store = await modelCacheStore();
+  if (persist) await requestPersistentStorage();
+  if (store && !force) {
+    try {
+      const hit = await store.match(url);
+      if (hit) {
+        const data = new Uint8Array(await hit.arrayBuffer());
+        onProgress?.({ phase: "cache", file: String(url).split("/").pop(), loaded: data.length, total: data.length });
+        return data;
+      }
+    } catch {
+      /* cache read failed — fall through to the network */
+    }
+  }
+  const res = await fetchImpl(url);
+  if (!res?.ok) throw new Error(`fetch ${url} failed (HTTP ${res?.status ?? "no response"})`);
+  const data = await readBytesWithProgress(res, url, onProgress);
+  if (store) {
+    try {
+      await store.put(url, new Response(data, {
+        headers: { "content-type": "application/octet-stream", "content-length": String(data.length) },
+      }));
+    } catch {
+      /* quota exceeded / private mode — caching is best-effort */
+    }
+  }
+  return data;
+}
+
+async function readBytesWithProgress(res, url, onProgress) {
+  const total = Number(res.headers?.get?.("content-length")) || 0;
+  const file = String(url).split("/").pop();
+  if (res.body?.getReader) {
+    const reader = res.body.getReader();
+    let received = 0;
+    let data = total ? new Uint8Array(total) : null;
+    const chunks = total ? null : [];
+    for (;;) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      if (data) {
+        if (received + value.length > data.length) {
+          // content-length can describe the *encoded* size (gzip/br), so the
+          // decoded stream may be larger — grow instead of overflowing.
+          const grown = new Uint8Array(Math.max(data.length * 2, received + value.length));
+          grown.set(data.subarray(0, received));
+          data = grown;
+        }
+        data.set(value, received);
+      } else {
+        chunks.push(value);
+      }
+      received += value.length;
+      onProgress?.({ phase: "download", file, loaded: received, total: total || null });
+    }
+    if (!data) {
+      data = new Uint8Array(received);
+      let off = 0;
+      for (const c of chunks) { data.set(c, off); off += c.length; }
+    } else {
+      // the buffer may have grown past the decoded length (content-length is
+      // the *encoded* size) — always trim to what was actually received
+      data = data.subarray(0, received);
+    }
+    return data;
+  }
+  let data;
+  if (typeof res.arrayBuffer === "function") {
+    data = new Uint8Array(await res.arrayBuffer());
+  } else if (typeof res.text === "function") {
+    data = new TextEncoder().encode(await res.text());
+  } else if (typeof res.json === "function") {
+    data = new TextEncoder().encode(JSON.stringify(await res.json()));
+  } else {
+    throw new Error(`fetch ${url}: response has no readable body`);
+  }
+  onProgress?.({ phase: "download", file, loaded: data.length, total: data.length });
+  return data;
+}
+
+/**
+ * ONNX external-data loader.
+ *
+ * The pinned Laya export stores weights in `<model>.onnx.data` siblings
+ * (encoder ~468 MB, head ~53 MB). onnxruntime-web cannot resolve those from a
+ * URL by itself (`Module.MountedFiles is not available`); it needs the bytes
+ * passed via `session_options.externalData`. Fetch the sibling (streaming,
+ * with progress) and hand it over.
+ *
+ * @returns {Promise<{path:string, data:Uint8Array}|null>} null when the model
+ *   has no external data (404) or the fetch fails.
+ */
+export async function fetchOnnxExternalData(url, { fetchImpl = fetch, onProgress = null, force = false } = {}) {
+  const dataUrl = `${url}.data`;
+  const path = `${String(url).split("/").pop()}.data`;
+  try {
+    // Probe cheaply first: a 404 means the model has no external data.
+    const head = await fetchImpl(dataUrl, { method: "HEAD" }).catch(() => null);
+    if (head && !head.ok) return null;
+    const data = await fetchCachedBytes(dataUrl, { fetchImpl, onProgress, force, persist: true });
+    return { path, data };
+  } catch {
+    return null;
+  }
+}
+
 /**
  * Create a Laya decider.
  *
@@ -256,6 +446,7 @@ export async function createLayaDecider({
   tokenizerSubfolder = null,
   ort = null,
   transformers = null,
+  fetchImpl = fetch,
   wasmPaths = null,
   device = "auto",
   // (url, sessionOptions) => Promise<InferenceSession>. Hosts that need custom
@@ -283,21 +474,40 @@ export async function createLayaDecider({
   const report = (phase) => (p) =>
     onProgress?.({ phase, file: p?.file ?? null, loaded: p?.loaded ?? null, total: p?.total ?? null });
 
-  const tokModel = tokenizerModel ?? model;
-  const tokSubfolder = tokenizerSubfolder === null
-    ? undefined
-    : (tokenizerSubfolder ?? subfolder);
-  const tok = await tf.AutoTokenizer.from_pretrained(tokModel, {
-    revision: tokenizerModel ? undefined : revision,
-    ...(tokSubfolder ? { subfolder: tokSubfolder } : {}),
-    progress_callback: report("tokenizer"),
-  });
+  let tok;
+  if (tokenizerModel) {
+    // Host-provided path/repo (local dir in Node, or a repo whose tokenizer
+    // lives at the root) — use the library loader.
+    const tokSubfolder = tokenizerSubfolder === null
+      ? undefined
+      : (tokenizerSubfolder ?? subfolder);
+    tok = await tf.AutoTokenizer.from_pretrained(tokenizerModel, {
+      revision: tokenizerSubfolder ? undefined : revision,
+      ...(tokSubfolder ? { subfolder: tokSubfolder } : {}),
+      progress_callback: report("tokenizer"),
+    });
+  } else {
+    // Default: the pinned web export keeps the tokenizer under `v1/`, which
+    // AutoTokenizer cannot address (no `subfolder` support) — load it directly.
+    tok = await loadLayaTokenizer({
+      urlBase: hfUrl(model, revision, `${subfolder}/`),
+      transformers: tf,
+      fetchImpl,
+      onProgress,
+    });
+  }
   const adapter = layaTokenizerAdapter(tok);
 
   const encUrl = encoderUrl ?? hfUrl(model, revision, `${subfolder}/encoder_q8.onnx`);
   const headUrlResolved = headUrl ?? hfUrl(model, revision, `${subfolder}/head_q8.onnx`);
   const createSession = sessionFactory
-    ?? ((url, opts) => ortMod.InferenceSession.create(url, opts));
+    ?? (async (url, opts) => {
+      const report2 = (p) => onProgress?.(p);
+      const modelBytes = await fetchCachedBytes(url, { fetchImpl, onProgress: report2, persist: true });
+      const external = await fetchOnnxExternalData(url, { fetchImpl, onProgress: report2 });
+      const sessionOpts = external ? { ...opts, externalData: [external] } : opts;
+      return ortMod.InferenceSession.create(modelBytes, sessionOpts);
+    });
   const loadPair = async (providers) => {
     const opts = { executionProviders: providers, graphOptimizationLevel: "all" };
     const enc = await createSession(encUrl, { ...opts, progress_callback: report("encoder") });
@@ -341,22 +551,27 @@ export async function createLayaDecider({
         Object.entries(temperatureByOptions).map(([k, v]) => [k, clampTemperature(v)])),
     },
     async decide(state, questions) {
+      const t0 = now();
       const normalized = normalizeQuestions(questions);
       const items = normalized.map((q) =>
         buildLayaSequence({ encode: adapter.encode, ids: adapter.ids, state, question: q, maxLen, headMaxLen, maxPrefixes }));
       const b = collateLayaItems(items, adapter.pad);
 
+      const tEncode = now();
       const encOut = await encoder.run({
         input_ids: int64(b.inputIds, [b.n, b.L]),
         attention_mask: int64(b.attention, [b.n, b.L]),
       });
+      const encoderMs = now() - tEncode;
       const hidden = encOut.hidden ?? encOut.last_hidden_state ?? encOut[Object.keys(encOut)[0]];
 
       // The q8 web head is exported batch-1 only (despite dynamic batch
       // axes), so run it once per question over the shared encoder output.
       // The encoder — the expensive half — still runs as one batched pass.
       const headLogits = [];
+      const headPerQuestionMs = [];
       const markerCounts = items.map((it) => it.markers.length);
+      const tHead = now();
       for (let i = 0; i < b.n; i++) {
         const k = markerCounts[i];
         const row = hidden.data.subarray(i * b.L * hidden.dims[2], (i + 1) * b.L * hidden.dims[2]);
@@ -364,6 +579,7 @@ export async function createLayaDecider({
         const rowMask = b.attention.subarray(i * b.L, (i + 1) * b.L);
         const rowPos = b.markerPos.subarray(i * b.M, i * b.M + k);
         const rowMarkerMask = b.markerMask.subarray(i * b.M, i * b.M + k);
+        const tHeadTotal = now();
         const headOut = await head.run({
           hidden: rowHidden,
           attention_mask: int64(rowMask, [1, b.L]),
@@ -373,6 +589,7 @@ export async function createLayaDecider({
         });
         const rowLogits = headOut.logits ?? headOut[Object.keys(headOut)[0]];
         headLogits.push(...Array.from(rowLogits.data, Number));
+        headPerQuestionMs.push(Math.round(now() - tHeadTotal));
       }
 
       const answers = layaAnswersFromLogits({
@@ -386,11 +603,18 @@ export async function createLayaDecider({
       // Exact, human-readable view of what the encoder consumed: each
       // question's token sequence decoded back to text (special tokens kept).
       const prompts = items.map((it) => adapter.decode(it.ids));
+      const headMs = now() - tHead;
       return {
         answers,
         truncated: items.some((it) => it.ids.length >= maxLen),
         length,
         prompts,
+        timings: {
+          totalMs: Math.round(now() - t0),
+          encoderMs: Math.round(encoderMs),
+          headMs: Math.round(headMs),
+          headPerQuestionMs,
+        },
       };
     },
     async dispose() {

package/src/registry.mjs

@@ -0,0 +1,102 @@
+/**
+ * registry.mjs — multi-family runtime for JEV-style typed-decision models.
+ *
+ * jev-web ships two families today (open-jev's fused DeBERTa graph and Laya's
+ * ModernBERT encoder + typed head) and any third party can add its own without
+ * touching core:
+ *
+ *   registerDecisionFamily("my-model", {
+ *     defaults: { model: "org/my-typed-decision-ONNX", revision: "main" },
+ *     create: async (opts) => ({ info, decide(state, questions), dispose? }),
+ *   });
+ *
+ *   const runtime = await createDecisionRuntime({ family: "my-model" });
+ *
+ * `createDecisionRuntime` tries the requested family, then any `fallback`
+ * families in order, reporting each failure through `onFallback`. The returned
+ * object always exposes the shared contract:
+ *   { info: { family, ... }, decide(state, questions) -> DecideResult, dispose? }
+ * where DecideResult is `{ answers, truncated, length, prompts?, timings? }`.
+ */
+
+import { createDecider } from "./session.mjs";
+import { DEFAULT_MODEL, DEFAULT_REVISION } from "./config.mjs";
+import {
+  createLayaDecider,
+  LAYA_DEFAULT_MODEL,
+  LAYA_DEFAULT_REVISION,
+  LAYA_DEFAULT_SUBFOLDER,
+} from "./laya.mjs";
+
+const families = new Map();
+
+/**
+ * Register (or replace) a decision family.
+ * @param {string} id short stable identifier, e.g. "open-jev"
+ * @param {{create:(opts:object)=>Promise<object>, defaults?:object}} descriptor
+ */
+export function registerDecisionFamily(id, descriptor) {
+  if (!id || typeof id !== "string") throw new TypeError("registerDecisionFamily: id is required");
+  if (!descriptor || typeof descriptor.create !== "function") {
+    throw new TypeError(`registerDecisionFamily(${id}): descriptor.create is required`);
+  }
+  families.set(id, { defaults: {}, ...descriptor });
+  return id;
+}
+
+/** @returns {string[]} registered family ids in registration order. */
+export function listDecisionFamilies() {
+  return [...families.keys()];
+}
+
+export function getDecisionFamily(id) {
+  return families.get(id) ?? null;
+}
+
+registerDecisionFamily("open-jev", {
+  defaults: { model: DEFAULT_MODEL, revision: DEFAULT_REVISION },
+  create: createDecider,
+});
+
+registerDecisionFamily("laya", {
+  defaults: {
+    model: LAYA_DEFAULT_MODEL,
+    revision: LAYA_DEFAULT_REVISION,
+    subfolder: LAYA_DEFAULT_SUBFOLDER,
+  },
+  create: createLayaDecider,
+});
+
+/**
+ * Create a decision runtime, preferring `family` and falling back in order.
+ * @param {{family?:string, fallback?:string[], onFallback?:Function}} [opts]
+ * @returns {Promise<{info:object, decide:Function, dispose?:Function}>}
+ */
+export async function createDecisionRuntime({
+  family = "laya",
+  fallback = [],
+  onFallback = null,
+  ...opts
+} = {}) {
+  const order = [family, ...(Array.isArray(fallback) ? fallback : [])].filter(Boolean);
+  let lastError = null;
+  for (const id of order) {
+    const descriptor = families.get(id);
+    if (!descriptor) {
+      lastError = new Error(`unknown decision family "${id}"`);
+      onFallback?.({ family: id, error: lastError });
+      continue;
+    }
+    try {
+      const runtime = await descriptor.create({ ...descriptor.defaults, ...opts });
+      return {
+        ...runtime,
+        info: { family: id, ...(runtime?.info ?? {}) },
+      };
+    } catch (error) {
+      lastError = error;
+      onFallback?.({ family: id, error });
+    }
+  }
+  throw lastError ?? new Error("createDecisionRuntime: no family could be created");
+}

package/src/session.mjs

@@ -5,6 +5,8 @@
 // builds). Everything heavy stays out of module scope.
 
 import { normalizeQuestions } from "./questions.mjs";
+
+const now = () => (typeof performance !== "undefined" ? performance.now() : Date.now());
 import { buildDecisionInput, typedMarkerIds } from "./encode.mjs";
 import { answersFromScores } from "./answers.mjs";
 import {
@@ -77,6 +79,7 @@ export async function createDecider({
 
   return {
     info: {
+      family: "open-jev",
       model,
       revision,
       device: resolvedDevice,
@@ -88,6 +91,7 @@ export async function createDecider({
       configSource: remote.source,
     },
     async decide(state, questions) {
+      const t0 = now();
       const normalized = normalizeQuestions(questions);
       const tokenized = normalized.map((q) => ({
         type: q.type,
@@ -109,7 +113,9 @@ export async function createDecider({
         pair_q: i64(input.pairQ, [1, input.pairQ.length]),
         pair_opt: i64(input.pairOpt, [1, input.pairOpt.length]),
       };
+      const tInfer = now();
       const out = await net(feeds);
+      const inferMs = now() - tInfer;
       const logits = out?.logits;
       if (!logits) throw new Error("model output has no `logits` — not a typed-decision export?");
       const scores = Array.from(logits.to("float32").data, Number);
@@ -125,7 +131,13 @@ export async function createDecider({
       } catch {
         prompt = input.inputIds.join(" ");
       }
-      return { answers, truncated: input.truncated, length: input.length, prompts: [prompt] };
+      return {
+        answers,
+        truncated: input.truncated,
+        length: input.length,
+        prompts: [prompt],
+        timings: { totalMs: Math.round(now() - t0), inferMs: Math.round(inferMs) },
+      };
     },
     async decideMany(cases) {
       const out = [];