Something went wrong. Try again.
a tiny pet with a GPT-2 brain, living in your browser
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156import { fit, initCore, stop } from "./core.ts";import type { Model, Progress, Tokenizer, Transformers } from "./runtime-types.ts";import type { Generation, WorkerRequest, WorkerResponse } from "./types.ts";// Tokenling brain: a module worker that runs the fine-tuned GPT-2 (small, 4-bit)// 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 = <K extends WorkerResponse["type"]>( type: K, data: Omit<Extract<WorkerResponse, { type: K }>, "type">,) => self.postMessage({ type, ...data });
async function load({ runtime, models, modelId, dtype, threads,}: Extract<WorkerRequest, { type: "load" }>) { // Stop rules and context fitting come from the shared Rust rules (core/). await initCore(); T = await import(`${runtime}/transformers.min.js`); if (threads) T.env.backends.onnx.wasm.numThreads = threads; T.env.allowRemoteModels = false; T.env.allowLocalModels = true; T.env.localModelPath = models; T.env.useBrowserCache = true; T.env.backends.onnx.wasm.wasmPaths = `${runtime}/`; T.env.backends.onnx.logLevel = "error"; const files = new Map<string, [number, number]>(); 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 }); // The model stores 4-bit weights behind DequantizeLinear nodes. // disable_quant_qdq makes ONNX Runtime fold them to fp32 at load instead of // fusing them into MatMulNBits, which is about 5x slower in WebAssembly // (train/quantize.py); without it the model does not load in a usable way. // disable_prepacking skips a second packed copy of every weight (-160 MB). model = await T.AutoModelForCausalLM.from_pretrained(modelId, { device: "wasm", dtype, progress_callback, session_options: { extra: { session: { disable_quant_qdq: "1", disable_prepacking: "1" } } }, }); 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<WorkerRequest, { type: "generate" }>) { 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; // 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<WorkerRequest>) => { try { if (data.type === "load") { await load(data); // 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), }); }};