| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| import { readFileSync, writeFileSync, mkdirSync } from "node:fs"; |
| import { createHash } from "node:crypto"; |
|
|
| const here = new URL("..", import.meta.url).pathname; |
| const args = process.argv.slice(2); |
| const imgAt = args.indexOf("--image"); |
| const image = imgAt >= 0 ? args.splice(imgAt, 2)[1] : null; |
| const dir = (args[0] || here + "../..").replace(/\/$/, ""); |
| const meta = JSON.parse(readFileSync(dir + "/meta.json", "utf8")); |
| const bin = readFileSync(dir + "/model.bin"); |
| if (bin.length !== meta.bytes) throw new Error(`model.bin is ${bin.length} bytes, meta.json says ${meta.bytes}`); |
| const T = Object.fromEntries(meta.tensors.map((t) => [t.name, t])); |
| const need = (n) => { if (!T[n]) throw new Error("missing tensor " + n); return T[n]; }; |
|
|
| for (const t of meta.tensors) { |
| if (t.dtype === "q4" && (t.offset % 16 || t.scale_offset % 2 || t.shape[1] % 32)) |
| throw new Error(`Q4 tensor ${t.name} breaks the kernel's alignment rules`); |
| if (t.offset % 4) throw new Error(`tensor ${t.name} is not 4-byte aligned`); |
| } |
|
|
| |
| const tk = meta.tokenizer; |
| if (tk.kind !== "wordpiece" || !tk.lowercase) throw new Error("expected a lowercase WordPiece tokenizer"); |
| const entries = []; |
| tk.vocab.forEach((p, i) => { if (p !== null && p !== undefined) entries.push([Buffer.from(p, "utf8"), i]); }); |
| entries.sort((a, b) => Buffer.compare(a[0], b[0])); |
| for (let i = 1; i < entries.length; i++) |
| if (Buffer.compare(entries[i - 1][0], entries[i][0]) === 0) throw new Error("duplicate piece " + entries[i][0]); |
| const n = entries.length, blobLen = entries.reduce((s, e) => s + e[0].length, 0); |
| const head = Buffer.alloc(12 + 4 * (n + 1) + 2 * n); |
| head.write("TDV1", 0, "ascii"); head.writeUInt32LE(n, 4); head.writeUInt32LE(blobLen, 8); |
| let off = 0; |
| entries.forEach(([b], i) => { head.writeUInt32LE(off, 12 + 4 * i); off += b.length; }); |
| head.writeUInt32LE(off, 12 + 4 * n); |
| entries.forEach(([, id], i) => head.writeUInt16LE(id, 12 + 4 * (n + 1) + 2 * i)); |
| const vocab = Buffer.concat([head, ...entries.map((e) => e[0])]); |
|
|
| |
| const c = meta.cfg; |
| if (c.arch !== "electra" || c.fusion !== "all" || (c.L_a || 0) !== 0) throw new Error("engine supports electra / fusion all / L_a 0"); |
| if (!T["h.P.weight"]) throw new Error("engine needs a build with the h.P prototype projection"); |
| const q4 = (n) => { const t = need(n); if (t.dtype !== "q4") throw new Error(`${n} is ${t.dtype}, the ESP32 engine needs q4`); return `{${t.offset}u, ${t.scale_offset}u, ${t.shape[0]}, ${t.shape[1]}}`; }; |
| const i8 = (n) => { const t = need(n); if (t.dtype !== "int8") throw new Error(n + " is not int8"); return `{${t.offset}u, ${t.scale_offset}u, ${t.shape[0]}, ${t.shape[1]}}`; }; |
| const f32 = (n) => { const t = need(n); if (t.dtype !== "f32") throw new Error(n + " is not f32"); return `${t.offset}u`; }; |
| const ffn = need("m.blocks.0.fc.weight").shape[0]; |
| if (c.d / c.heads !== 64 || ffn % 64) throw new Error("engine expects 64-wide heads and an FFN width that is a multiple of 64"); |
| const blocks = []; |
| for (let l = 0; l < c.layers; l++) { |
| const p = `m.blocks.${l}.`; |
| blocks.push(` {${q4(p + "q.weight")}, ${q4(p + "k.weight")}, ${q4(p + "v.weight")}, ${q4(p + "o.weight")},\n` + |
| ` ${q4(p + "fc.weight")}, ${q4(p + "fc2.weight")},\n` + |
| ` ${f32(p + "q.bias")}, ${f32(p + "k.bias")}, ${f32(p + "v.bias")}, ${f32(p + "o.bias")}, ${f32(p + "fc.bias")}, ${f32(p + "fc2.bias")},\n` + |
| ` ${f32(p + "ln1.weight")}, ${f32(p + "ln1.bias")}, ${f32(p + "ln2.weight")}, ${f32(p + "ln2.bias")}}`); |
| } |
| const sp = meta.specials, F = meta.format; |
| const fl = (x) => { const s = String(Math.fround(x)); return (s.includes(".") || s.includes("e") ? s : s + ".0") + "f"; }; |
| |
| let fnv = 0x811c9dc5; |
| for (let i = 0; i < Math.min(65536, bin.length); i++) { fnv ^= bin[i]; fnv = Math.imul(fnv, 0x01000193) >>> 0; } |
| const sha = createHash("sha256").update(bin).digest("hex"); |
| const h = `// GENERATED by tools/pack.mjs from meta.json (${meta.name}, model.bin sha256 ${sha.slice(0, 16)}...). Do not edit. |
| #pragma once |
| #include <stdint.h> |
| |
| namespace tdm { |
| |
| constexpr uint32_t MODEL_BYTES = ${meta.bytes}u; |
| constexpr uint32_t MODEL_FNV64K = 0x${fnv.toString(16).padStart(8, "0")}u; // FNV-1a of the first 64 KB |
| constexpr const char* MODEL_SHA256 = "${sha}"; |
| constexpr int D = ${c.d}, LAYERS = ${c.layers}, HEADS = ${c.heads}, DHEAD = ${c.d / c.heads}; |
| constexpr int EMB = ${c.emb_dim}, FFN = ${ffn}, DH = ${c.dh_head}, VOCAB = ${tk.vocab.length}; |
| constexpr float LN_EPS = ${fl(c.ln_eps)}; |
| constexpr int MAX_CHARS = ${tk.max_chars}; |
| |
| struct Q4 { uint32_t nib, sc; int rows, cols; }; // nibble plane, bf16 scale plane |
| struct I8 { uint32_t q, sc; int rows, cols; }; // int8 rows, one f32 scale per row |
| |
| struct Block { |
| Q4 q, k, v, o, fc, fc2; |
| uint32_t qb, kb, vb, ob, fcb, fc2b; |
| uint32_t ln1w, ln1b, ln2w, ln2b; |
| }; |
| |
| constexpr uint32_t TYPE0 = ${f32("m.type0")}; |
| constexpr Q4 WORD = ${q4("m.word.weight")}; |
| constexpr I8 POS = ${i8("m.posemb.weight")}; |
| constexpr uint32_t ELN_W = ${f32("m.eln.weight")}, ELN_B = ${f32("m.eln.bias")}; |
| constexpr Q4 PROJ = ${q4("m.proj.weight")}; |
| constexpr uint32_t PROJ_B = ${f32("m.proj.bias")}; |
| |
| constexpr Block BLOCKS[LAYERS] = { |
| ${blocks.join(",\n")} |
| }; |
| |
| // heads |
| constexpr uint32_t H_SCALE = ${f32("h.scale")}; |
| constexpr uint32_t H_NORM_W = ${f32("h.norm.weight")}, H_NORM_B = ${f32("h.norm.bias")}; |
| constexpr I8 H_A[2] = {${i8("h.A.0.weight")}, ${i8("h.A.1.weight")}}; |
| constexpr I8 H_O[2] = {${i8("h.O.0.weight")}, ${i8("h.O.1.weight")}}; |
| constexpr I8 H_P = ${i8("h.P.weight")}; |
| constexpr uint32_t H_NOUL_W = ${f32("h.noul.weight")}, H_NOUL_B = ${f32("h.noul.bias")}; |
| constexpr I8 H_NOUL_Q = ${i8("h.noul_q.weight")}; // unused since protos without a center use the count-weighted centre |
| constexpr I8 H_SQ = ${i8("h.sq.weight")}, H_SK = ${i8("h.sk.weight")}; |
| constexpr uint32_t H_SNULL_W = ${f32("h.snull.weight")}, H_SNULL_B = ${f32("h.snull.bias")}; |
| constexpr I8 H_EQ = ${i8("h.eq.weight")}, H_ES = ${i8("h.es.weight")}, H_EK = ${i8("h.ek.weight")}; |
| |
| constexpr int TS_MAX = ${F.ts_max}, P_Q = ${F.p_q}, Q_MAX = ${F.q_max}, K_MAX = ${F.k_max}, SPAN_MAX = ${F.span_max}; |
| constexpr int POS_ROWS = ${need("m.posemb.weight").shape[0]}; |
| |
| constexpr int SP_STATE = ${sp["<|state|>"]}, SP_SEP = ${sp["<|sep|>"]}, SP_O = ${sp["<|o|>"]}, SP_LV = ${sp["<|lv|>"]}, SP_ANS = ${sp["<|ans|>"]}; |
| constexpr int SP_TYPE[4] = {${sp["<|choice|>"]}, ${sp["<|noul|>"]}, ${sp["<|score|>"]}, ${sp["<|span|>"]}}; // choice, noul, score, span |
| constexpr int SP_UNK = ${tk.vocab.indexOf(tk.unk)}; |
| |
| constexpr float TEMP[4] = {${meta.temp.map(fl).join(", ")}}; // choice, noul, score, span |
| constexpr float BETA[${meta.beta.length}] = {${meta.beta.map(fl).join(", ")}}; |
| |
| } // namespace tdm |
| `; |
| mkdirSync(here + "tinydecide/model", { recursive: true }); |
| writeFileSync(here + "tinydecide/model/vocab.bin", vocab); |
| writeFileSync(here + "tinydecide/src/td_meta.h", h); |
| let msg = `td_meta.h: ${meta.name}, ${c.layers} layers, FFN ${ffn}, vocab ${tk.vocab.length}; vocab.bin: ${vocab.length} B`; |
| if (image) { |
| const pad = Buffer.alloc((16 - (bin.length % 16)) % 16); |
| const img = Buffer.concat([bin, pad, vocab]); |
| writeFileSync(image, img); |
| msg += `; ${image}: ${img.length} B (0x${img.length.toString(16)})`; |
| } |
| console.log(msg); |
|
|
