aliasweb.mjs · aliasrobotics/alias3-web at main

alias3-web / aliasweb.mjs

alias3-web: 4-bit WebGPU graph, readout files, library and demo

2ab2e60 verified about 12 hours ago

10.1 kB

// alias3-web (alias-web in the paper) · a cybersecurity transformer in the browser (WebGPU).
//
// Two regimes from one loaded 4-bit graph:
//   System 1  ask()       one typed question, one prefill, no token decoded: a probability per option
//   System 2  generate()  text from the same weights, token by token, on request
//
// The prefill engine is the vendored runner (vendor/ct0.js: ct0-edge by Eduardo Worrel, MIT, see
// vendor/ct0-edge.LICENSE), unchanged. This file adds the readout the model was trained with:
//
//   · the System 1 markers as the model's own special tokens, by id:
//       <state> = <|fim_prefix|>   <q> = <|fim_middle|>   <opt> = <|box_start|>
//       </opt>  = <|box_end|>      <decide> = <|fim_suffix|>
//   · the final RMSNorm (heads/norm.bin): the graph exports hidden states before it
//   · the complete pointer head (heads/head.bin), read as in training:
//       z_i = q(h_<decide>) · k(h_</opt>_i) / (sqrt(256) · T),  softmax over the options
//   · yes/no and severity as choices over ("Yes", "No") and ("1" … "5"): one mechanism for every question
//
// The page that imports this module needs an import map for the runner and its two runtimes; see index.html.

export { AliasWeb };

const MARKERS = {
  state: "<|fim_prefix|>",
  q: "<|fim_middle|>",
  opt: "<|box_start|>",
  optEnd: "<|box_end|>",
  decide: "<|fim_suffix|>",
};
const LEVELS = ["1", "2", "3", "4", "5"];
export const DEFAULT_SCORE_Q = "How severe is the security risk in this content?";

// ── head.bin: [u32 lenJSON][json][fp16 q][fp16 qb][fp16 k][fp16 kb] ────────
function f16(h) {
  const s = (h & 0x8000) >> 15, e = (h & 0x7c00) >> 10, f = h & 0x03ff;
  if (e === 0) return (s ? -1 : 1) * Math.pow(2, -14) * (f / 1024);
  if (e === 0x1f) return f ? NaN : s ? -Infinity : Infinity;
  return (s ? -1 : 1) * Math.pow(2, e - 15) * (1 + f / 1024);
}
export function loadHead(buf) {
  const dv = new DataView(buf);
  const jl = dv.getUint32(0, true);
  const meta = JSON.parse(new TextDecoder().decode(new Uint8Array(buf, 4, jl)));
  let off = 4 + jl;
  const take = (n) => { const f = new Float32Array(n); for (let i = 0; i < n; i++) f[i] = f16(dv.getUint16(off + i * 2, true)); off += n * 2; return f; };
  const d = meta.headDim, h = meta.hiddenSize;
  return { q: take(d * h), qb: take(d), k: take(d * h), kb: take(d), headDim: d, hiddenSize: h,
    temperature: meta.temperature ?? 1.0, scale: meta.scale ?? 1 / Math.sqrt(d), scoreOffset: meta.score_offset ?? 0 };
}
// norm.bin: [u32 lenJSON][json {hiddenSize}][fp32 hiddenSize]
async function loadNorm(buf) {
  const dv = new DataView(buf);
  const jl = dv.getUint32(0, true);
  const meta = JSON.parse(new TextDecoder().decode(new Uint8Array(buf, 4, jl)));
  const n = meta.hiddenSize, off = 4 + jl, f = new Float32Array(n);
  for (let i = 0; i < n; i++) f[i] = dv.getFloat32(off + i * 4, true);
  return f;
}
function proj(W, b, x, dim, hidden) {
  const y = new Float32Array(dim);
  for (let d = 0; d < dim; d++) { let a = b[d]; const r = d * hidden; for (let i = 0; i < hidden; i++) a += W[r + i] * x[i]; y[d] = a; }
  return y;
}
const dot = (a, b) => { let s = 0; for (let i = 0; i < a.length; i++) s += a[i] * b[i]; return s; };
const headLogit = (hd, hD, hO) => dot(proj(hd.q, hd.qb, hD, hd.headDim, hd.hiddenSize), proj(hd.k, hd.kb, hO, hd.headDim, hd.hiddenSize)) * hd.scale / hd.temperature + (hd.scoreOffset || 0);
const softmax = (xs) => {
  const mx = Math.max(...xs.filter(Number.isFinite));
  const e = xs.map((x) => (Number.isFinite(x) ? Math.exp(x - mx) : 0));
  const z = e.reduce((a, b) => a + b, 0) || 1;
  return e.map((v) => v / z);
};

// text supplied by the caller can never produce a delimiter: <|name|> becomes <¦name¦>, as in training
const SPECIAL_RE = /<\|([A-Za-z0-9_]+)\|>/g;
const sanitize = (t) => t.replace(SPECIAL_RE, "<¦$1¦>");

class AliasWeb {
  constructor(tok, prefill, head, mid, normW) { this.tok = tok; this.prefill = prefill; this.head = head; this.mid = mid; this.normW = normW; }

  // the graph exports hidden_states before the final norm; the head was trained on the normalised
  // state, so the model's final RMSNorm is applied here
  norm(h) {
    if (!this.normW) return h;
    let s = 0; for (let i = 0; i < h.length; i++) s += h[i] * h[i];
    const inv = 1 / Math.sqrt(s / h.length);
    const o = new Float32Array(h.length);
    for (let i = 0; i < h.length; i++) o[i] = h[i] * inv * this.normW[i];
    return o;
  }

  static async load(opts = {}) {
    const base = opts.base ?? "./";
    const status = opts.status ?? (() => { });
    status("loading the runner and the 4-bit graph (2.4 GB)…");
    const edge = await import("ct0-edge");
    const inner = await edge.CT0.load({
      runtime: "genai",
      genaiBaseUrl: base + "model/",
      genaiModel: "model.onnx",
      baseUrl: base + "model/",
      status: () => { },
      verifyEp: false,          // placement was verified once against this graph in Chrome
      skipVersionCheck: true,   // the import map pins the ONNX Runtime version
      genai: { warmup: false },
    });
    const be = inner.backend ?? inner.be;
    if (!be?.prefill) throw new Error("the runner did not expose its prefill backend");
    status("loading the head and the final norm…");
    const head = loadHead(await (await fetch(base + "heads/head.bin")).arrayBuffer());
    const normW = await loadNorm(await (await fetch(base + "heads/norm.bin")).arrayBuffer());
    // marker ids from the tokenizer's added_tokens
    const tk = await (await fetch(base + "model/tokenizer.json")).json();
    const added = new Map((tk.added_tokens ?? []).map((t) => [t.content, t.id]));
    const mid = {};
    for (const [k, m] of Object.entries(MARKERS)) {
      const id = added.get(m);
      if (typeof id !== "number") throw new Error("marker without a token id: " + m);
      mid[k] = id;
    }
    const model = new AliasWeb(inner.tok?.tokenizer ?? inner.tok, (ids, positions) => be.prefill(ids, { positions }), head, mid, normW);
    model.be = be; model.base = base;
    status("warm-up…");
    const w = model.encode("warmup", "ok?", ["Yes", "No"]);
    await model.prefill(w.ids, []);
    status("ready");
    return model;
  }

  encode(state, question, options) {
    const t = this.tok;
    const idsOf = (text) => Array.from(t.encode(sanitize(text), { add_special_tokens: false }), Number);
    const ids = [this.mid.state, ...idsOf(state), this.mid.q, ...idsOf(question)];
    const optEnds = [];
    for (const o of options) {
      ids.push(this.mid.opt, ...idsOf(o));
      optEnds.push(ids.length);
      ids.push(this.mid.optEnd);
    }
    ids.push(this.mid.decide);
    return { ids, decide: ids.length - 1, optEnds };
  }

  async ask(q) {
    const t0 = performance.now();
    let options, question, labels;
    if (q.kind === "noul") { options = ["Yes", "No"]; question = q.q ?? ""; labels = ["yes", "no"]; }
    else if (q.kind === "score") { options = LEVELS; question = q.q ?? DEFAULT_SCORE_Q; labels = LEVELS; }
    else { options = q.options; question = q.q ?? ""; labels = options; }
    const enc = this.encode(q.state, question, options);
    const p0 = performance.now();
    const { hiddens } = await this.prefill(enc.ids, [enc.decide, ...enc.optEnds]);
    const prefillMs = performance.now() - p0;
    const hDecide = this.norm(hiddens[0]), perOpt = hiddens.slice(1).map((h) => this.norm(h));
    const raws = perOpt.map((h) => headLogit(this.head, hDecide, h));
    const ps = softmax(raws);
    const out = { kind: q.kind, ms: performance.now() - t0, prefillMs, nTokens: enc.ids.length, raws: {} };
    labels.forEach((l, i) => out.raws[l] = raws[i]);
    if (q.kind === "noul") out.p = ps[0];
    else if (q.kind === "score") {
      out.per = Object.fromEntries(labels.map((l, i) => [l, ps[i]]));
      out.level = 1 + ps.indexOf(Math.max(...ps));
      out.expected = ps.reduce((a, p, i) => a + p * (i + 1), 0);
    } else {
      out.choice = Object.fromEntries(labels.map((l, i) => [l, ps[i]]));
      out.top = labels[ps.indexOf(Math.max(...ps))];
      out.p = Math.max(...ps);
    }
    return out;
  }

  // ── System 2: text from the same 4-bit weights ─────────────────────────
  // The graph exposes logits; the runner's backend samples from them (end tokens <|im_end|> and
  // <|endoftext|>).
  async generate(q) {
    const t0 = performance.now();
    const state = q.state ?? "";
    const prompt = (q.prompt ?? "").trim();
    const user = state ? prompt + "\n\n" + state : prompt;
    const t = this.tok;
    const idsOf = (text) => Array.from(t.encode(sanitize(text), { add_special_tokens: false }), Number);
    if (!this._added) {
      const tk = await (await fetch(this.base + "model/tokenizer.json")).json();
      this._added = new Map((tk.added_tokens ?? []).map((x) => [x.content, x.id]));
    }
    const imS = this._added.get("<|im_start|>"), imE = this._added.get("<|im_end|>");
    // The base family deliberates by default and this checkpoint in 4-bit does not close the block:
    // the assistant turn is forced to start after the closing marker. The special tokens go in by id,
    // never through idsOf(): the sanitiser would rewrite them.
    const body = idsOf("user\n" + user);
    const asst = q.forceNoThink === false
      ? [imS, ...idsOf("assistant\n")]
      : [imS, ...idsOf("assistant\n</think>\n\n")];
    const ids = [imS, ...body, imE, ...asst];
    const out = [];
    const tt0 = performance.now();
    for await (const tokId of this.be.generate(ids, {
      maxNewTokens: q.maxNewTokens ?? 96,
      temperature: q.temperature ?? 0.7,
      topP: q.topP ?? 0.9,
    })) {
      out.push(tokId);
      if (q.onToken) q.onToken(t.decode([tokId], { skip_special_tokens: true }), out.length);
    }
    const genMs = performance.now() - tt0;
    return {
      text: t.decode(out, { skip_special_tokens: true, clean_up_tokenization_spaces: false }).trim(),
      tokens: out.length, genMs,
      tokPerS: +(out.length / (genMs / 1000)).toFixed(1),
      ms: performance.now() - t0,
    };
  }
}

	// alias3-web (alias-web in the paper) · a cybersecurity transformer in the browser (WebGPU).
	//
	// Two regimes from one loaded 4-bit graph:
	// System 1 ask() one typed question, one prefill, no token decoded: a probability per option
	// System 2 generate() text from the same weights, token by token, on request
	//
	// The prefill engine is the vendored runner (vendor/ct0.js: ct0-edge by Eduardo Worrel, MIT, see
	// vendor/ct0-edge.LICENSE), unchanged. This file adds the readout the model was trained with:
	//
	// · the System 1 markers as the model's own special tokens, by id:
	// <state> = <\|fim_prefix\|> <q> = <\|fim_middle\|> <opt> = <\|box_start\|>
	// </opt> = <\|box_end\|> <decide> = <\|fim_suffix\|>
	// · the final RMSNorm (heads/norm.bin): the graph exports hidden states before it
	// · the complete pointer head (heads/head.bin), read as in training:
	// z_i = q(h_<decide>) · k(h_</opt>_i) / (sqrt(256) · T), softmax over the options
	// · yes/no and severity as choices over ("Yes", "No") and ("1" … "5"): one mechanism for every question
	//
	// The page that imports this module needs an import map for the runner and its two runtimes; see index.html.

	export { AliasWeb };

	const MARKERS = {
	state: "<\|fim_prefix\|>",
	q: "<\|fim_middle\|>",
	opt: "<\|box_start\|>",
	optEnd: "<\|box_end\|>",
	decide: "<\|fim_suffix\|>",
	};
	const LEVELS = ["1", "2", "3", "4", "5"];
	export const DEFAULT_SCORE_Q = "How severe is the security risk in this content?";

	// ── head.bin: [u32 lenJSON][json][fp16 q][fp16 qb][fp16 k][fp16 kb] ────────
	function f16(h) {
	const s = (h & 0x8000) >> 15, e = (h & 0x7c00) >> 10, f = h & 0x03ff;
	if (e === 0) return (s ? -1 : 1) * Math.pow(2, -14) * (f / 1024);
	if (e === 0x1f) return f ? NaN : s ? -Infinity : Infinity;
	return (s ? -1 : 1) * Math.pow(2, e - 15) * (1 + f / 1024);
	}
	export function loadHead(buf) {
	const dv = new DataView(buf);
	const jl = dv.getUint32(0, true);
	const meta = JSON.parse(new TextDecoder().decode(new Uint8Array(buf, 4, jl)));
	let off = 4 + jl;
	const take = (n) => { const f = new Float32Array(n); for (let i = 0; i < n; i++) f[i] = f16(dv.getUint16(off + i * 2, true)); off += n * 2; return f; };
	const d = meta.headDim, h = meta.hiddenSize;
	return { q: take(d * h), qb: take(d), k: take(d * h), kb: take(d), headDim: d, hiddenSize: h,
	temperature: meta.temperature ?? 1.0, scale: meta.scale ?? 1 / Math.sqrt(d), scoreOffset: meta.score_offset ?? 0 };
	}
	// norm.bin: [u32 lenJSON][json {hiddenSize}][fp32 hiddenSize]
	async function loadNorm(buf) {
	const dv = new DataView(buf);
	const jl = dv.getUint32(0, true);
	const meta = JSON.parse(new TextDecoder().decode(new Uint8Array(buf, 4, jl)));
	const n = meta.hiddenSize, off = 4 + jl, f = new Float32Array(n);
	for (let i = 0; i < n; i++) f[i] = dv.getFloat32(off + i * 4, true);
	return f;
	}
	function proj(W, b, x, dim, hidden) {
	const y = new Float32Array(dim);
	for (let d = 0; d < dim; d++) { let a = b[d]; const r = d * hidden; for (let i = 0; i < hidden; i++) a += W[r + i] * x[i]; y[d] = a; }
	return y;
	}
	const dot = (a, b) => { let s = 0; for (let i = 0; i < a.length; i++) s += a[i] * b[i]; return s; };
	const headLogit = (hd, hD, hO) => dot(proj(hd.q, hd.qb, hD, hd.headDim, hd.hiddenSize), proj(hd.k, hd.kb, hO, hd.headDim, hd.hiddenSize)) * hd.scale / hd.temperature + (hd.scoreOffset \|\| 0);
	const softmax = (xs) => {
	const mx = Math.max(...xs.filter(Number.isFinite));
	const e = xs.map((x) => (Number.isFinite(x) ? Math.exp(x - mx) : 0));
	const z = e.reduce((a, b) => a + b, 0) \|\| 1;
	return e.map((v) => v / z);
	};

	// text supplied by the caller can never produce a delimiter: <\|name\|> becomes <¦name¦>, as in training
	const SPECIAL_RE = /<\\|([A-Za-z0-9_]+)\\|>/g;
	const sanitize = (t) => t.replace(SPECIAL_RE, "<¦$1¦>");

	class AliasWeb {
	constructor(tok, prefill, head, mid, normW) { this.tok = tok; this.prefill = prefill; this.head = head; this.mid = mid; this.normW = normW; }

	// the graph exports hidden_states before the final norm; the head was trained on the normalised
	// state, so the model's final RMSNorm is applied here
	norm(h) {
	if (!this.normW) return h;
	let s = 0; for (let i = 0; i < h.length; i++) s += h[i] * h[i];
	const inv = 1 / Math.sqrt(s / h.length);
	const o = new Float32Array(h.length);
	for (let i = 0; i < h.length; i++) o[i] = h[i] * inv * this.normW[i];
	return o;
	}

	static async load(opts = {}) {
	const base = opts.base ?? "./";
	const status = opts.status ?? (() => { });
	status("loading the runner and the 4-bit graph (2.4 GB)…");
	const edge = await import("ct0-edge");
	const inner = await edge.CT0.load({
	runtime: "genai",
	genaiBaseUrl: base + "model/",
	genaiModel: "model.onnx",
	baseUrl: base + "model/",
	status: () => { },
	verifyEp: false, // placement was verified once against this graph in Chrome
	skipVersionCheck: true, // the import map pins the ONNX Runtime version
	genai: { warmup: false },
	});
	const be = inner.backend ?? inner.be;
	if (!be?.prefill) throw new Error("the runner did not expose its prefill backend");
	status("loading the head and the final norm…");
	const head = loadHead(await (await fetch(base + "heads/head.bin")).arrayBuffer());
	const normW = await loadNorm(await (await fetch(base + "heads/norm.bin")).arrayBuffer());
	// marker ids from the tokenizer's added_tokens
	const tk = await (await fetch(base + "model/tokenizer.json")).json();
	const added = new Map((tk.added_tokens ?? []).map((t) => [t.content, t.id]));
	const mid = {};
	for (const [k, m] of Object.entries(MARKERS)) {
	const id = added.get(m);
	if (typeof id !== "number") throw new Error("marker without a token id: " + m);
	mid[k] = id;
	}
	const model = new AliasWeb(inner.tok?.tokenizer ?? inner.tok, (ids, positions) => be.prefill(ids, { positions }), head, mid, normW);
	model.be = be; model.base = base;
	status("warm-up…");
	const w = model.encode("warmup", "ok?", ["Yes", "No"]);
	await model.prefill(w.ids, []);
	status("ready");
	return model;
	}

	encode(state, question, options) {
	const t = this.tok;
	const idsOf = (text) => Array.from(t.encode(sanitize(text), { add_special_tokens: false }), Number);
	const ids = [this.mid.state, ...idsOf(state), this.mid.q, ...idsOf(question)];
	const optEnds = [];
	for (const o of options) {
	ids.push(this.mid.opt, ...idsOf(o));
	optEnds.push(ids.length);
	ids.push(this.mid.optEnd);
	}
	ids.push(this.mid.decide);
	return { ids, decide: ids.length - 1, optEnds };
	}

	async ask(q) {
	const t0 = performance.now();
	let options, question, labels;
	if (q.kind === "noul") { options = ["Yes", "No"]; question = q.q ?? ""; labels = ["yes", "no"]; }
	else if (q.kind === "score") { options = LEVELS; question = q.q ?? DEFAULT_SCORE_Q; labels = LEVELS; }
	else { options = q.options; question = q.q ?? ""; labels = options; }
	const enc = this.encode(q.state, question, options);
	const p0 = performance.now();
	const { hiddens } = await this.prefill(enc.ids, [enc.decide, ...enc.optEnds]);
	const prefillMs = performance.now() - p0;
	const hDecide = this.norm(hiddens[0]), perOpt = hiddens.slice(1).map((h) => this.norm(h));
	const raws = perOpt.map((h) => headLogit(this.head, hDecide, h));
	const ps = softmax(raws);
	const out = { kind: q.kind, ms: performance.now() - t0, prefillMs, nTokens: enc.ids.length, raws: {} };
	labels.forEach((l, i) => out.raws[l] = raws[i]);
	if (q.kind === "noul") out.p = ps[0];
	else if (q.kind === "score") {
	out.per = Object.fromEntries(labels.map((l, i) => [l, ps[i]]));
	out.level = 1 + ps.indexOf(Math.max(...ps));
	out.expected = ps.reduce((a, p, i) => a + p * (i + 1), 0);
	} else {
	out.choice = Object.fromEntries(labels.map((l, i) => [l, ps[i]]));
	out.top = labels[ps.indexOf(Math.max(...ps))];
	out.p = Math.max(...ps);
	}
	return out;
	}

	// ── System 2: text from the same 4-bit weights ─────────────────────────
	// The graph exposes logits; the runner's backend samples from them (end tokens <\|im_end\|> and
	// <\|endoftext\|>).
	async generate(q) {
	const t0 = performance.now();
	const state = q.state ?? "";
	const prompt = (q.prompt ?? "").trim();
	const user = state ? prompt + "\n\n" + state : prompt;
	const t = this.tok;
	const idsOf = (text) => Array.from(t.encode(sanitize(text), { add_special_tokens: false }), Number);
	if (!this._added) {
	const tk = await (await fetch(this.base + "model/tokenizer.json")).json();
	this._added = new Map((tk.added_tokens ?? []).map((x) => [x.content, x.id]));
	}
	const imS = this._added.get("<\|im_start\|>"), imE = this._added.get("<\|im_end\|>");
	// The base family deliberates by default and this checkpoint in 4-bit does not close the block:
	// the assistant turn is forced to start after the closing marker. The special tokens go in by id,
	// never through idsOf(): the sanitiser would rewrite them.
	const body = idsOf("user\n" + user);
	const asst = q.forceNoThink === false
	? [imS, ...idsOf("assistant\n")]
	: [imS, ...idsOf("assistant\n</think>\n\n")];
	const ids = [imS, ...body, imE, ...asst];
	const out = [];
	const tt0 = performance.now();
	for await (const tokId of this.be.generate(ids, {
	maxNewTokens: q.maxNewTokens ?? 96,
	temperature: q.temperature ?? 0.7,
	topP: q.topP ?? 0.9,
	})) {
	out.push(tokId);
	if (q.onToken) q.onToken(t.decode([tokId], { skip_special_tokens: true }), out.length);
	}
	const genMs = performance.now() - tt0;
	return {
	text: t.decode(out, { skip_special_tokens: true, clean_up_tokenization_spaces: false }).trim(),
	tokens: out.length, genMs,
	tokPerS: +(out.length / (genMs / 1000)).toFixed(1),
	ms: performance.now() - t0,
	};
	}
	}