Something went wrong. Try again.
a tiny pet with a GPT-2 brain, living in your browser
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304import type { Generated, Generation, WorkerRequest, WorkerResponse } from "../types.ts";import { CFG } from "./config.ts";import { pump, scheduleIdle } from "./conversation.ts";import { debugTurn } from "./debug.ts";import { keepOffline } from "./pwa.ts";import { brain, emit, store, toast, ui, upgrade } from "./store.ts";import { errorMessage, fmtMB } from "./utils.ts";
let genId = 0;const pending = new Map< number, { resolve: (result: Generated) => void; reject: (error: Error) => void; onToken?: (text: string) => void; }>();// What transformers.js fetches for a model, and where it keeps it: Cache// Storage, keyed by the absolute file URL.export const BRAIN_CACHE = "transformers-cache";export const MODEL_FILES = [ "tokenizer_config.json", "tokenizer.json", "config.json", "generation_config.json", "onnx/model_q4.onnx",];// The ONNX Runtime build the vendored transformers.js loads (asyncify, 26.9 MB),// fetched once alongside the first brain.export const RUNTIME_BYTES = 27e6;export const modelsBase = new URL(CFG.models, location.href).href;const modelUrl = (id: string, f: string) => `${modelsBase}${id}/${f}`;// Brains this app can talk to, newest first. Older ones speak the same prompt// format but may lack features trained into newer ones.export const BRAINS = [ { id: CFG.modelId, size: CFG.size, features: CFG.features || [] }, ...(CFG.older || []),];const brainInfo = (id: string | null) => BRAINS.find((b) => b.id === id);// Feature name -> what it lets the pet do, for the upgrade dialog. The page// config lists which brains have which features.
export const brainHas = (f?: string) => !f || (brain.id ? brainInfo(brain.id)?.features || [] : BRAINS[0].features).includes(f);export const olderBrain = () => !!brain.id && brain.id !== CFG.modelId;
export async function cachedBrains() { try { if (!("caches" in self)) return []; const keys = await (await caches.open(BRAIN_CACHE)).keys(); return BRAINS.map((b) => b.id).filter((id) => keys.some((r) => r.url === modelUrl(id, "onnx/model_q4.onnx")), ); } catch { return []; }}
// Delete model files of every brain except the ones in keep (only under this// page's models URL; the cache holds nothing else of ours).export async function forgetBrains(keep: string[]) { try { const cache = await caches.open(BRAIN_CACHE); for (const r of await cache.keys()) { if (!r.url.startsWith(modelsBase)) continue; const id = r.url.slice(modelsBase.length).split("/")[0]; if (!keep.includes(id)) await cache.delete(r); } } catch { /* optional */ } void updateStorage();}
// The upgrade to CFG.modelId downloads into the cache while the older brain// keeps running, then swaps brains. Nothing about the pet changes.
async function fetchBrain(id: string, onBytes: (bytes: number) => void) { const cache = await caches.open(BRAIN_CACHE); for (const f of MODEL_FILES) { const url = modelUrl(id, f); if (await cache.match(url)) continue; const res = await fetch(url); if (!res.ok || !res.body) throw new Error(`${f}: HTTP ${res.status}`); const [keep, count] = res.body.tee(); const put = cache.put(url, new Response(keep, { status: res.status, headers: res.headers })); const reader = count.getReader(); for (;;) { const { done, value } = await reader.read(); if (done) break; onBytes(value.length); } await put; }}export async function startUpgrade() { closeUpgrade(); if (upgrade.state === "downloading") return; if (!brain.worker) { loadBrain(CFG.modelId); return; } // nothing running: a plain download Object.assign(upgrade, { state: "downloading", loaded: 0, error: "" }); try { navigator.storage?.persist?.().catch(() => {}); } catch { /* optional */ } emit(); let shown = 0; try { await fetchBrain(CFG.modelId, (n) => { upgrade.loaded += n; if (upgrade.loaded - shown > 1e6) { shown = upgrade.loaded; emit(); } }); } catch (e) { Object.assign(upgrade, { state: "error", error: errorMessage(e) }); emit(); return; } // Swap once the old brain is between turns. while (brain.busy) await new Promise((r) => setTimeout(r, 200)); brain.worker?.terminate(); clearTimeout(watchdog); Object.assign(brain, { worker: null, state: "off", busy: false, progress: null, id: null }); upgrade.state = "swapping"; loadBrain(CFG.modelId);}
export function onMobileData() { const c = navigator.connection; return !!c && (c.saveData || c.type === "cellular");}
let upgradeOffered = false;export function offerUpgrade() { if (upgradeOffered || !olderBrain() || upgrade.state !== "idle") return; upgradeOffered = true; ui.upgradeOpen = true; emit();}export function closeUpgrade() { ui.upgradeOpen = false; emit();}// If the runtime stalls (no progress, no error, no ready), retry once with a// single thread, which needs no SharedArrayBuffer, then give up.let watchdog = 0;function armWatchdog() { clearTimeout(watchdog); watchdog = setTimeout(() => { if (brain.state === "ready" || brain.state === "off" || !brain.worker) return; console.warn("brain load stalled", brain.threads ? "on one thread" : ""); brain.worker.terminate(); brain.worker = null; if (brain.threads === 1) { brain.state = "error"; brain.error = "the model runtime did not start"; emit(); return; } brain.threads = 1; loadBrain(brain.id || CFG.modelId); }, 90e3);}
export async function loadBrain(id = CFG.modelId) { if (brain.worker || !store.tabOwned) return; // Firefox may leave its persistence permission prompt unanswered. Storage // persistence is optional and must not hold up the download (or a retry). try { navigator.storage?.persist?.().catch(() => {}); } catch { /* optional */ } Object.assign(brain, { state: "downloading", id, progress: null, error: "" }); emit(); let w: Worker | null = null; const fail = async (message: string) => { clearTimeout(watchdog); w?.terminate(); for (const p of pending.values()) p.reject(new Error(message)); pending.clear(); // An interrupted body can leave a padded, invalid ONNX in the runtime's // cache. Retry must fetch that file again, without discarding other brains // or a valid model after unrelated runtime failures (e.g. out of memory). if (/error in input stream|protobuf parsing failed/i.test(message)) { try { await (await caches.open(BRAIN_CACHE)).delete(modelUrl(id, "onnx/model_q4.onnx")); } catch { /* optional */ } } if (brain.worker !== w) return; Object.assign(brain, { worker: null, busy: false, state: "error", error: message }); emit(); }; try { w = new Worker(new URL("./brain.js", import.meta.url), { type: "module" }); brain.worker = w; } catch (e) { await fail(errorMessage(e) || "worker failed"); return; } w.onmessage = ({ data }: MessageEvent<WorkerResponse>) => { if (brain.worker !== w) return; if (data.type === "progress") { brain.progress = data; armWatchdog(); // fires only after 90 s without progress or ready emit(); } else if (data.type === "loading") { brain.state = "warming"; armWatchdog(); emit(); } else if (data.type === "ready") { clearTimeout(watchdog); brain.state = "ready"; if (upgrade.state === "swapping") { upgrade.state = "idle"; toast(`${store.state.pet?.name || "Your pet"} has its new brain!`, "sparkle"); } // Keep only the brain that runs (and the current one, if it is half there). forgetBrains(olderBrain() ? [id, CFG.modelId] : [id]); emit(); scheduleIdle(4000); pump(); keepOffline(data.resources || []); if (olderBrain()) setTimeout(offerUpgrade, 1500); } else if (data.type === "token") pending.get(data.id)?.onToken?.(data.text); else if (data.type === "done") { pending.get(data.id)?.resolve(data); pending.delete(data.id); } else if (data.type === "error") { if (data.id && pending.has(data.id)) { pending.get(data.id)!.reject(new Error(data.message)); pending.delete(data.id); } else fail(data.message); } }; w.onerror = (e) => { if (brain.worker === w) fail(e.message || "worker failed"); }; try { w.postMessage({ type: "load", runtime: CFG.runtime, models: CFG.models, modelId: id, dtype: "q4", threads: brain.threads || undefined, } satisfies WorkerRequest); armWatchdog(); } catch (e) { await fail(errorMessage(e) || "worker failed"); }}
export function gen(opts: Generation, onToken?: (text: string) => void): Promise<Generated> { const id = ++genId; return new Promise<Generated>((resolve, reject) => { pending.set(id, { resolve: (r) => { debugTurn(opts, r); resolve(r); }, reject, onToken, }); brain.worker!.postMessage({ type: "generate", id, ...opts } satisfies WorkerRequest); });}
export async function updateStorage() { const out: string[] = []; try { if (navigator.storage?.estimate) out.push( `This site uses ${fmtMB((await navigator.storage.estimate()).usage || 0)} of browser storage.`, ); if (navigator.storage?.persisted) out.push( (await navigator.storage.persisted()) ? "Your browser keeps it until you delete it." : "Your browser may clear it if you don't visit for a while; installing the app or exporting a save keeps your pet safe.", ); } catch { /* optional */ } ui.storage = out.join(" "); emit();}export function stopBrain() { clearTimeout(watchdog); brain.worker?.terminate(); for (const p of pending.values()) p.reject(new Error("brain stopped")); pending.clear(); Object.assign(brain, { worker: null, state: "off", busy: false, progress: null, id: null }); emit();}export const getBrain = () => (olderBrain() ? startUpgrade() : loadBrain());