import { fit, initCore, stop } from "./core.ts"; import type { Model, Progress, Tokenizer, Transformers } from "./runtime-types.ts"; import type { Generation, WorkerRequest, WorkerResponse } from "./types.ts"; // Subtokenling brain: a module worker that runs the fine-tuned language model // with transformers.js on WebAssembly. Weights are fetched once from this site // and kept in the browser's Cache Storage; nothing is sent anywhere. let T: Transformers, tokenizer: Tokenizer, model: Model, busy = false; const post = ( type: K, data: Omit, "type">, ) => self.postMessage({ type, ...data }); // This handshake exists only on an explicitly selected private worker URL. // It retains shapes/scalars, never tensor objects or graph byte buffers. const memoryProfile = new URLSearchParams(self.location.search).get("profile-memory") === "1"; let releaseProfile: (() => void) | undefined; let lastForward: object | null = null; const profile = (value: object) => self.postMessage({ type: "private-memory", ...value }); function tensorShapes(value: unknown, depth = 0): Record { if (!value || typeof value !== "object" || depth > 3) return {}; return Object.fromEntries( Object.entries(value).flatMap<[string, unknown]>(([name, tensor]) => { if (!tensor || typeof tensor !== "object") return []; if ("dims" in tensor && Array.isArray(tensor.dims)) return [[name, { dims: [...tensor.dims], type: "type" in tensor ? tensor.type : null }]]; const nested = tensorShapes(tensor, depth + 1); return Object.keys(nested).length ? [[name, nested]] : []; }), ); } async function load({ runtime, models, modelId, dtype, threads, }: Extract) { // Stop rules and context fitting come from the shared Rust rules (core/). await initCore(); T = await import(`${runtime}/transformers.min.js`); if (memoryProfile) { await new Promise((resolve) => { releaseProfile = resolve; profile({ phase: "runtime_imported_no_session" }); }); } if (threads) T.env.backends.onnx.wasm.numThreads = threads; // 4.3.1's metadata registry skips HTTP localModelPath values. Use the // configured static model origin as its remote host instead: files still // come only from this deployment's models URL, never Hugging Face or an // inference API. AutoTokenizer retains each model's own tokenizer class. T.env.allowRemoteModels = true; T.env.allowLocalModels = false; T.env.remoteHost = models; T.env.remotePathTemplate = "{model}/"; T.env.useBrowserCache = true; T.env.backends.onnx.wasm.wasmPaths = `${runtime}/`; T.env.backends.onnx.logLevel = "error"; const files = new Map(); const progress_callback = (p: Progress) => { if (p.status === "progress" && p.total) { files.set(p.file, [p.loaded, p.total]); let loaded = 0, total = 0; for (const [l, t] of files.values()) { loaded += l; total += t; } post("progress", { loaded, total, file: p.file }); } else if (p.status === "done") { post("file", { file: p.file }); } }; tokenizer = await T.AutoTokenizer.from_pretrained(modelId, { progress_callback }); // q4 graphs store weights behind DequantizeLinear nodes. disable_quant_qdq // lets ONNX Runtime fold them to fp32 instead of fusing MatMulNBits, which // is much slower in WebAssembly (train/quantize.py). The same session flags // also work for q8/fp32; disable_prepacking avoids a second packed weight copy. model = await T.AutoModelForCausalLM.from_pretrained(modelId, { device: "wasm", dtype, progress_callback, session_options: { extra: { session: { disable_quant_qdq: "1", disable_prepacking: "1" } } }, }); if (memoryProfile) { if (!("forward" in model) || typeof model.forward !== "function") throw new Error("private memory profile cannot observe model.forward cache tensors"); const forward = model.forward; model.forward = async (...args: unknown[]) => { const result = await forward.apply(model, args); lastForward = { inputs: tensorShapes(args[0]), outputs: tensorShapes(result) }; return result; }; profile({ phase: "session_loaded", model_config: "config" in model ? model.config : null, effective_threads: T.env.backends.onnx.wasm.numThreads || null, }); } post("loading", { stage: "warming up" }); await model.generate({ ...tokenizer("hi"), max_new_tokens: 1, do_sample: false }); } // Stop at end of line (chat) or right after a lookup tag so the app can // append its {result}; diaries stop at a new header. The rule is shared with // the host's native brain (core/src/world.rs `stop`). function makeStopper( promptLen: number, mode: Generation["mode"], state: { text: string; reason: string }, ) { class Stop extends T.StoppingCriteria { _call(input_ids: (number | bigint)[][]) { const ids = input_ids[0].slice(promptLen).map(Number); state.text = tokenizer.decode(ids, { skip_special_tokens: true }); const stopped = stop(mode, state.text); if (!stopped) return [false]; state.reason = stopped.reason; state.text = stopped.text; return [true]; } } return new Stop(); } async function generate(request: Extract) { const { id, mode, maxNew, temperature, topK, topP } = request; // Drops the oldest chat lines until the prompt fits; idle turns pass a // smaller ctx because prefill dominates their cost. const prompt = fit(request, (text) => tokenizer.encode(text).length); const inputs = tokenizer(prompt); const promptLen = inputs.input_ids.dims[1]; const state = { text: "", reason: "length" }; let first = 0; const streamer = new T.TextStreamer(tokenizer, { skip_prompt: true, callback_function: () => { first ||= performance.now(); post("token", { id, text: state.text }); }, }); const t0 = performance.now(); await model.generate({ ...inputs, max_new_tokens: maxNew, do_sample: true, temperature, top_k: topK, top_p: topP, stopping_criteria: [makeStopper(promptLen, mode, state)], streamer, }); const text = state.text; if (memoryProfile) profile({ phase: "generation", id, context: request.ctx, prompt, prompt_tokens: promptLen, requested_history_lines: request.lines.length, max_new_tokens: maxNew, cache_tensors: lastForward, cache_observation_scope: "last actual model.forward input/output tensor dimensions, including nested KV; no tensor objects retained", }); // prefillMs: prompt processing up to the first new token; tps: decoding speed after it. const end = performance.now(), n = tokenizer.encode(text).length; const prefillMs = (first || end) - t0; post("done", { id, text, prompt, reason: state.reason, promptTokens: promptLen, prefillMs, tps: n > 1 && first ? (n - 1) / ((end - first) / 1000) : 0, }); } self.onmessage = async ({ data, }: MessageEvent) => { if (data.type === "private-memory-continue") { if (memoryProfile) { releaseProfile?.(); releaseProfile = undefined; } return; } try { if (data.type === "load") { await load(data); if (memoryProfile) profile({ phase: "loaded_warmed", cache_tensors: lastForward, resource_timing: performance.getEntriesByType("resource").map((entry) => { // The resource entry-type filter supplies the browser's resource subtype. const resource = entry as PerformanceResourceTiming; return { name: resource.name, decoded_bytes: resource.decodedBodySize, encoded_bytes: resource.encodedBodySize, transfer_bytes: resource.transferSize, }; }), }); // The runtime files this worker fetched, so the page can keep them for offline use. post("ready", { dtype: data.dtype, resources: performance.getEntriesByType("resource").map((e) => e.name), }); } else if (data.type === "generate") { if (busy) return post("error", { id: data.id, message: "busy" }); busy = true; try { await generate(data); } finally { busy = false; } } } catch (err) { post("error", { id: data.type === "generate" ? data.id : undefined, message: err instanceof Error ? err.message : String(err), }); } };