Consolidate the release into ten files, fix the decoder pool race, and measure the figures with the merged engine
75bf8f3 verified | """Oasis reference: Mirage's forward pass in PyTorch, the 2-bit quantiser, the text reader's | |
| quantisation-aware training and evaluation, and the export to .oasis. | |
| python reference.py COMMAND [options] (python reference.py COMMAND --help describes one) | |
| Model export, checkpoint, run, det-eval | |
| Reader train, eval-reader, eval-detected, pack-words, pack-detected, synth-val | |
| The forward pass follows the equations Cactus Compute published for Mirage, and its decoder follows the | |
| Needle 3 reference in cactus-compute (Apache-2.0). Commands read Mirage's mirage.safetensors and mirage.cact | |
| from mirage/ beside this file (--mirage), and data from data/ (--data): COCO val2017 photos in | |
| data/coco/val300 (python oasis.py fetch-coco) with data/coco/annotations/instances_val2017.json, and word | |
| crops in parquet files. The engine's benchmarks against Mirage are in oasis.py. | |
| """ | |
| import argparse | |
| import glob | |
| import io | |
| import json | |
| import math | |
| import os | |
| import random | |
| import re | |
| import struct | |
| import sys | |
| import time | |
| import unicodedata | |
| import zlib | |
| from dataclasses import dataclass | |
| from math import prod | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from PIL import Image, ImageDraw, ImageFilter, ImageFont | |
| from safetensors.torch import load_file, save_file | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| sys.path.insert(0, HERE) | |
| from oasis import OasisEngine, coco_truth, edit_distance, iou, sha256, val300, word_key, word_scene # noqa: E402 | |
| CHARSET_ENTRY, VOCAB_ENTRY = 804, 805 # mirage.cact entries holding the reader's charset and the tokenizer | |
| EMBEDDING_ROWS = 8299 | |
| LINUX_FONTS = ["/usr/share/fonts/truetype/dejavu/*.ttf", "/usr/share/fonts/truetype/liberation/*.ttf"] | |
| # ===================================================================================== shared helpers | |
| def parser(name, doc): | |
| """A command's argument parser, with the --mirage and --data directories every command accepts.""" | |
| ap = argparse.ArgumentParser(prog=f"reference.py {name}", description=doc, | |
| formatter_class=argparse.RawDescriptionHelpFormatter) | |
| ap.add_argument("--mirage", default=os.path.join(HERE, "mirage"), help="directory of mirage.safetensors and mirage.cact") | |
| ap.add_argument("--data", default=os.path.join(HERE, "data"), help="directory of the benchmark data") | |
| return ap | |
| def mirage_weights(args): | |
| return os.path.join(args.mirage, "mirage.safetensors") | |
| def mirage_archive(args): | |
| return os.path.join(args.mirage, "mirage.cact") | |
| def charset_of(arc): | |
| return "".join(chr(int(c)) for c in arc.array(CHARSET_ENTRY)) | |
| def vocab_of(arc, rows=EMBEDDING_ROWS): | |
| return Vocab(arc.raw(VOCAB_ENTRY), embedding_rows=rows) | |
| def dead_classes(weights): | |
| """Reader classes whose output weights are zero in `weights`: characters Mirage's reader never learned.""" | |
| k = weights["reader/out/kernel"] | |
| return [c for c in range(1, k.shape[1]) if k[:, c].norm() < 1e-6] | |
| def vocab_words(vocab): | |
| """Alphabetic pieces of at least two letters, the word list of the synthetic text.""" | |
| return sorted({p.replace("▁", "") for p, t in zip(vocab.pieces, vocab.types) | |
| if t == 0 and len(p.replace("▁", "")) >= 2 and p.replace("▁", "").isalpha()}) | |
| def ctc_text(labels, charset): | |
| """Collapse a CTC label path: drop blanks (0) and repeats.""" | |
| s, prev = [], 0 | |
| for lab in labels: | |
| if lab != 0 and lab != prev: | |
| s.append(charset[lab - 1]) | |
| prev = lab | |
| return "".join(s) | |
| def cq_plan(): | |
| """Mirage's quantised matrices in mirage.cact order: the archive entry, the float32 tensor, its decoder | |
| layer for tensors stacked over the eight layers, how its [out, in] matrix is read (T: transposed; | |
| T_per_layer: transposed per layer and stacked; reshape: as stored), the bits and the [out, in] shape.""" | |
| plan = [] | |
| def add(entry, name, shape, how="T", bits=2, layer=None): | |
| plan.append(dict(entry=entry, name=name, layer=layer, how=how, bits=bits, shape=list(shape))) | |
| add(0, "embedding/embedding", (8299, 512), "reshape", 4) | |
| for l in range(8): | |
| for k, proj, rows in ((0, "q", 384), (1, "k", 96), (2, "v", 128), (8, "gate", 512), (9, "out", 512)): | |
| add(2 + 27 * l + k, f"stack/layers/block/self_attn/{proj}_proj/kernel", (rows, 512), layer=l) | |
| for k, part, rows in ((0, "pre", 32), (1, "post", 32), (2, "res", 128)): | |
| add(223 + k, f"stack/mhc_phi_{part}", (rows, 2048), "T_per_layer", 4) | |
| for i in range(2): | |
| add(228 + 4 * i, f"engrams_{i}/embedding", (73728, 128), "reshape") | |
| add(229 + 4 * i, f"engrams_{i}/key_proj/kernel", (512, 512)) | |
| add(230 + 4 * i, f"engrams_{i}/value_proj/kernel", (512, 512)) | |
| for l in range(8): | |
| for k, proj, rows in ((0, "q", 384), (1, "k", 384), (2, "v", 512), (3, "gate", 512), (4, "out", 512)): | |
| add(239 + 10 * l + k, f"stack/layers/block/cross_attn/{proj}_proj/kernel", (rows, 512), layer=l) | |
| for t in range(12): | |
| for k, proj, shape in ((0, "q", (128, 512)), (1, "k", (32, 512)), (2, "v", (64, 512)), (3, "gate", (256, 512)), | |
| (4, "out", (512, 256))): | |
| add(542 + 15 * t + k, f"top_{t}/self_attn/{proj}_proj/kernel", shape) | |
| add(773, "slot_pool/kernel", (512, 512)) | |
| for j, shape in enumerate(((64, 32), (128, 64), (384, 128))): | |
| add(777 + 3 * j, f"reader/pw_{j}/kernel", shape) | |
| for b in range(4): | |
| add(787 + 4 * b, f"reader/up_{b}/kernel", (768, 384)) | |
| add(788 + 4 * b, f"reader/down_{b}/kernel", (384, 384)) | |
| add(802, "reader/out/kernel", (307, 384)) | |
| return plan | |
| # ===================================================================================== tokenizer | |
| # SentencePiece vocabulary carried in a .cact RAW entry, plus Mirage's added language and reference tokens. | |
| # RAW layout: u32 piece count, u32 pad, eos, bos, unk ids, u8 add-dummy-prefix flag, u8 byte-fallback flag, | |
| # u16 padding, then per piece f32 score, u8 type, u16 length, UTF-8 bytes. Types: 0 normal, 1 unknown, | |
| # 2 control, 3 user-defined, 4 byte. Mirage's embedding has 96 rows past the 8,203 pieces: reference tokens | |
| # for 64 object regions and 32 text regions. | |
| N_OBJECT_REFS = 64 | |
| N_TEXT_REFS = 32 | |
| class Vocab: | |
| def __init__(self, raw: bytes, embedding_rows: int): | |
| count, self.pad, self.eos, self.bos, self.unk = struct.unpack_from("<5I", raw, 0) | |
| self.add_dummy_prefix, self.byte_fallback = bool(raw[20]), bool(raw[21]) | |
| off = 24 | |
| self.pieces, self.types, self.scores = [], [], [] | |
| for _ in range(count): | |
| score, typ, ln = struct.unpack_from("<fBH", raw, off) | |
| off += 7 | |
| self.pieces.append(raw[off:off + ln].decode("utf-8", "replace")) | |
| self.types.append(typ) | |
| self.scores.append(score) | |
| off += ln | |
| self.n_pieces = count | |
| self.ref_base = count | |
| self.size = embedding_rows | |
| self.id_of = {p: i for i, p in enumerate(self.pieces)} | |
| self.lang_ids = {p[2:-2]: i for i, p in enumerate(self.pieces) if re.fullmatch(r"<\|[a-z]{2}\|>", p)} | |
| def ref_id(self, kind, k): | |
| return self.ref_base + (k if kind == "object" else N_OBJECT_REFS + k) | |
| def piece(self, i): | |
| if i >= self.n_pieces: | |
| k = i - self.ref_base | |
| return f"<obj{k}>" if k < N_OBJECT_REFS else f"<txt{k - N_OBJECT_REFS}>" | |
| return self.pieces[i] | |
| def decode(self, ids, show_special=True): | |
| out, buf = [], bytearray() | |
| def flush(): | |
| if buf: | |
| out.append(buf.decode("utf-8", "replace")) | |
| buf.clear() | |
| for i in ids: | |
| if i >= self.n_pieces: | |
| flush() | |
| out.append(self.piece(i)) | |
| continue | |
| t, p = self.types[i], self.pieces[i] | |
| if t == 4: | |
| m = re.fullmatch(r"<0x([0-9A-Fa-f]{2})>", p) | |
| if m: | |
| buf.append(int(m.group(1), 16)) | |
| continue | |
| flush() | |
| if t in (1, 2, 3): | |
| if show_special: | |
| out.append(p) | |
| else: | |
| out.append(p.replace("▁", " ")) | |
| flush() | |
| return "".join(out) | |
| # ===================================================================================== Cactus archive | |
| # Cactus .cact archives (tag 0x05E12A84): a 196-byte header of 49 little-endian 32-bit fields (the last is | |
| # rope_theta as float32), codebook floats, then a 44-byte record per tensor: u8 dtype, u8 ndim, u16 pad, | |
| # u32 shape[4], u64 offset, u64 nbytes, u32 group, u32 bits. Dtypes: 1 fp16, 2 fp32, 3 CQ, 4 RAW. | |
| # A CQ tensor is stored [out, in]; each row is split into groups of `group` weights (zero-padded), every | |
| # group is rotated by the natural-order normalised Walsh-Hadamard matrix, its L2 norm is kept in fp16, and | |
| # each rotated coordinate of the unit direction is replaced by the index of the nearest Lloyd-Max codebook | |
| # entry. The blob holds every row's packed indices (LSB first), then every group norm. | |
| TAG_NEEDLE3 = 0x05E12A84 | |
| DTYPES = {1: "fp16", 2: "fp32", 3: "cq", 4: "raw"} | |
| CODEBOOK_SLICES = {2: slice(0, 4), 3: slice(4, 12), 4: slice(12, 28)} | |
| class Entry: | |
| index: int | |
| dtype: str | |
| shape: tuple | |
| offset: int | |
| nbytes: int | |
| group: int | |
| bits: int | |
| def hadamard(n): | |
| h = np.array([[1.0]]) | |
| while h.shape[0] < n: | |
| h = np.block([[h, h], [h, -h]]) | |
| return h / np.sqrt(n) | |
| class Archive: | |
| def __init__(self, path): | |
| with open(path, "rb") as f: | |
| self.data = f.read() | |
| b = self.data | |
| self.fields = struct.unpack_from("<48I", b, 0) | |
| self.rope_theta = struct.unpack_from("<f", b, 192)[0] | |
| self.tag, n_tensors, cb_len = self.fields[0], self.fields[1], self.fields[2] | |
| if self.tag != TAG_NEEDLE3: | |
| raise ValueError(f"unsupported archive tag 0x{self.tag:08X}") | |
| self.codebook = np.array(struct.unpack_from(f"<{cb_len}f", b, 196), dtype=np.float64) | |
| off = 196 + 4 * cb_len | |
| self.entries = [] | |
| for i in range(n_tensors): | |
| dtype, ndim = b[off], b[off + 1] | |
| shape = struct.unpack_from("<4I", b, off + 4)[:ndim] | |
| start, nbytes = struct.unpack_from("<QQ", b, off + 20) | |
| group, bits = struct.unpack_from("<II", b, off + 36) | |
| self.entries.append(Entry(i, DTYPES[dtype], tuple(shape), start, nbytes, group, bits)) | |
| off += 44 | |
| def raw(self, i): | |
| e = self.entries[i] | |
| return self.data[e.offset:e.offset + e.nbytes] | |
| def array(self, i): | |
| """Return entry i as float64 in its stored layout (CQ dequantised).""" | |
| e = self.entries[i] | |
| if e.dtype in ("fp16", "fp32"): | |
| dt = np.float16 if e.dtype == "fp16" else np.float32 | |
| a = np.frombuffer(self.data, dtype=dt, count=e.nbytes // np.dtype(dt).itemsize, offset=e.offset) | |
| return a.astype(np.float64).reshape(e.shape) | |
| if e.dtype == "cq": | |
| return self._dequant(e) | |
| raise TypeError("RAW entries have no array form") | |
| def cq_parts(self, e): | |
| rows, cols = prod(e.shape[:-1]), e.shape[-1] | |
| ng = -(-cols // e.group) | |
| padded = ng * e.group | |
| blob = np.frombuffer(self.data, dtype=np.uint8, count=e.nbytes, offset=e.offset) | |
| ib = rows * padded * e.bits // 8 | |
| per = 8 // e.bits | |
| idx = np.empty(ib * per, dtype=np.uint8) | |
| for k in range(per): | |
| idx[k::per] = (blob[:ib] >> (e.bits * k)) & ((1 << e.bits) - 1) | |
| idx = idx.reshape(rows, ng, e.group) | |
| norms = blob[ib:].view(np.float16).astype(np.float64).reshape(rows, ng) | |
| return idx, norms, rows, cols, padded | |
| def _dequant(self, e): | |
| idx, norms, rows, cols, padded = self.cq_parts(e) | |
| cb = self.codebook[CODEBOOK_SLICES[e.bits]] | |
| rot = cb[idx] * norms[..., None] | |
| w = rot @ hadamard(e.group) | |
| return w.reshape(rows, padded)[:, :cols].reshape(e.shape) | |
| # ===================================================================================== quantisation | |
| # A quantised matrix is held [out, in]; each row is cut into 128-weight groups, every group is rotated by | |
| # the normalised natural-order Walsh-Hadamard matrix, and each rotated coordinate is replaced by a codebook | |
| # entry times one fp16 scale per group. "stock" sets the scale to the group's L2 norm and picks the nearest | |
| # entry of the unit-norm direction (the published scheme); "refit" alternates nearest-entry assignment and | |
| # the least-squares scale. | |
| H128 = torch.tensor(hadamard(128), dtype=torch.float64) | |
| def nearest(x, cb): | |
| """Index of the nearest codebook entry for every element of x (codebook sorted ascending).""" | |
| mids = (cb[1:] + cb[:-1]) / 2 | |
| return torch.bucketize(x.contiguous(), mids) | |
| def quantize_rows(W, cb, method="stock", iters=8, chunk=8192): | |
| rows = W.shape[0] | |
| if rows > chunk: | |
| return torch.cat([quantize_rows(W[i:i + chunk], cb, method, iters, chunk) for i in range(0, rows, chunk)]) | |
| rows, cols = W.shape | |
| pc = -(-cols // 128) * 128 | |
| Wp = torch.zeros(rows, pc, dtype=torch.float64) | |
| Wp[:, :cols] = W.double() | |
| g = Wp.view(rows, pc // 128, 128) @ H128 | |
| n = g.norm(dim=-1, keepdim=True) | |
| idx = nearest(g / n.clamp_min(1e-30), cb) | |
| s = n | |
| if method == "refit": | |
| for _ in range(iters): | |
| c = cb[idx] | |
| s = (g * c).sum(-1, keepdim=True) / (c * c).sum(-1, keepdim=True).clamp_min(1e-30) | |
| safe = torch.where(s.abs() > 1e-30, s, torch.full_like(s, 1e-30)) | |
| idx = nearest(g / safe, cb) | |
| c = cb[idx] | |
| s = (g * c).sum(-1, keepdim=True) / (c * c).sum(-1, keepdim=True).clamp_min(1e-30) | |
| s = s.half().double() | |
| rec = (cb[idx] * s) @ H128 | |
| return rec.reshape(rows, pc)[:, :cols].float() | |
| def plan_get(p, item): | |
| """The [out, in] matrix of a plan item, read from the weight dictionary p.""" | |
| a = p[item["name"]] | |
| if item["layer"] is not None: | |
| a = a[item["layer"]] | |
| if item["how"] == "T": | |
| return a.T | |
| if item["how"] == "T_per_layer": | |
| L, i_, o_ = a.shape | |
| return a.transpose(1, 2).reshape(L * o_, i_) | |
| return a.reshape(item["shape"]) | |
| def plan_put(p, item, W): | |
| """Write the [out, in] matrix W of a plan item back into the weight dictionary p.""" | |
| a = p[item["name"]] | |
| if item["how"] == "T": | |
| new = W.T.contiguous() | |
| elif item["how"] == "T_per_layer": | |
| L, i_, o_ = a.shape | |
| new = W.reshape(L, o_, i_).transpose(1, 2).contiguous() | |
| else: | |
| new = W.reshape(a.shape[1:] if item["layer"] is not None else a.shape) | |
| if item["layer"] is not None: | |
| a[item["layer"]] = new | |
| else: | |
| p[item["name"]] = new | |
| def variant(p32, plan, codebook, method="stock", bits_rules=()): | |
| """Return a copy of p32 with every planned tensor quantised and everything else rounded to fp16. | |
| bits_rules: sequence of (regex on tensor name, bits) applied in order to override the plan. | |
| """ | |
| p = {k: v.half().float().clone() for k, v in p32.items()} | |
| size_bits = 0 | |
| for item in plan: | |
| bits = item["bits"] | |
| for pat, b in bits_rules: | |
| if re.search(pat, item["name"]): | |
| bits = b | |
| cb = torch.tensor(codebook[CODEBOOK_SLICES[bits]], dtype=torch.float64) | |
| W = plan_get(p32, item) | |
| plan_put(p, item, quantize_rows(W, cb, method)) | |
| rows, cols = W.shape | |
| size_bits += rows * (-(-cols // 128) * 128) * bits + rows * (-(-cols // 128)) * 16 | |
| quant_names = {it["name"] for it in plan} | |
| fp16_params = sum(v.numel() for k, v in p32.items() if k not in quant_names and not k.startswith(("mtp_", "name_head"))) | |
| return p, dict(cq_mb=round(size_bits / 8 / 1e6, 3), fp16_mb=round(fp16_params * 2 / 1e6, 3)) | |
| class Quantizer: | |
| """Quantisation-aware training for Hadamard-codebook weights: fake_quant_kernel(W) applies the stock | |
| quantiser in the forward pass (each 128-weight input group of a row rotated by the Walsh-Hadamard matrix, | |
| scaled by its fp16 L2 norm, every coordinate snapped to the nearest codebook entry) and passes gradients | |
| straight through to the latent weight.""" | |
| def __init__(self, codebook, device): | |
| self.H = torch.tensor(hadamard(128), dtype=torch.float32, device=device) | |
| self.cb = {b: torch.tensor(codebook[CODEBOOK_SLICES[b]], dtype=torch.float32, device=device) for b in (2, 3, 4)} | |
| self.mid = {b: (c[1:] + c[:-1]) / 2 for b, c in self.cb.items()} | |
| def quantize_rows(self, R, bits): | |
| """R: [rows, cols] in [out, in] layout. Returns the dequantised quantised matrix.""" | |
| rows, cols = R.shape | |
| pc = -(-cols // 128) * 128 | |
| if pc != cols: | |
| R = torch.nn.functional.pad(R, (0, pc - cols)) | |
| g = R.reshape(rows, pc // 128, 128) @ self.H | |
| n = g.norm(dim=-1, keepdim=True) | |
| idx = torch.bucketize((g / n.clamp_min(1e-30)).contiguous(), self.mid[bits]) | |
| s = n.half().float() | |
| rec = (self.cb[bits][idx] * s) @ self.H | |
| return rec.reshape(rows, pc)[:, :cols] | |
| def fake_quant_kernel(self, W, bits): | |
| """W: a Dense kernel stored [in, out]; quantised along rows of W.T.""" | |
| with torch.no_grad(): | |
| Q = self.quantize_rows(W.T.float(), bits).T | |
| return W + (Q - W).detach() | |
| # ===================================================================================== forward pass | |
| # Mirage's forward pass, float32 on the CPU, from the weights of mirage.safetensors. | |
| EPS = 1e-6 | |
| D_MODEL = 512 | |
| CANVAS = 512 | |
| N_OBJ, N_TXT = 64, 32 | |
| ENGRAM_SITES = (3, 7) | |
| ENGRAM_ORDERS = (2, 3) | |
| ENGRAM_HEADS = 2 | |
| ENGRAM_SLOTS = 18432 | |
| ENGRAM_SEED = 0x9E3779B9 | |
| ENGRAM_PRIME = 0x01000193 | |
| ENGRAM_TAPS, ENGRAM_DILATION = 4, 3 | |
| # Settings the published equations leave open; kept together so they can be varied. | |
| OPTS = { | |
| "pad_mode": "sym", # "sym": pad 1 on both sides; "same": TF-style (0, 1) for stride 2 | |
| "canvas_fill": 0.0, # normalised value of letterbox padding | |
| "rope2d_rows_first": True, | |
| "rope2d_layout": "pairs16", # "pairs16": pairs (r, r+16), rows on r<8; "blocks": two 16-d halves | |
| "walsh": "sequency", # or "natural" | |
| "level_gate": "raw", # or "sigmoid" | |
| "resample": Image.BILINEAR, | |
| } | |
| def zc_norm(x, scale, dim=-1): | |
| shape = [1] * x.dim() | |
| shape[dim] = -1 | |
| return x * torch.rsqrt(x.pow(2).mean(dim, keepdim=True) + EPS) * (1 + scale.view(shape)) | |
| def rms_unit(x): | |
| return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS) | |
| def _pad(x, stride_hw): | |
| if OPTS["pad_mode"] == "sym" or stride_hw == (1, 1): | |
| return F.pad(x, (1, 1, 1, 1)) | |
| pads = [] | |
| for size, s in ((x.shape[3], stride_hw[1]), (x.shape[2], stride_hw[0])): | |
| out = -(-size // s) | |
| total = max((out - 1) * s + 3 - size, 0) | |
| pads += [total // 2, total - total // 2] | |
| return F.pad(x, pads) | |
| def conv3(x, w_hwio, stride=(1, 1)): | |
| w = w_hwio.permute(3, 2, 0, 1).contiguous() | |
| return F.conv2d(_pad(x, stride), w, stride=stride) | |
| def dwconv(x, w_hw1c, stride=(1, 1)): | |
| w = w_hw1c.permute(3, 2, 0, 1).contiguous() | |
| if w.shape[-1] == 3 and w.shape[-2] == 3: | |
| return F.conv2d(_pad(x, stride), w, stride=stride, groups=x.shape[1]) | |
| kh, kw = w.shape[-2:] | |
| return F.conv2d(F.pad(x, (kw // 2, kw // 2, kh // 2, kh // 2)), w, stride=stride, groups=x.shape[1]) | |
| def pw(x, k): | |
| return (x.permute(0, 2, 3, 1) @ k).permute(0, 3, 1, 2) | |
| def cnorm(x, scale): | |
| return zc_norm(x, scale, dim=1) | |
| def walsh_sequency(n): | |
| q = n.bit_length() - 1 | |
| j = np.arange(n) | |
| g = j ^ (j >> 1) | |
| s = np.array([int(format(int(v), f"0{q}b")[::-1], 2) for v in g]) | |
| pc = np.vectorize(lambda v: bin(int(v)).count("1"))(s[:, None] & np.arange(n)[None, :]) | |
| return torch.tensor(((-1.0) ** pc) / math.sqrt(n), dtype=torch.float32) | |
| def band(u, w): | |
| r = (w.shape[1] - 1) // 2 | |
| win = F.pad(u, (r, r)).unfold(-1, 2 * r + 1, 1) | |
| return (win * w).sum(-1) | |
| def sinkhorn(logits, iters=20): | |
| k = logits | |
| for _ in range(iters): | |
| k = k - torch.logsumexp(k, dim=-1, keepdim=True) | |
| k = k - torch.logsumexp(k, dim=-2, keepdim=True) | |
| return k.exp() | |
| def rope_half(x, pos, theta): | |
| hd = x.shape[-1] | |
| freqs = 1.0 / (theta ** (torch.arange(0, hd, 2, dtype=torch.float32) / hd)) | |
| ang = pos.float()[..., None] * freqs | |
| c, s = ang.cos(), ang.sin() | |
| x1, x2 = x[..., : hd // 2], x[..., hd // 2:] | |
| return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], -1) | |
| def sincos_2d(rows, cols, d=D_MODEL, theta=10000.0): | |
| def one(pos, dd): | |
| omega = 1.0 / (theta ** (torch.arange(dd // 2, dtype=torch.float32) / (dd // 2))) | |
| a = pos.float()[:, None] * omega[None] | |
| return torch.cat([a.sin(), a.cos()], -1) | |
| return torch.cat([one(rows, d // 2), one(cols, d // 2)], -1) | |
| def engram_index(tokens, t, oi, order, h): | |
| acc = (ENGRAM_SEED * (oi * ENGRAM_HEADS + h + 1)) & 0xFFFFFFFF | |
| for j in range(order): | |
| tok = tokens[t - j] if t - j >= 0 else 0 | |
| acc = ((acc ^ tok) * ENGRAM_PRIME) & 0xFFFFFFFF | |
| acc ^= acc >> 15 | |
| return acc % ENGRAM_SLOTS | |
| class Mirage: | |
| def __init__(self, weights_path): | |
| self.p = {k: v.float() for k, v in load_file(weights_path).items()} | |
| self.W = walsh_sequency(D_MODEL) | |
| self.W_nat = torch.tensor(hadamard(D_MODEL), dtype=torch.float32) | |
| p = self.p | |
| self.perm1 = torch.tensor(np.random.RandomState(11).permutation(D_MODEL)) | |
| self.perm2 = torch.tensor(np.random.RandomState(13).permutation(D_MODEL)) | |
| self.T = p["embedding/embedding"] | |
| # ---------------------------------------------------------------- image | |
| def letterbox(self, img): | |
| img = img.convert("RGB") | |
| W0, H0 = img.size | |
| r = CANVAS / max(W0, H0) | |
| w, h = max(1, round(W0 * r)), max(1, round(H0 * r)) | |
| rs = np.asarray(img.resize((w, h), OPTS["resample"]), dtype=np.float32) / 127.5 - 1.0 | |
| canvas = np.full((CANVAS, CANVAS, 3), OPTS["canvas_fill"], dtype=np.float32) | |
| canvas[:h, :w] = rs | |
| x = torch.from_numpy(canvas).permute(2, 0, 1)[None] | |
| return x, dict(W=W0, H=H0, r=r, w=w, h=h) | |
| def valid_mask(meta, stride): | |
| n = CANVAS // stride | |
| m = torch.zeros(1, 1, n, n) | |
| m[..., : -(-meta["h"] // stride), : -(-meta["w"] // stride)] = 1.0 | |
| return m | |
| def hconvglu(self, x, pre): | |
| p = self.p | |
| W = self.W_nat if OPTS["walsh"] == "natural" else self.W | |
| u = x @ W | |
| g = F.gelu(band(u, p[pre + "/gate_conv"]) + p[pre + "/gate_bias"], approximate="tanh") | |
| v = band(u, p[pre + "/value_conv"]) | |
| return band(g * v, p[pre + "/out_conv"]) @ W | |
| # --------------------------------------------------------------- encoder | |
| def encode(self, img): | |
| p = self.p | |
| x, meta = self.letterbox(img) | |
| ms = {s: self.valid_mask(meta, s) for s in (2, 4, 8, 16, 32)} | |
| f = ms[2] * cnorm(conv3(x, p["stem/w0"], (2, 2)), p["stem/norm_0/scale"]) | |
| feats = {2: f} | |
| for j in range(1, 5): | |
| s = 2 ** (j + 1) | |
| y = dwconv(F.silu(f), p[f"stem/dw_{j}"], (2, 2)) | |
| f = ms[s] * cnorm(pw(y, p[f"stem/pw_{j}/kernel"]), p[f"stem/norm_{j}/scale"]) | |
| feats[s] = f | |
| S4, S8, S16 = F.silu(feats[4]), F.silu(feats[8]), F.silu(feats[16]) | |
| X = F.silu(feats[32]) | |
| m32 = ms[32] | |
| for i in range(16): | |
| X = self.body_block(X, i, m32) | |
| Z = X.flatten(2).transpose(1, 2) # [1, 256, 512] | |
| Z = torch.cat([Z, p["registers"][None]], 1) # [1, 260, 512] | |
| valid = torch.cat([m32.flatten().bool(), torch.ones(4, dtype=torch.bool)]) | |
| rows = torch.arange(256) // 16 | |
| cols = torch.arange(256) % 16 | |
| levels = [Z] | |
| for l in range(12): | |
| Z = self.top_block(Z, l, valid, rows, cols) | |
| levels.append(Z) | |
| mv = m32.flatten()[None, :, None] | |
| E = zc_norm(levels[12][:, :256], p["enc_norm/scale"]) | |
| for l in range(12): | |
| g = p["level_gates"][l] if OPTS["level_gate"] == "raw" else torch.sigmoid(p["level_gates"][l]) | |
| E = E + g * zc_norm(levels[l][:, :256], p[f"level_norm_{l}/scale"]) | |
| E = mv * E # [1, 256, 512] | |
| D = self.neck(E, S16, S8, ms) | |
| return dict(meta=meta, ms=ms, E=E, D=D, S4=S4, valid=m32.flatten().bool()) | |
| def body_block(self, X, i, m): | |
| p, pre = self.p, f"body_{i}" | |
| Z = cnorm(X, p[pre + "/mix_norm/scale"]) | |
| v = dwconv(Z, p[pre + "/mix/dw_value"]) + p[pre + "/mix/b_value"].view(1, -1, 1, 1) | |
| g = dwconv(Z, p[pre + "/mix/dw_gate"]) + p[pre + "/mix/b_gate"].view(1, -1, 1, 1) | |
| Xt = X + m * dwconv(v * F.silu(g), p[pre + "/mix/dw_down"]) | |
| xt = Xt.permute(0, 2, 3, 1) | |
| H = self.hconvglu(zc_norm(xt, p[pre + "/hconv_norm/scale"]), pre + "/hconv") | |
| mm = m.permute(0, 2, 3, 1) | |
| r = torch.sqrt(((mm * H) ** 2).sum((1, 2), keepdim=True) + EPS) | |
| G = H + p[pre + "/grn/gamma"] * H * r / (r.mean(-1, keepdim=True) + EPS) + p[pre + "/grn/beta"] | |
| return (xt + mm * G).permute(0, 3, 1, 2) | |
| def top_block(self, Z, l, valid, rows, cols): | |
| p, pre = self.p, f"top_{l}/" | |
| T = Z.shape[1] | |
| z = zc_norm(Z, p[pre + "attn_norm/scale"]) | |
| q = (z @ p[pre + "self_attn/q_proj/kernel"]).view(1, T, 4, 32).transpose(1, 2) | |
| k = (z @ p[pre + "self_attn/k_proj/kernel"]).view(1, T, 1, 32).transpose(1, 2) | |
| v = (z @ p[pre + "self_attn/v_proj/kernel"]).view(1, T, 1, 64).transpose(1, 2) | |
| q = zc_norm(q, p[pre + "self_attn/q_norm/scale"]) | |
| k = zc_norm(k, p[pre + "self_attn/k_norm/scale"]) | |
| a, b = (rows, cols) if OPTS["rope2d_rows_first"] else (cols, rows) | |
| def rope2(t): | |
| c = t[:, :, :256] | |
| if OPTS["rope2d_layout"] == "blocks": | |
| c = torch.cat([rope_half(c[..., :16], a, 10000.0), rope_half(c[..., 16:], b, 10000.0)], -1) | |
| else: | |
| freqs = 1.0 / (10000.0 ** (torch.arange(0, 16, 2, dtype=torch.float32) / 16)) | |
| ang = torch.cat([a.float()[:, None] * freqs, b.float()[:, None] * freqs], -1) | |
| cs, sn = ang.cos(), ang.sin() | |
| x1, x2 = c[..., :16], c[..., 16:] | |
| c = torch.cat([x1 * cs - x2 * sn, x2 * cs + x1 * sn], -1) | |
| return torch.cat([c, t[:, :, 256:]], 2) | |
| q, k = rope2(q), rope2(k) | |
| s = (q @ k.transpose(-1, -2)) / math.sqrt(32) | |
| s = s.masked_fill(~valid.view(1, 1, 1, T), float("-inf")) | |
| o = (s.softmax(-1) @ v).transpose(1, 2).reshape(1, T, 256) | |
| A = (o * torch.sigmoid(z @ p[pre + "self_attn/gate_proj/kernel"])) @ p[pre + "self_attn/out_proj/kernel"] | |
| mv = valid.view(1, T, 1).float() | |
| Zt = Z + mv * torch.sigmoid(p[pre + "attn_gate"]) * zc_norm(A, p[pre + "post_attn_norm/scale"]) | |
| return Zt + mv * self.hconvglu(zc_norm(Zt, p[pre + "hconv_norm/scale"]), pre + "hconv") | |
| def res_block(self, Fm, norm, dw, pwk, m): | |
| p = self.p | |
| return m * (Fm + pw(dwconv(F.silu(cnorm(Fm, p[norm])), p[dw]), p[pwk])) | |
| def neck(self, E, S16, S8, ms): | |
| p = self.p | |
| Eg = E.transpose(1, 2).reshape(1, D_MODEL, 16, 16) | |
| t = F.interpolate(pw(Eg, p["dense/coarse/kernel"]), scale_factor=2, mode="nearest") | |
| t = ms[16] * (t + pw(S16, p["dense/mid/kernel"])) | |
| t = self.res_block(t, "dense/neck_norm/scale", "dense/neck_dw", "dense/neck_pw/kernel", ms[16]) | |
| d = ms[8] * (F.interpolate(t, scale_factor=2, mode="nearest") + pw(S8, p["dense/fine/kernel"])) | |
| for i in range(3): | |
| d = self.res_block(d, f"dense/block_norm_{i}/scale", f"dense/dw_{i}", f"dense/pw_{i}/kernel", ms[8]) | |
| return ms[8] * F.silu(cnorm(d, p["dense/final_norm/scale"])) # [1, 128, 64, 64] | |
| # ------------------------------------------------------------- detection | |
| def detect(self, enc, conf=0.3, text_conf=0.3, nms_iou=None): | |
| p = self.p | |
| out = torch.einsum("bchw,cd->bhwd", enc["D"], p["dense/out/kernel"])[0] + p["dense/bias"] | |
| m8 = enc["ms"][8][0, 0].bool() | |
| meta = enc["meta"] | |
| result = {} | |
| for kind, ch, limit, thr in (("object", 0, N_OBJ, conf), ("text", 1, N_TXT, text_conf)): | |
| heat = torch.sigmoid(out[..., ch]) | |
| peak = (heat == F.max_pool2d(heat[None, None], 3, 1, 1)[0, 0]) & m8 & (heat >= thr) | |
| ii, jj = torch.nonzero(peak, as_tuple=True) | |
| scores = heat[ii, jj] | |
| order = torch.argsort(scores, descending=True) | |
| boxes = [] | |
| for o in order.tolist(): | |
| i, j = int(ii[o]), int(jj[o]) | |
| w = abs(float(out[i, j, 2 + 2 * ch])) | |
| h = abs(float(out[i, j, 3 + 2 * ch])) | |
| cx = j + float(out[i, j, 6 + 2 * ch]) | |
| cy = i + float(out[i, j, 7 + 2 * ch]) | |
| b = np.array([cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2]) / 64.0 | |
| b[[0, 2]] = b[[0, 2]].clip(0, meta["w"] / CANVAS) | |
| b[[1, 3]] = b[[1, 3]].clip(0, meta["h"] / CANVAS) | |
| if nms_iou and any(iou(b, x["box"]) > nms_iou for x in boxes): | |
| continue | |
| boxes.append(dict(box=b, score=float(scores[o]), cell=(i, j))) | |
| if len(boxes) >= limit: | |
| break | |
| result[kind] = boxes | |
| result["raw"] = out | |
| return result | |
| def to_1000(self, b, meta): | |
| sx = CANVAS / meta["r"] / meta["W"] * 1000 | |
| sy = CANVAS / meta["r"] / meta["H"] * 1000 | |
| return [round(float(b[0] * sx)), round(float(b[1] * sy)), round(float(b[2] * sx)), round(float(b[3] * sy))] | |
| # ----------------------------------------------------------------- masks | |
| def head(self, Fm, pre, m): | |
| p = self.p | |
| for i in range(2): | |
| Fm = self.res_block(Fm, f"dense/{pre}_norm_{i}/scale", f"dense/{pre}_dw_{i}", f"dense/{pre}_pw_{i}/kernel", m) | |
| return pw(F.silu(cnorm(Fm, p[f"dense/{pre}_final_norm/scale"])), p[f"dense/{pre}s/kernel" if pre == "proto" else f"dense/{pre}/kernel"]) | |
| def masks(self, enc, boxes): | |
| p, ms, D = self.p, enc["ms"], enc["D"] | |
| U = ms[4] * (F.interpolate(D, scale_factor=2, mode="nearest") + pw(enc["S4"], p["dense/edge/kernel"])) | |
| P = self.head(U, "proto", ms[4])[0] # [32, 128, 128] | |
| C = self.head(D, "coef", ms[8])[0] # [32, 64, 64] | |
| out = [] | |
| ys = (torch.arange(128) + 0.5) / 128 | |
| for bx in boxes: | |
| i, j = bx["cell"] | |
| logit = torch.einsum("r,rhw->hw", C[:, i, j], P) | |
| b = bx["box"] | |
| inside = ((ys[:, None] >= b[1]) & (ys[:, None] <= b[3]) & (ys[None, :] >= b[0]) & (ys[None, :] <= b[2])) | |
| ci, cj = min(127, int((b[1] + b[3]) / 2 * 128)), min(127, int((b[0] + b[2]) / 2 * 128)) | |
| inside[ci, cj] = True | |
| out.append(torch.sigmoid(logit) * inside) | |
| return out | |
| # ---------------------------------------------------------------- reader | |
| def read(self, img, box_img): | |
| """box_img: (x0, y0, x1, y1) in original pixels. Returns frame label ids and per-frame probs.""" | |
| img = img.convert("RGB") | |
| x0, y0, x1, y1 = box_img | |
| crop = img.crop((int(math.floor(x0)), int(math.floor(y0)), int(math.ceil(x1)), int(math.ceil(y1)))) | |
| cw, ch = crop.size | |
| if cw < 1 or ch < 1: | |
| return [], [] | |
| U = max(1, int(round(48 * cw / ch))) | |
| strip = np.asarray(crop.resize((U, 48), OPTS["resample"]), dtype=np.float32) | |
| labels, probs = [], [] | |
| starts = [0] if U <= 320 else list(range(0, U - 320 + 223, 224)) | |
| starts = [s for s in starts if s == 0 or s < U] | |
| n_frames = -(-U // 4) | |
| for xs in starts: | |
| win = np.full((48, 320, 3), 127.0, dtype=np.float32) | |
| seg = strip[:, xs:xs + 320] | |
| win[:, : seg.shape[1]] = seg | |
| J = torch.from_numpy(win / 127.5 - 1.0).permute(2, 0, 1)[None] | |
| q = self.reader_frames(J) # [80, 307] | |
| lo = 0 if xs == 0 else 12 | |
| hi = (-(-(U - xs) // 4)) if xs + 320 >= U else 68 | |
| for t in range(lo, min(hi, 80)): | |
| if xs // 4 + t < n_frames: | |
| labels.append(int(q[t].argmax())) | |
| probs.append(q[t]) | |
| return labels, probs | |
| def reader_frames(self, J): | |
| p = self.p | |
| G = cnorm(conv3(J, p["reader/w0"], (2, 2)), p["reader/stage_norm_0/scale"]) | |
| for j, stride in enumerate(((2, 2), (2, 1), (2, 1))): | |
| y = dwconv(F.silu(G), p[f"reader/dw_{j}"], stride) | |
| G = cnorm(pw(y, p[f"reader/pw_{j}/kernel"]), p[f"reader/stage_norm_{j + 1}/scale"]) | |
| Z = F.silu(G).mean(2)[0].transpose(0, 1) # [80, 384] | |
| for b in range(4): | |
| h = zc_norm(Z, p[f"reader/norm_{b}/scale"]) | |
| k = p[f"reader/mix_{b}"][0, :, 0, :] # [5, 384] | |
| h = F.conv1d(F.pad(h.T[None], (2, 2)), k.T[:, None, :], groups=384)[0].T | |
| vg = h @ p[f"reader/up_{b}/kernel"] | |
| Z = Z + (vg[:, :384] * F.silu(vg[:, 384:])) @ p[f"reader/down_{b}/kernel"] | |
| logits = zc_norm(Z, p["reader/final_norm/scale"]) @ p["reader/out/kernel"] + p["reader/out/bias"] | |
| return logits.softmax(-1) | |
| # ------------------------------------------------------------ slots/memory | |
| def fourier(self, b): | |
| r = torch.arange(8, dtype=torch.float32) | |
| parts = [] | |
| for c in range(4): | |
| a = (2.0 ** r) * math.pi * float(b[c]) | |
| parts += [a.sin(), a.cos()] | |
| return torch.cat(parts) | |
| def memory(self, enc, objects, texts, vocab): | |
| p = self.p | |
| E, D, valid = enc["E"][0], enc["D"][0], enc["valid"] | |
| rows, cols = torch.arange(256) // 16, torch.arange(256) % 16 | |
| cells = E + torch.sigmoid(p["pe_gate"]) * sincos_2d(rows, cols) | |
| slots, refs = [], [] | |
| centres = (torch.arange(16) + 0.5) / 16 | |
| for kind, boxes in (("object", objects), ("text", texts)): | |
| for k, bx in enumerate(boxes): | |
| b = bx["box"] | |
| i, j = bx["cell"] | |
| inside = ((centres[:, None] >= b[1]) & (centres[:, None] <= b[3]) | |
| & (centres[None, :] >= b[0]) & (centres[None, :] <= b[2])).flatten() | |
| ci = min(15, int((b[1] + b[3]) / 2 * 16)) * 16 + min(15, int((b[0] + b[2]) / 2 * 16)) | |
| inside[ci] = True | |
| inside &= valid | |
| if not inside.any(): | |
| inside[ci] = True | |
| pooled = E[inside].mean(0) | |
| z = torch.cat([D[:, i, j], self.fourier(b)]) @ p["slot_proj/kernel"] + pooled @ p["slot_pool/kernel"] | |
| rid = vocab.ref_id(kind, k) | |
| slots.append(zc_norm(z, p["slot_norm/scale"]) + math.sqrt(D_MODEL) * self.T[rid]) | |
| refs.append(rid) | |
| M = torch.cat([cells] + ([torch.stack(slots)] if slots else []), 0) | |
| mvalid = torch.cat([valid, torch.ones(len(slots), dtype=torch.bool)]) | |
| return M, mvalid, refs | |
| class Decoder: | |
| """Incremental decoder with self-attention, tap, engram and cross-attention caches.""" | |
| def __init__(self, model: Mirage, M, mvalid): | |
| self.m, self.p = model, model.p | |
| p = self.p | |
| pre = "stack/layers/block/" | |
| self.ck, self.cv = [], [] | |
| for l in range(8): | |
| k = (M @ p[pre + "cross_attn/k_proj/kernel"][l]).view(-1, 8, 48).transpose(0, 1) | |
| self.ck.append(zc_norm(k, p[pre + "cross_attn/k_norm/scale"][l])) | |
| self.cv.append((M @ p[pre + "cross_attn/v_proj/kernel"][l]).view(-1, 8, 64).transpose(0, 1)) | |
| self.mvalid = mvalid | |
| self.tokens = [] | |
| self.K = [[] for _ in range(8)] | |
| self.V = [[] for _ in range(8)] | |
| self.raw = [{"q": [], "k": [], "v": []} for _ in range(8)] | |
| self.engram_v = {s: [] for s in ENGRAM_SITES} | |
| def _engram(self, site_i, layer, t): | |
| p = self.p | |
| pre = f"engrams_{site_i}/" | |
| rows = [] | |
| for oi, order in enumerate(ENGRAM_ORDERS): | |
| for h in range(ENGRAM_HEADS): | |
| tbl = p[pre + "embedding"][oi * ENGRAM_HEADS + h] | |
| if t - (order - 1) < 0: | |
| rows.append(torch.zeros(tbl.shape[1])) | |
| else: | |
| rows.append(tbl[engram_index(self.tokens, t, oi, order, h)]) | |
| e = torch.cat(rows) | |
| k = e @ p[pre + "key_proj/kernel"] | |
| vraw = e @ p[pre + "value_proj/kernel"] | |
| hist = self.engram_v[layer] | |
| hist.append(vraw) | |
| taps = p[pre + "taps"] | |
| v = torch.zeros_like(vraw) | |
| for j in range(ENGRAM_TAPS): | |
| if t - ENGRAM_DILATION * j >= 0: | |
| v = v + taps[j] * hist[t - ENGRAM_DILATION * j] | |
| return k, v | |
| def _hmlp(self, x, l): | |
| p, pre = self.p, "stack/layers/block/hadamard_mlp/" | |
| def kron(z, a, b): | |
| return (a.T @ z.view(16, 32) @ b).reshape(-1) | |
| cond = 1 + torch.softmax(x @ p[pre + "cond_v"][l], -1) @ p[pre + "cond_u"][l] | |
| z = kron(p[pre + "d1"][l] * x, p[pre + "w1a"][l], p[pre + "w1b"][l])[self.m.perm1] | |
| z = kron(F.silu(p[pre + "d2"][l] * cond * z + p[pre + "b2"][l]), p[pre + "w2a"][l], p[pre + "w2b"][l])[self.m.perm2] | |
| z = kron(p[pre + "d3"][l] * z, p[pre + "w3a"][l], p[pre + "w3b"][l]) | |
| return p[pre + "d4"][l] * z | |
| def _block(self, l, u, t): | |
| p, pre = self.p, "stack/layers/block/" | |
| x = u | |
| if l in ENGRAM_SITES: | |
| k, v = self._engram(ENGRAM_SITES.index(l), l, t) | |
| alpha = torch.sigmoid((rms_unit(x) * rms_unit(k)).sum() / math.sqrt(D_MODEL)) | |
| x = x + alpha * v | |
| skip = x | |
| xn = zc_norm(x, p[pre + "ZCRMSNorm_0/scale"][l]) | |
| R = self.raw[l] | |
| out = {} | |
| for name, n_heads, hd in (("q", 8, 48), ("k", 2, 48), ("v", 2, 64)): | |
| r = xn @ p[pre + f"self_attn/{name}_proj/kernel"][l] | |
| R[name].append(r) | |
| taps = p[pre + f"self_attn/{name}_taps"][l] | |
| acc = torch.zeros_like(r) | |
| for j in range(3): | |
| if t - j >= 0: | |
| acc = acc + taps[j] * R[name][t - j] | |
| out[name] = acc.view(n_heads, hd) | |
| q = zc_norm(out["q"], p[pre + "self_attn/q_norm/scale"][l]) | |
| k = zc_norm(out["k"], p[pre + "self_attn/k_norm/scale"][l]) | |
| pos = torch.tensor(t) | |
| q, k = rope_half(q, pos, 100000.0), rope_half(k, pos, 100000.0) | |
| self.K[l].append(k) | |
| self.V[l].append(out["v"]) | |
| Ks, Vs = torch.stack(self.K[l], 1), torch.stack(self.V[l], 1) # [2, t+1, d] | |
| Ks, Vs = Ks.repeat_interleave(4, 0), Vs.repeat_interleave(4, 0) # [8, t+1, d] | |
| att = torch.softmax((Ks @ q[:, :, None])[..., 0] / math.sqrt(48), -1) # [8, t+1] | |
| o = (att[:, None, :] @ Vs)[:, 0, :].reshape(-1) | |
| A = (o * torch.sigmoid(xn @ p[pre + "self_attn/gate_proj/kernel"][l])) @ p[pre + "self_attn/out_proj/kernel"][l] | |
| x = skip + torch.sigmoid(p[pre + "attn_gate"][l]) * zc_norm(A, p[pre + "post_attn_norm/scale"][l]) | |
| skip = x | |
| xc = zc_norm(x, p[pre + "cross_norm/scale"][l]) | |
| qc = zc_norm((xc @ p[pre + "cross_attn/q_proj/kernel"][l]).view(8, 48), p[pre + "cross_attn/q_norm/scale"][l]) | |
| s = (self.ck[l] @ qc[:, :, None])[..., 0] / math.sqrt(48) | |
| s = s.masked_fill(~self.mvalid[None], float("-inf")) | |
| oc = (torch.softmax(s, -1)[:, None, :] @ self.cv[l])[:, 0, :].reshape(-1) | |
| CA = (oc * torch.sigmoid(xc @ p[pre + "cross_attn/gate_proj/kernel"][l])) @ p[pre + "cross_attn/out_proj/kernel"][l] | |
| x = skip + torch.sigmoid(p[pre + "cross_gate"][l]) * zc_norm(CA, p[pre + "post_cross_norm/scale"][l]) | |
| return x + self._hmlp(zc_norm(x, p[pre + "pre_hada_norm/scale"][l]), l) | |
| def step(self, tok): | |
| p = self.p | |
| t = len(self.tokens) | |
| self.tokens.append(int(tok)) | |
| X = (math.sqrt(D_MODEL) * self.m.T[tok]).expand(4, D_MODEL).clone() | |
| for l in range(8): | |
| nx = rms_unit(X.reshape(-1)) | |
| lane = torch.zeros(4) | |
| lane[l % 4] = 1.0 | |
| hpre = torch.sigmoid(p["stack/mhc_a_pre"][l] * (nx @ p["stack/mhc_phi_pre"][l]) + p["stack/mhc_b_pre"][l] + 8 * lane - 4) | |
| u = (hpre[:, None] * X).sum(0) | |
| y = self._block(l, u, t) - u | |
| hpost = 2 * torch.sigmoid(p["stack/mhc_a_post"][l] * (nx @ p["stack/mhc_phi_post"][l]) + p["stack/mhc_b_post"][l] - 4 * (1 - lane)) | |
| hres = sinkhorn(p["stack/mhc_a_res"][l] * (nx @ p["stack/mhc_phi_res"][l]).view(4, 4) + p["stack/mhc_b_res"][l]) | |
| X = hres @ X + hpost[:, None] * y | |
| h = zc_norm(X.mean(0), p["stack/final_norm/scale"]) | |
| return h @ self.m.T.T | |
| # Full-sequence (teacher-forced) decoder: logits for every position of a token sequence in one pass, | |
| # equivalent to Decoder stepped token by token; for training and for scoring weight variants. | |
| PRE = "stack/layers/block/" | |
| def shift(x, j): | |
| """Row t of the result is row t - j of x (zeros before the start).""" | |
| if j == 0: | |
| return x | |
| return torch.cat([x.new_zeros((j,) + x.shape[1:]), x[:-j]], 0) if j < x.shape[0] else torch.zeros_like(x) | |
| def engram_rows(p, tokens, site_i): | |
| T = len(tokens) | |
| tok = [int(t) for t in tokens] | |
| parts = [] | |
| for oi, order in enumerate(ENGRAM_ORDERS): | |
| for h in range(ENGRAM_HEADS): | |
| tbl = p[f"engrams_{site_i}/embedding"][oi * ENGRAM_HEADS + h] | |
| idx = torch.tensor([engram_index(tok, t, oi, order, h) for t in range(T)]) | |
| rows = tbl[idx] | |
| ok = (torch.arange(T) >= order - 1).float()[:, None] | |
| parts.append(rows * ok) | |
| return torch.cat(parts, -1) | |
| def hmlp(p, x, l, perm1, perm2): | |
| pre = PRE + "hadamard_mlp/" | |
| def kron(z, a, b): | |
| return torch.einsum("tij,ik,jl->tkl", z.reshape(z.shape[0], 16, 32), a, b).reshape(z.shape[0], -1) | |
| cond = 1 + torch.softmax(x @ p[pre + "cond_v"][l], -1) @ p[pre + "cond_u"][l] | |
| z = kron(p[pre + "d1"][l] * x, p[pre + "w1a"][l], p[pre + "w1b"][l])[:, perm1] | |
| z = kron(F.silu(p[pre + "d2"][l] * cond * z + p[pre + "b2"][l]), p[pre + "w2a"][l], p[pre + "w2b"][l])[:, perm2] | |
| z = kron(p[pre + "d3"][l] * z, p[pre + "w3a"][l], p[pre + "w3b"][l]) | |
| return p[pre + "d4"][l] * z | |
| def cross_kv(p, M): | |
| ck, cv = [], [] | |
| for l in range(8): | |
| k = (M @ p[PRE + "cross_attn/k_proj/kernel"][l]).view(-1, 8, 48).transpose(0, 1) | |
| ck.append(zc_norm(k, p[PRE + "cross_attn/k_norm/scale"][l])) | |
| cv.append((M @ p[PRE + "cross_attn/v_proj/kernel"][l]).view(-1, 8, 64).transpose(0, 1)) | |
| return ck, cv | |
| def decoder_forward(model, M, mvalid, tokens): | |
| """tokens: 1-D LongTensor. Returns logits [T, vocab]; row t predicts token t + 1.""" | |
| p = model.p | |
| T = tokens.shape[0] | |
| ck, cv = cross_kv(p, M) | |
| X = (math.sqrt(D_MODEL) * model.T[tokens])[:, None, :].expand(T, 4, D_MODEL) | |
| pos = torch.arange(T) | |
| causal = torch.ones(T, T, dtype=torch.bool).tril() | |
| eng = {s: engram_rows(p, tokens, i) for i, s in enumerate(ENGRAM_SITES)} | |
| for l in range(8): | |
| nx = rms_unit(X.reshape(T, -1)) | |
| lane = torch.zeros(4) | |
| lane[l % 4] = 1.0 | |
| hpre = torch.sigmoid(p["stack/mhc_a_pre"][l] * (nx @ p["stack/mhc_phi_pre"][l]) + p["stack/mhc_b_pre"][l] + 8 * lane - 4) | |
| u = (hpre[..., None] * X).sum(1) | |
| x = u | |
| if l in ENGRAM_SITES: | |
| i = ENGRAM_SITES.index(l) | |
| e = eng[l] | |
| k = e @ p[f"engrams_{i}/key_proj/kernel"] | |
| vraw = e @ p[f"engrams_{i}/value_proj/kernel"] | |
| taps = p[f"engrams_{i}/taps"] | |
| v = sum(taps[j] * shift(vraw, ENGRAM_DILATION * j) for j in range(ENGRAM_TAPS)) | |
| alpha = torch.sigmoid((rms_unit(x) * rms_unit(k)).sum(-1, keepdim=True) / math.sqrt(D_MODEL)) | |
| x = x + alpha * v | |
| skip = x | |
| xn = zc_norm(x, p[PRE + "ZCRMSNorm_0/scale"][l]) | |
| qkv = {} | |
| for name, heads, hd in (("q", 8, 48), ("k", 2, 48), ("v", 2, 64)): | |
| raw = xn @ p[PRE + f"self_attn/{name}_proj/kernel"][l] | |
| taps = p[PRE + f"self_attn/{name}_taps"][l] | |
| qkv[name] = sum(taps[j] * shift(raw, j) for j in range(3)).view(T, heads, hd) | |
| q = rope_half(zc_norm(qkv["q"], p[PRE + "self_attn/q_norm/scale"][l]), pos[:, None], 100000.0) | |
| k = rope_half(zc_norm(qkv["k"], p[PRE + "self_attn/k_norm/scale"][l]), pos[:, None], 100000.0) | |
| q, k, v = q.transpose(0, 1), k.transpose(0, 1).repeat_interleave(4, 0), qkv["v"].transpose(0, 1).repeat_interleave(4, 0) | |
| s = (q @ k.transpose(-1, -2)) / math.sqrt(48) | |
| o = (s.masked_fill(~causal, float("-inf")).softmax(-1) @ v).transpose(0, 1).reshape(T, -1) | |
| A = (o * torch.sigmoid(xn @ p[PRE + "self_attn/gate_proj/kernel"][l])) @ p[PRE + "self_attn/out_proj/kernel"][l] | |
| x = skip + torch.sigmoid(p[PRE + "attn_gate"][l]) * zc_norm(A, p[PRE + "post_attn_norm/scale"][l]) | |
| skip = x | |
| xc = zc_norm(x, p[PRE + "cross_norm/scale"][l]) | |
| qc = zc_norm((xc @ p[PRE + "cross_attn/q_proj/kernel"][l]).view(T, 8, 48), p[PRE + "cross_attn/q_norm/scale"][l]).transpose(0, 1) | |
| sc = (qc @ ck[l].transpose(-1, -2)) / math.sqrt(48) | |
| oc = (sc.masked_fill(~mvalid[None, None], float("-inf")).softmax(-1) @ cv[l]).transpose(0, 1).reshape(T, -1) | |
| CA = (oc * torch.sigmoid(xc @ p[PRE + "cross_attn/gate_proj/kernel"][l])) @ p[PRE + "cross_attn/out_proj/kernel"][l] | |
| x = skip + torch.sigmoid(p[PRE + "cross_gate"][l]) * zc_norm(CA, p[PRE + "post_cross_norm/scale"][l]) | |
| x = x + hmlp(p, zc_norm(x, p[PRE + "pre_hada_norm/scale"][l]), l, model.perm1, model.perm2) | |
| y = x - u | |
| hpost = 2 * torch.sigmoid(p["stack/mhc_a_post"][l] * (nx @ p["stack/mhc_phi_post"][l]) + p["stack/mhc_b_post"][l] - 4 * (1 - lane)) | |
| hres = sinkhorn(p["stack/mhc_a_res"][l] * (nx @ p["stack/mhc_phi_res"][l]).view(T, 4, 4) + p["stack/mhc_b_res"][l]) | |
| X = hres @ X + hpost[..., None] * y[:, None, :] | |
| h = zc_norm(X.mean(1), p["stack/final_norm/scale"]) | |
| return h @ model.T.T | |
| # ===================================================================================== text reader | |
| # Trainable, batched copy of Mirage's convolutional CTC text reader. Input: [B, 3, 48, 320] strips scaled to | |
| # [-1, 1]. Output: [B, 80, 307] logits, class 0 the CTC blank, class c > 0 the charset entry c - 1. | |
| PREFIX = "reader/" | |
| QUANTISED = {"pw_0/kernel", "pw_1/kernel", "pw_2/kernel", "out/kernel"} | \ | |
| {f"{k}_{b}/kernel" for k in ("up", "down") for b in range(4)} | |
| class Reader(nn.Module): | |
| """quantizer: a Quantizer; when given, the 2-bit matrices run through fake quantisation.""" | |
| def __init__(self, weights, quantizer=None, bits=2, act_quant=False): | |
| super().__init__() | |
| self.names = sorted(k for k in weights if k.startswith(PREFIX)) | |
| self.params = nn.ParameterDict({k.replace("/", "__"): nn.Parameter(weights[k].clone().float()) for k in self.names}) | |
| self.quantizer, self.bits, self.act_quant = quantizer, bits, act_quant | |
| def aq(self, x): | |
| """Engine-style activation quantisation: each 128-group rotated, int8 per group, rotated back.""" | |
| if not self.act_quant or self.quantizer is None: | |
| return x | |
| d = x.shape[-1] | |
| pc = -(-d // 128) * 128 | |
| with torch.no_grad(): | |
| xp = F.pad(x, (0, pc - d)) if pc != d else x | |
| g = xp.reshape(*xp.shape[:-1], pc // 128, 128) @ self.quantizer.H | |
| s = g.abs().amax(-1, keepdim=True).clamp_min(1e-12) / 127.0 | |
| q = (torch.round(g / s) * s) @ self.quantizer.H | |
| q = q.reshape(xp.shape)[..., :d] | |
| return x + (q - x).detach() | |
| def w(self, name): | |
| p = self.params[(PREFIX + name).replace("/", "__")] | |
| if self.quantizer is not None and name in QUANTISED: | |
| return self.quantizer.fake_quant_kernel(p, self.bits) | |
| return p | |
| def forward(self, J): | |
| G = cnorm(conv3(J, self.w("w0"), (2, 2)), self.w("stage_norm_0/scale")) | |
| for j, stride in enumerate(((2, 2), (2, 1), (2, 1))): | |
| y = dwconv(F.silu(G), self.w(f"dw_{j}"), stride) | |
| G = cnorm((self.aq(y.permute(0, 2, 3, 1)) @ self.w(f"pw_{j}/kernel")).permute(0, 3, 1, 2), self.w(f"stage_norm_{j + 1}/scale")) | |
| Z = F.silu(G).mean(2).transpose(1, 2) # [B, 80, 384] | |
| for b in range(4): | |
| h = zc_norm(Z, self.w(f"norm_{b}/scale")) | |
| k = self.w(f"mix_{b}")[0, :, 0, :] # [5, 384] | |
| h = F.conv1d(F.pad(h.transpose(1, 2), (2, 2)), k.T[:, None, :], groups=384).transpose(1, 2) | |
| vg = self.aq(h) @ self.w(f"up_{b}/kernel") | |
| Z = Z + self.aq(vg[..., :384] * F.silu(vg[..., 384:])) @ self.w(f"down_{b}/kernel") | |
| return self.aq(zc_norm(Z, self.w("final_norm/scale"))) @ self.w("out/kernel") + self.w("out/bias") | |
| def state(self): | |
| return {k: self.params[k.replace("/", "__")].detach().cpu().contiguous() for k in self.names} | |
| def greedy(logits, charset, lengths=None): | |
| ids = logits.argmax(-1).cpu() | |
| out = [] | |
| for b in range(ids.shape[0]): | |
| n = ids.shape[1] if lengths is None else int(lengths[b]) | |
| out.append(ctc_text(ids[b, :n].tolist(), charset)) | |
| return out | |
| # ===================================================================================== record assembly | |
| # Turns the decoder's compact lines plus boxes, masks and OCR strings into one JSON record. Decoder grammar, | |
| # as emitted by the model: | |
| # <think>trace</think> | |
| # S key=value; key=value | |
| # E<n> name @position *<objK><objJ> {key=value; key=value} | |
| # T *<txtK> | |
| # R <subject n> predicate <object n> | |
| # Unrecognised lines are kept under "unparsed". | |
| REF = re.compile(r"<(obj|txt)(\d+)>") | |
| ENTITY = re.compile(r"^E(\d+)\s+(.*)$") | |
| def _attrs(s): | |
| out = {} | |
| for part in s.split(";"): | |
| if "=" in part: | |
| k, v = part.split("=", 1) | |
| out[k.strip()] = v.strip() | |
| elif part.strip(): | |
| out.setdefault("note", part.strip()) | |
| return out | |
| def parse_decoded(text): | |
| trace = "" | |
| body = text | |
| m = re.search(r"<think>(.*?)</think>", text, flags=re.S) | |
| if m: | |
| trace = m.group(1).strip() | |
| body = text[m.end():] | |
| body = body.replace("</s>", "") | |
| rec = dict(scene={}, entities=[], text_refs=[], relations=[], events=[], unparsed=[], reasoning=trace) | |
| for line in body.split("\n"): | |
| line = line.strip() | |
| if not line: | |
| continue | |
| if line.startswith("S "): | |
| rec["scene"].update(_attrs(line[2:])) | |
| continue | |
| em = ENTITY.match(line) | |
| if em: | |
| n, rest = int(em.group(1)), em.group(2) | |
| attrs = {} | |
| am = re.search(r"\{(.*)\}\s*$", rest) | |
| if am: | |
| attrs = _attrs(am.group(1)) | |
| rest = rest[:am.start()].strip() | |
| refs = [(k, int(i)) for k, i in REF.findall(rest)] | |
| rest = REF.sub("", rest).replace("*", " ").strip() | |
| pos = None | |
| if "@" in rest: | |
| rest, pos = rest.split("@", 1) | |
| pos = pos.strip() | |
| name, count = rest.strip(), None | |
| cm = re.search(r"\s*#(\S+)$", name) | |
| if cm: | |
| name, count = name[:cm.start()].strip(), cm.group(1) | |
| ent = dict(n=n, name=name, position=pos, refs=refs, description=attrs) | |
| if count is not None: | |
| ent["count"] = int(count) if count.isdigit() else count | |
| rec["entities"].append(ent) | |
| continue | |
| if line.startswith("T"): | |
| rec["text_refs"] += [int(i) for k, i in REF.findall(line) if k == "txt"] | |
| continue | |
| rm = re.match(r"^R\s+(\d+)\s+(\S+)\s+(\d+)\s*$", line) | |
| if rm: | |
| rec["relations"].append(dict(subject=int(rm.group(1)), predicate=rm.group(2), object=int(rm.group(3)))) | |
| continue | |
| if line.startswith("V ") or line.startswith("EV"): | |
| rec["events"].append(line.split(" ", 1)[1] if " " in line else line) | |
| continue | |
| rec["unparsed"].append(line) | |
| return rec | |
| def assemble(decoded, objects, texts, ocr, to_1000, masks=None, fields=None): | |
| """objects/texts: detector proposals (score, box); ocr: strings per text proposal; to_1000: box mapper.""" | |
| rec = parse_decoded(decoded) | |
| out = dict(scene=rec["scene"], objects=[], regions=[], text=[], relations=rec["relations"], | |
| events=rec["events"], reasoning=rec["reasoning"]) | |
| for e in rec["entities"]: | |
| item = dict(id=e["n"], name=e["name"]) | |
| if "count" in e: | |
| item["count"] = e["count"] | |
| if e["position"]: | |
| item["position"] = e["position"] | |
| if e["description"]: | |
| item["description"] = e["description"] | |
| objs = [i for k, i in e["refs"] if k == "obj" and i < len(objects)] | |
| if objs: | |
| item["bbox"] = to_1000(objects[objs[0]]["box"]) | |
| item["score"] = round(objects[objs[0]]["score"], 3) | |
| if len(objs) > 1: | |
| item["count"] = len(objs) | |
| item["instances"] = [dict(bbox=to_1000(objects[i]["box"]), score=round(objects[i]["score"], 3)) for i in objs] | |
| if masks is not None and objs[0] in masks: | |
| item["mask"] = masks[objs[0]] | |
| out["objects"].append(item) | |
| else: | |
| out["regions"].append(item) | |
| for k, (b, s) in enumerate(zip(texts, ocr)): | |
| if not s: | |
| continue | |
| out["text"].append(dict(string=s, bbox=to_1000(b["box"]), score=round(b["score"], 3), | |
| mentioned=k in rec["text_refs"])) | |
| if rec["unparsed"]: | |
| out["unparsed"] = rec["unparsed"] | |
| return out | |
| class Pipeline: | |
| """describe(): image in, structured record out, using the reference model.""" | |
| def __init__(self, weights, archive, reader_weights=None): | |
| self.model = Mirage(weights) | |
| if reader_weights: | |
| for k, v in load_file(reader_weights).items(): | |
| self.model.p[k] = v.float() | |
| arc = Archive(archive) | |
| self.vocab = vocab_of(arc, self.model.T.shape[0]) | |
| self.charset = charset_of(arc) | |
| def _ocr(self, img, box_1000, meta): | |
| px = (box_1000[0] / 1000 * meta["W"], box_1000[1] / 1000 * meta["H"], | |
| box_1000[2] / 1000 * meta["W"], box_1000[3] / 1000 * meta["H"]) | |
| labels, probs = self.model.read(img, px) | |
| return ctc_text(labels, self.charset) | |
| def _polygons(self, prob, meta, fmt): | |
| import cv2 | |
| cv2.setNumThreads(1) | |
| m = (prob.numpy() > 0.5).astype(np.uint8) | |
| if not m.any(): | |
| return [] | |
| contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| scale = CANVAS / 128 / meta["r"] | |
| rings = [] | |
| for c in sorted(contours, key=cv2.contourArea, reverse=True): | |
| c = cv2.approxPolyDP(c, 0.6, True)[:, 0, :].astype(np.float64) | |
| if len(c) < 3: | |
| continue | |
| xs = np.clip((c[:, 0] + 0.5) * scale / meta["W"] * 1000, 0, 1000) | |
| ys = np.clip((c[:, 1] + 0.5) * scale / meta["H"] * 1000, 0, 1000) | |
| rings.append([int(round(v)) for xy in zip(xs, ys) for v in xy]) | |
| return rings[:1][0] if fmt == "largest" and rings else rings | |
| def describe(self, image, language="en", conf=0.5, text_conf=0.35, max_objects=20, max_text_lines=20, | |
| max_tokens=512, masks=True, mask_format="largest"): | |
| t0 = time.time() | |
| m = self.model | |
| enc = m.encode(image) | |
| meta = enc["meta"] | |
| det = m.detect(enc, conf=conf, text_conf=text_conf, nms_iou=0.5) | |
| objects, texts = det["object"][:max_objects], det["text"][:max_text_lines] | |
| to_1000 = lambda b: m.to_1000(b, meta) # noqa: E731 | |
| ocr = [self._ocr(image, to_1000(b["box"]), meta) for b in texts] | |
| t1 = time.time() | |
| M, mvalid, refs = m.memory(enc, objects, texts, self.vocab) | |
| dec = Decoder(m, M, mvalid) | |
| v = self.vocab | |
| banned = torch.zeros(v.size, dtype=torch.bool) | |
| banned[list(v.lang_ids.values())] = True | |
| banned[v.ref_base:] = True | |
| banned[refs] = False | |
| logits = None | |
| for tid in (v.bos, v.lang_ids[language]): | |
| logits = dec.step(tid) | |
| gen = [] | |
| for _ in range(max_tokens): | |
| nxt = int(logits.masked_fill(banned, float("-inf")).argmax()) | |
| if nxt == v.eos: | |
| break | |
| gen.append(nxt) | |
| logits = dec.step(nxt) | |
| t2 = time.time() | |
| decoded = v.decode(gen) | |
| mask_out = None | |
| if masks and objects: | |
| probs = m.masks(enc, objects) | |
| mask_out = {i: self._polygons(p, meta, mask_format) for i, p in enumerate(probs)} | |
| out = assemble(decoded, objects, texts, ocr, to_1000, masks=mask_out) | |
| out["timing_ms"] = dict(perception=round(1000 * (t1 - t0)), decode=round(1000 * (t2 - t1)), | |
| tokens=len(gen), total=round(1000 * (time.time() - t0))) | |
| return out | |
| # ===================================================================================== synthetic text | |
| # Synthetic text-line strips for the CTC reader: sampled strings rendered with system fonts and augmented. | |
| # Each sample is a [48, 320, 3] uint8 strip (text left-aligned, grey 127 padding), the used width U and the | |
| # label ids (charset index + 1). | |
| STRIP_H, STRIP_W = 48, 320 | |
| # Accented letters and marks of the seven description languages (en, de, fr, es, it, nl, pl). | |
| COMMON = set("àâäæçéèêëîïôœùûüÿÀÂÄÆÇÉÈÊËÎÏÔŒÙÛÜŸßáíñóúÁÍÑÓÚìòÌÒijIJąćęłńśźżĄĆĘŁŃŚŹŻ¡¿«»–—") | |
| # Letters found only in romanisations and obsolete orthographies. | |
| RARE = set("ĔĕĬĭŎŏĸʼnĿŀ") | |
| def letter_weight(c): | |
| return 4.0 if c in COMMON else 0.2 if c in RARE else 1.0 | |
| def font_paths(patterns): | |
| paths = [] | |
| for p in patterns: | |
| paths += glob.glob(p) | |
| return sorted(set(paths)) | |
| class FontPool: | |
| def __init__(self, paths, charset): | |
| self.paths, self.cover = [], [] | |
| probe = "\U0010fffd" | |
| for path in paths: | |
| try: | |
| f = ImageFont.truetype(path, 32) | |
| except OSError: | |
| continue | |
| missing = np.asarray(f.getmask(probe)) | |
| ok = set() | |
| for c in charset: | |
| if c == " ": | |
| ok.add(c) | |
| continue | |
| m = f.getmask(c) | |
| if m.size[0] == 0: | |
| continue | |
| arr = np.asarray(m) | |
| if arr.shape == missing.shape and np.array_equal(arr, missing): | |
| continue | |
| ok.add(c) | |
| self.paths.append(path) | |
| self.cover.append(ok) | |
| self.cache = {} | |
| def get(self, rng, text, size): | |
| need = set(text) | |
| cands = [i for i, c in enumerate(self.cover) if need <= c] | |
| if not cands: | |
| return None | |
| i = rng.choice(cands) | |
| key = (i, size) | |
| if key not in self.cache: | |
| self.cache[key] = ImageFont.truetype(self.paths[i], size) | |
| return self.cache[key] | |
| class TextSampler: | |
| def __init__(self, charset, dead, words, seed=0): | |
| self.charset = charset | |
| self.chars = [c for c in charset if c != " "] | |
| self.dead = [c for c in dead if c in charset] | |
| self.words = [w for w in words if w and all(ch in charset for ch in w)] | |
| self.variants = {} | |
| for c in self.chars: | |
| base = unicodedata.normalize("NFD", c)[0] | |
| if base != c and base.isalpha(): | |
| self.variants.setdefault(base, []).append(c) | |
| self.dead_by_base = {} | |
| for c in self.dead: | |
| base = unicodedata.normalize("NFD", c)[0] | |
| self.dead_by_base.setdefault(base, []).append(c) | |
| self.letters_with_dead = [b for b in self.dead_by_base] | |
| self.rng = random.Random(seed) | |
| def _accent(self, word, p_dead): | |
| r = self.rng | |
| out = [] | |
| for ch in word: | |
| if ch in self.dead_by_base and r.random() < p_dead: | |
| out.append(r.choice(self.dead_by_base[ch])) | |
| elif ch in self.variants and r.random() < 0.15: | |
| out.append(r.choice(self.variants[ch])) | |
| else: | |
| out.append(ch) | |
| return "".join(out) | |
| def _case(self, s): | |
| r = self.rng.random() | |
| if r < 0.3: | |
| return s.upper() | |
| if r < 0.45: | |
| return s.title() | |
| if r < 0.6: | |
| return s.lower() | |
| return s | |
| def sample(self): | |
| r = self.rng | |
| mode = r.random() | |
| if mode < 0.40: | |
| s = " ".join(self._case(self._accent(r.choice(self.words), 0.08)) for _ in range(r.randint(1, 3))) | |
| elif mode < 0.60: | |
| pool = self.chars + self.dead * 4 | |
| s = "".join(r.choice(pool) for _ in range(r.randint(2, 12))) | |
| elif mode < 0.75: | |
| words = [] | |
| for _ in range(r.randint(1, 3)): | |
| w = list(self._case(r.choice(self.words))) | |
| if not w: | |
| continue | |
| k = r.randrange(len(w)) | |
| w[k] = r.choice(self.dead) | |
| words.append("".join(w)) | |
| s = " ".join(words) or r.choice(self.dead) | |
| elif mode < 0.85: | |
| num = f"{r.randint(0, 9999)}{r.choice(['', ',', '.'])}{r.randint(0, 99):02d}" | |
| s = r.choice([ | |
| f"{num} {r.choice('€£¥$')}", f"{r.choice('€£¥$')}{num}", f"{r.randint(1, 31)}–{r.randint(1, 31)}", | |
| f"¡{self._case(r.choice(self.words))}!", f"¿{self._case(r.choice(self.words))}?", | |
| f"«{r.choice(self.words)}»", f"“{r.choice(self.words)}”", f"{r.choice(self.words)} – {r.choice(self.words)}", | |
| f"{r.randint(1, 12)}:{r.randint(0, 59):02d}", f"{r.randint(0, 100)}%", f"№{num}" if "№" in self.charset else f"#{num}", | |
| ]) | |
| else: | |
| s = " ".join(self._accent(r.choice(self.words), 0.15).upper() for _ in range(r.randint(1, 3))) | |
| s = "".join(ch for ch in s if ch in self.charset).strip() | |
| return s or r.choice(self.dead) | |
| class RealisticSampler(TextSampler): | |
| """Text with letter frequencies close to real writing in the charset's languages. | |
| Real words keep their own accents; contrast lines swap one letter for a sibling with the same base | |
| letter (Ê for E, Ĕ for E, ...) weighted by how common it is, so accents have to be read rather than | |
| guessed; a small share of random strings covers every class. | |
| """ | |
| def __init__(self, charset, dead, words, seed=0): | |
| super().__init__(charset, dead, words, seed) | |
| self.siblings = {} | |
| for c in self.chars: | |
| base = unicodedata.normalize("NFD", c)[0] | |
| if base.isalpha() and base in charset: | |
| self.siblings.setdefault(base, []).append(c) | |
| self.siblings = {b: [c for c in v if c != b] for b, v in self.siblings.items()} | |
| self.siblings = {b: v for b, v in self.siblings.items() if v} | |
| self.pool = self.chars | |
| self.pool_w = [letter_weight(c) for c in self.chars] | |
| def _swap(self, word): | |
| spots = [i for i, ch in enumerate(word) if ch in self.siblings or unicodedata.normalize("NFD", ch)[0] in self.siblings] | |
| if not spots: | |
| return word | |
| i = self.rng.choice(spots) | |
| base = unicodedata.normalize("NFD", word[i])[0] | |
| opts = self.siblings[base] | |
| c = self.rng.choices(opts, weights=[letter_weight(o) for o in opts])[0] | |
| return word[:i] + c + word[i + 1:] | |
| def sample(self): | |
| r = self.rng | |
| mode = r.random() | |
| if mode < 0.45: | |
| s = " ".join(self._case(r.choice(self.words)) for _ in range(r.randint(1, 3))) | |
| elif mode < 0.80: | |
| words = [self._case(r.choice(self.words)) for _ in range(r.randint(1, 3))] | |
| k = r.randrange(len(words)) | |
| words[k] = self._swap(words[k]) | |
| s = " ".join(words) | |
| elif mode < 0.90: | |
| s = "".join(r.choices(self.pool, weights=self.pool_w, k=r.randint(2, 12))) | |
| else: | |
| return super().sample() | |
| s = "".join(ch for ch in s if ch in self.charset).strip() | |
| return s or r.choice(self.chars) | |
| def _contrast_colors(rng): | |
| while True: | |
| a = np.array([rng.randint(0, 255) for _ in range(3)]) | |
| b = np.array([rng.randint(0, 255) for _ in range(3)]) | |
| la, lb = a @ [0.299, 0.587, 0.114], b @ [0.299, 0.587, 0.114] | |
| if abs(la - lb) > 70: | |
| return tuple(int(x) for x in a), tuple(int(x) for x in b) | |
| def _background(rng, w, h, color): | |
| kind = rng.random() | |
| base = np.zeros((h, w, 3), np.float32) + np.array(color, np.float32) | |
| if kind < 0.4: | |
| pass | |
| elif kind < 0.65: | |
| other = np.array([rng.randint(0, 255) for _ in range(3)], np.float32) | |
| t = np.linspace(0, 1, w if rng.random() < 0.5 else h) | |
| t = t[None, :, None] if t.shape[0] == w else t[:, None, None] | |
| base = base * (1 - 0.35 * t) + other * 0.35 * t | |
| elif kind < 0.85: | |
| base = base + np.random.default_rng(rng.randint(0, 2**31)).normal(0, rng.uniform(5, 25), (h, w, 1)) | |
| else: | |
| period = rng.randint(4, 20) | |
| stripes = ((np.arange(w) // period) % 2)[None, :, None] * rng.uniform(10, 30) | |
| base = base + stripes | |
| return Image.fromarray(np.clip(base, 0, 255).astype(np.uint8)) | |
| def render(text, fonts, rng, tight=0.0): | |
| """tight: share of lines cut as closely as a text detector's box, so accents can touch the edge.""" | |
| for _ in range(4): | |
| size = rng.randint(20, 64) | |
| font = fonts.get(rng, text, size) | |
| if font is None: | |
| return None | |
| x0, y0, x1, y1 = font.getbbox(text) | |
| tw, th = x1 - x0, y1 - y0 | |
| if tw <= 0 or th <= 0: | |
| return None | |
| if rng.random() < tight: | |
| mx, my = int(size * rng.uniform(0.0, 0.12)), int(size * rng.uniform(-0.03, 0.04)) | |
| else: | |
| mx, my = int(size * rng.uniform(0.05, 0.35)), int(size * rng.uniform(0.05, 0.3)) | |
| W, H = tw + 2 * mx, th + 2 * my | |
| U = max(1, round(STRIP_H * W / H)) | |
| if U <= STRIP_W: | |
| break | |
| text = text[: max(1, len(text) * STRIP_W // U - 1)].strip() or text[:1] | |
| else: | |
| return None | |
| fg, bg = _contrast_colors(rng) | |
| img = _background(rng, W, H, bg) | |
| ImageDraw.Draw(img).text((mx - x0, my - y0), text, font=font, fill=fg) | |
| if rng.random() < 0.3: | |
| img = img.rotate(rng.uniform(-2.5, 2.5), resample=Image.BILINEAR, expand=False, fillcolor=bg) | |
| if rng.random() < 0.3: | |
| img = img.filter(ImageFilter.GaussianBlur(rng.uniform(0.3, 1.2))) | |
| if rng.random() < 0.2: | |
| f = rng.uniform(0.35, 0.7) | |
| small = img.resize((max(1, int(W * f)), max(1, int(H * f))), Image.BILINEAR) | |
| img = small.resize((W, H), Image.BILINEAR) | |
| if rng.random() < 0.3: | |
| buf = io.BytesIO() | |
| img.save(buf, "JPEG", quality=rng.randint(30, 90)) | |
| img = Image.open(io.BytesIO(buf.getvalue())).convert("RGB") | |
| return text, img | |
| def to_strip(img): | |
| W, H = img.size | |
| U = max(1, round(STRIP_H * W / H)) | |
| if U > STRIP_W: | |
| return None, 0 | |
| arr = np.full((STRIP_H, STRIP_W, 3), 127, np.uint8) | |
| arr[:, :U] = np.asarray(img.resize((U, STRIP_H), Image.BILINEAR)) | |
| return arr, U | |
| class Generator: | |
| def __init__(self, charset, dead, words, font_globs, seed, style="uniform", tight=0.0): | |
| self.charset = charset | |
| self.index = {c: i + 1 for i, c in enumerate(charset)} | |
| self.fonts = FontPool(font_paths(font_globs), charset) | |
| self.text = (RealisticSampler if style == "realistic" else TextSampler)(charset, dead, words, seed) | |
| self.rng = random.Random(seed + 1) | |
| self.tight = tight | |
| def sample(self): | |
| while True: | |
| s = self.text.sample() | |
| out = render(s, self.fonts, self.rng, self.tight) | |
| if out is None: | |
| continue | |
| s, img = out | |
| strip, U = to_strip(img) | |
| if strip is None or not s: | |
| continue | |
| labels = [self.index[c] for c in s] | |
| frames = -(-U // 4) | |
| reps = sum(1 for a, b in zip(labels, labels[1:]) if a == b) | |
| if frames < len(labels) + reps: | |
| continue | |
| return strip, U, labels, s | |
| # ===================================================================================== reader evaluation | |
| # Word crops packed as archives of strips 48 pixels high (pix [48, total, 3], offsets, labels, sources), | |
| # read with the engine's window and frame-retention rules. | |
| def reader_windows(strip): | |
| U = strip.shape[1] | |
| starts = [0] if U <= 320 else list(range(0, U - 320 + 223, 224)) | |
| wins, keep = [], [] | |
| n_frames = -(-U // 4) | |
| for xs in starts: | |
| win = np.full((48, 320, 3), 127, np.uint8) | |
| seg = strip[:, xs:xs + 320] | |
| win[:, :seg.shape[1]] = seg | |
| lo = 0 if xs == 0 else 12 | |
| hi = -(-(U - xs) // 4) if xs + 320 >= U else 68 | |
| keep.append([t for t in range(lo, min(hi, 80)) if xs // 4 + t < n_frames]) | |
| wins.append(win) | |
| return wins, keep | |
| def load_packed(path): | |
| """The windows of a packed crop archive, with each window's word, its retained frames, the labels and sources.""" | |
| d = np.load(path, allow_pickle=False) | |
| pix, offs, labels, sources = d["pix"], d["offsets"], d["labels"], d["sources"] | |
| all_w, owner, keeps = [], [], [] | |
| for i in range(len(labels)): | |
| ws, ks = reader_windows(pix[:, offs[i]:offs[i + 1]]) | |
| all_w += ws | |
| owner += [i] * len(ws) | |
| keeps += ks | |
| return torch.from_numpy(np.stack(all_w)).permute(0, 3, 1, 2), owner, keeps, labels, sources | |
| def reader_paths(net, X, owner, keeps, n, dev): | |
| """The retained frame labels of every word, reading the windows X in batches of 512.""" | |
| ids = [] | |
| for b in range(0, X.shape[0], 512): | |
| ids.append(net(X[b:b + 512].to(dev).float() / 127.5 - 1.0).argmax(-1).cpu()) | |
| ids = torch.cat(ids).numpy() if ids else np.zeros((0, 80), np.int64) | |
| paths = [[] for _ in range(n)] | |
| for k, (o, kp) in enumerate(zip(owner, keeps)): | |
| paths[o] += ids[k, kp].tolist() | |
| return paths | |
| def evaluate_reader(net, sets, charset, dead_chars, dev): | |
| was = net.training | |
| net.eval() | |
| out = {} | |
| for X, owner, keeps, labels, sources in sets: | |
| for i, path in enumerate(reader_paths(net, X, owner, keeps, len(labels), dev)): | |
| pred, gt = ctc_text(path, charset), str(labels[i]) | |
| st = out.setdefault(str(sources[i]), dict(n=0, acc=0, exact=0, ed=0, len=0, dh=0, dt=0, fp=0, ah=0, at=0)) | |
| st["n"] += 1 | |
| st["acc"] += word_key(pred) == word_key(gt) | |
| st["exact"] += pred == gt | |
| st["ed"] += edit_distance(pred, gt) | |
| st["len"] += len(gt) | |
| for ch in set(gt) & dead_chars: | |
| st["dt"] += gt.count(ch) | |
| st["dh"] += min(gt.count(ch), pred.count(ch)) | |
| for ch in set(gt): | |
| if ord(ch) > 127 and ch.isalpha(): | |
| st["at"] += gt.count(ch) | |
| st["ah"] += min(gt.count(ch), pred.count(ch)) | |
| st["fp"] += sum(1 for ch in pred if ch in dead_chars and ch not in gt) | |
| if was: | |
| net.train() | |
| return {src: dict(n=st["n"], word_acc=round(st["acc"] / st["n"], 4), exact=round(st["exact"] / st["n"], 4), | |
| cer=round(st["ed"] / max(st["len"], 1), 4), dead_recall=round(st["dh"] / max(st["dt"], 1), 4), | |
| accent_recall=round(st["ah"] / max(st["at"], 1), 4), dead_false=st["fp"]) for src, st in out.items()} | |
| # ===================================================================================== commands: model | |
| def pack_indices(idx, bits): | |
| flat = idx.reshape(-1).astype(np.uint8) | |
| per = 8 // bits | |
| flat = np.concatenate([flat, np.zeros((-len(flat)) % per, np.uint8)]) | |
| out = np.zeros(len(flat) // per, np.uint8) | |
| for k in range(per): | |
| out |= (flat[k::per] << (bits * k)).astype(np.uint8) | |
| return out | |
| def cq_blob_from_latent(R, bits, quantizer): | |
| """R: [out, in] float32. Same arithmetic as Quantizer (the training forward).""" | |
| rows, cols = R.shape | |
| pc = -(-cols // 128) * 128 | |
| Rp = torch.nn.functional.pad(R.to(quantizer.H.device), (0, pc - cols)) | |
| g = Rp.reshape(rows, pc // 128, 128) @ quantizer.H | |
| n = g.norm(dim=-1, keepdim=True) | |
| idx = torch.bucketize((g / n.clamp_min(1e-30)).contiguous(), quantizer.mid[bits]) | |
| return pack_indices(idx.cpu().numpy(), bits).tobytes() + n.half().cpu().numpy().reshape(rows, -1).tobytes() | |
| def compact_scales(blob, rows, cols, bits): | |
| """Re-encode a cq blob's fp16 group scales as 8-bit log codes with the two f32 bounds appended.""" | |
| pc = -(-cols // 128) * 128 | |
| ib = rows * pc * bits // 8 | |
| nrm = np.frombuffer(blob, np.float16, count=rows * (pc // 128), offset=ib).astype(np.float64) | |
| nz = nrm[nrm > 0] | |
| lo, hi = (np.float32(np.log(nz.min())), np.float32(np.log(nz.max()))) if len(nz) else (np.float32(0), np.float32(0)) | |
| span = float(hi) - float(lo) | |
| pos = np.log(np.where(nrm > 0, nrm, 1.0)) - float(lo) | |
| steps = np.round(pos / span * 254.0) if span > 0 else np.zeros_like(pos) | |
| code = np.where(nrm > 0, np.clip(steps, 0, 254) + 1, 0).astype(np.uint8) | |
| return bytes(blob[:ib]) + code.tobytes() + np.array([lo, hi], np.float32).tobytes() | |
| def int_rows(a, bits, keep=None): | |
| """[rows, cols] float -> signed integers of `bits` bits in one LSB-first bit stream, then one f32 scale | |
| per row (the largest magnitude of a row maps to the largest integer). Entries outside the boolean mask | |
| `keep` are stored as zero.""" | |
| a = np.asarray(a, np.float64).reshape(a.shape[0], -1) | |
| top = 2 ** (bits - 1) - 1 | |
| s = np.abs(a).max(1) / top | |
| q = np.round(a / np.where(s > 0, s, 1.0)[:, None]).clip(-top, top).astype(np.int64) | |
| if keep is not None: | |
| q = np.where(keep, q, 0) | |
| u = (q & ((1 << bits) - 1)).reshape(-1) | |
| stream = ((u[:, None] >> np.arange(bits)) & 1).astype(np.uint8).reshape(-1) | |
| return np.packbits(stream, bitorder="little").tobytes() + s.astype(np.float32).tobytes() | |
| BAND = ("hconv/gate_conv", "hconv/value_conv", "hconv/out_conv") | |
| FP32 = {"dense/out/kernel", "dense/bias", "dense/protos/kernel", "dense/coef/kernel", "reader/out/bias"} | |
| def band_taps(rows, k): | |
| """[rows, k] mask of the taps a band applies to: tap t of output j reads channel j + t - (k - 1) / 2.""" | |
| j, t = np.arange(rows)[:, None], np.arange(k)[None, :] | |
| c = j + t - (k - 1) // 2 | |
| return (c >= 0) & (c < rows) | |
| def cmd_export(argv): | |
| """Write an .oasis weight file for the Oasis engine. | |
| Layout: 8-byte magic "OASIS\\0v1", u32 header length, UTF-8 header of "key value" lines, u32 tensor count, | |
| then per tensor: u16 name length, name, u8 dtype, u8 bits, u8 ndim, u8 flags, u32 shape[ndim], u64 offset, | |
| u64 nbytes; blobs follow. Flag 1 marks a deflated blob (u64 length of the stored bytes, then a zlib stream of | |
| them); other blobs are 64-byte aligned. Dtypes: 0 f32, 1 f16, 2 cq, 3 raw bytes, 4 integer rows (signed | |
| integers of `bits` bits in one LSB-first bit stream, then one f32 scale per row), 5 cq with 8-bit scales. A cq | |
| tensor is [out, in]: every row's packed indices (LSB first, columns padded to 128), then one scale per row and | |
| 128-column group, as fp16 (dtype 2) or as an 8-bit code (dtype 5) followed by two f32 values lo and hi: code 0 | |
| is zero and codes 1..255 step evenly in log scale from exp(lo) to exp(hi). --band-bits stores the HConvGLU | |
| bands as integer rows, with the taps that fall outside the channel vector, which the engine never reads, | |
| stored as zero; --scale-bits 8 stores the group scales of every quantised matrix outside --keep-scales as | |
| codes. | |
| Quantised tensors whose latent weights match Mirage's are copied from mirage.cact; others are quantised here | |
| with the training-time quantiser, on the device the weights were trained on (--device, CUDA when available): | |
| quantisation-aware training leaves many latent weights at codebook boundaries, where the backend's float32 | |
| rounding decides the index.""" | |
| ap = parser("export", cmd_export.__doc__) | |
| ap.add_argument("--weights", help="float32 weights to export (default: Mirage's)") | |
| ap.add_argument("--base", help="Mirage's float32 weights (default: from --mirage)") | |
| ap.add_argument("--archive", help="Mirage's archive (default: from --mirage)") | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--name", default="oasis") | |
| ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") | |
| ap.add_argument("--band-bits", type=int, choices=[8, 10, 12, 16], default=12) | |
| ap.add_argument("--scale-bits", type=int, choices=[8, 16], default=8) | |
| ap.add_argument("--keep-scales", default="reader/", help="comma list of name prefixes whose group scales stay fp16") | |
| ap.add_argument("--no-deflate", action="store_true", help="store every tensor as it is, without zlib") | |
| args = ap.parse_args(argv) | |
| keep_scales = tuple(p for p in args.keep_scales.split(",") if p) | |
| arc = Archive(args.archive or mirage_archive(args)) | |
| w = {k: v.float() for k, v in load_file(args.weights or mirage_weights(args)).items() if not k.startswith(("mtp_", "name_head"))} | |
| base = load_file(args.base or mirage_weights(args)) | |
| vocab = vocab_of(arc) | |
| quantizer = Quantizer(arc.codebook, args.device) | |
| tensors = [] # (name, dtype, bits, shape, bytes) | |
| quantised = {} | |
| for item in cq_plan(): | |
| key = item["name"] if item["layer"] is None else f"{item['name']}/{item['layer']}" | |
| src = w[item["name"]] if item["layer"] is None else w[item["name"]][item["layer"]] | |
| ref_src = base[item["name"]] if item["layer"] is None else base[item["name"]][item["layer"]] | |
| e = arc.entries[item["entry"]] | |
| if torch.equal(src.float(), ref_src.float()): | |
| blob = arc.raw(item["entry"]) | |
| else: | |
| R = plan_get({item["name"]: w[item["name"]]}, item).float() | |
| blob = cq_blob_from_latent(R, item["bits"], quantizer) | |
| quantised.setdefault(item["name"], []).append(item["layer"]) | |
| rows, cols = int(np.prod(e.shape[:-1])), int(e.shape[-1]) | |
| if args.scale_bits == 8 and not item["name"].startswith(keep_scales): | |
| tensors.append((key, 5, e.bits, list(e.shape), compact_scales(blob, rows, cols, e.bits))) | |
| else: | |
| tensors.append((key, 2, e.bits, list(e.shape), blob)) | |
| for k, v in sorted(w.items()): | |
| if k in quantised: | |
| continue | |
| if k == "embedding/embedding": | |
| continue | |
| if k in FP32: | |
| tensors.append((k, 0, 0, list(v.shape), v.float().numpy().tobytes())) | |
| continue | |
| if args.band_bits < 16 and k.endswith(BAND): | |
| keep = band_taps(*v.shape) | |
| tensors.append((k, 4, args.band_bits, list(v.shape), int_rows(v.numpy(), args.band_bits, keep))) | |
| continue | |
| if k.startswith("stack/layers/block/") and v.dim() >= 1 and v.shape[0] == 8: | |
| for l in range(8): | |
| tensors.append((f"{k}/{l}", 1, 0, list(v[l].shape) or [1], v[l].half().numpy().tobytes())) | |
| else: | |
| tensors.append((k, 1, 0, list(v.shape) or [1], v.half().numpy().tobytes())) | |
| tensors.append(("hmlp/perm1", 0, 0, [512], np.random.RandomState(11).permutation(512).astype(np.float32).tobytes())) | |
| tensors.append(("hmlp/perm2", 0, 0, [512], np.random.RandomState(13).permutation(512).astype(np.float32).tobytes())) | |
| tensors.append(("quant/codebook", 0, 0, [len(arc.codebook)], np.asarray(arc.codebook, np.float32).tobytes())) | |
| charset = arc.array(CHARSET_ENTRY).astype(np.float32) | |
| tensors.append(("reader/charset", 0, 0, [len(charset)], charset.tobytes())) | |
| pieces = bytearray(struct.pack("<I", vocab.n_pieces)) | |
| for p, t in zip(vocab.pieces, vocab.types): | |
| b = p.encode("utf-8") | |
| pieces += struct.pack("<BH", t, len(b)) + b | |
| tensors.append(("tokenizer/pieces", 3, 0, [len(pieces)], bytes(pieces))) | |
| langs = " ".join(f"{k}:{v}" for k, v in vocab.lang_ids.items()) | |
| # The vocabulary carries eleven language tags; Mirage was trained to describe in these seven. | |
| languages = "en de fr es it nl pl" | |
| header = "\n".join([ | |
| f"name {args.name}", "format 1", "d_model 512", f"vocab {vocab.size}", f"n_pieces {vocab.n_pieces}", | |
| f"bos {vocab.bos}", f"eos {vocab.eos}", f"pad {vocab.pad}", f"ref_base {vocab.ref_base}", "n_obj_refs 64", | |
| "n_txt_refs 32", f"langs {langs}", f"languages {languages}", "group 128", "canvas 512", "engram_sites 3 7", | |
| "engram_slots 18432", "rope_dec 100000", "rope_enc 10000", "codebook_bits 2:0 3:4 4:12", | |
| ]).encode("utf-8") | |
| # A tensor whose bytes deflate by more than half a percent is stored as flag 1: its u64 length, then the | |
| # zlib stream; the engine inflates it at load. Stored tensors start on 64-byte boundaries, deflated ones | |
| # anywhere. | |
| stored = [] | |
| for name, dt, bits, shape, blob in tensors: | |
| flags = 0 | |
| if not args.no_deflate: | |
| z = zlib.compress(bytes(blob), 9) | |
| if len(z) + 8 < len(blob) * 0.995: | |
| blob, flags = struct.pack("<Q", len(blob)) + z, 1 | |
| stored.append((name, dt, bits, shape, blob, flags)) | |
| offsets, pos = [], 0 | |
| for name, dt, bits, shape, blob, flags in stored: | |
| if not flags: | |
| pos = -(-pos // 64) * 64 | |
| offsets.append(pos) | |
| pos += len(blob) | |
| head = bytearray(b"OASIS\0v1") + struct.pack("<I", len(header)) + header + struct.pack("<I", len(stored)) | |
| for (name, dt, bits, shape, blob, flags), off in zip(stored, offsets): | |
| nb = name.encode("utf-8") | |
| head += struct.pack("<H", len(nb)) + nb + struct.pack("<BBBB", dt, bits, len(shape), flags) | |
| head += struct.pack(f"<{len(shape)}I", *shape) + struct.pack("<QQ", off, len(blob)) | |
| data_start = -(-len(head) // 64) * 64 | |
| with open(args.out, "wb") as f: | |
| f.write(head + b"\0" * (data_start - len(head))) | |
| for (name, dt, bits, shape, blob, flags), off in zip(stored, offsets): | |
| f.seek(data_start + off) | |
| f.write(blob) | |
| size = os.path.getsize(args.out) | |
| n_q = sum(1 for t in tensors if t[1] in (2, 5)) | |
| print(f"wrote {args.out}: {len(tensors)} tensors ({n_q} quantised on {args.device}), {size / 1e6:.2f} MB, data at {data_start}") | |
| def cmd_checkpoint(argv): | |
| """Assemble an Oasis float32 checkpoint: Mirage's weights with replacement components, such as a retrained | |
| reader. Training-only heads (multi-token prediction, name head) are dropped. The metadata records the base | |
| checkpoint's hash and every replacement file's hash.""" | |
| ap = parser("checkpoint", cmd_checkpoint.__doc__) | |
| ap.add_argument("--base", help="Mirage's float32 weights (default: from --mirage)") | |
| ap.add_argument("--replace", action="append", default=[], help="a safetensors file of replacement tensors; repeatable") | |
| ap.add_argument("--dead-offset", type=float, default=0.0, help="added to the reader biases of classes the base never trained") | |
| ap.add_argument("--out", required=True) | |
| args = ap.parse_args(argv) | |
| base_path = args.base or mirage_weights(args) | |
| w = {k: v.float() for k, v in load_file(base_path).items() if not k.startswith(("mtp_", "name_head"))} | |
| dead = dead_classes(w) | |
| meta = {"base": os.path.basename(base_path), "base_sha256": sha256(base_path), "replacements": []} | |
| for path in args.replace: | |
| r = load_file(path) | |
| unknown = [k for k in r if k not in w] | |
| if unknown: | |
| raise SystemExit(f"{path} has tensors the base lacks: {unknown[:5]}") | |
| w.update({k: v.float() for k, v in r.items()}) | |
| meta["replacements"].append({"file": os.path.basename(path), "sha256": sha256(path), "tensors": len(r)}) | |
| if args.dead_offset: | |
| w["reader/out/bias"][dead] += args.dead_offset | |
| meta["reader_dead_class_bias_offset"] = args.dead_offset | |
| os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True) | |
| save_file({k: v.contiguous() for k, v in w.items()}, args.out, metadata={"oasis": json.dumps(meta)}) | |
| print(json.dumps(meta, indent=1)) | |
| print(f"wrote {args.out}: {len(w)} tensors, {sum(v.numel() for v in w.values()):,} parameters") | |
| def cmd_run(argv): | |
| """Run the reference forward pass on one image: detections, OCR, masks and a greedy decoder probe. Writes | |
| annotated.png, the two heat maps and result.json to --out (default: runs/IMAGE_NAME).""" | |
| ap = parser("run", cmd_run.__doc__) | |
| ap.add_argument("image") | |
| ap.add_argument("--weights", help="float32 weights (default: Mirage's)") | |
| ap.add_argument("--archive", help="Mirage's archive (default: from --mirage)") | |
| ap.add_argument("--prefix", default="<s>", help="decoder prompt: comma list of pieces or token ids") | |
| ap.add_argument("--max-new", type=int, default=200) | |
| ap.add_argument("--conf", type=float, default=0.3) | |
| ap.add_argument("--text-conf", type=float, default=0.3) | |
| ap.add_argument("--out", default=None) | |
| ap.add_argument("--set", action="append", default=[], help="OPTS override key=value") | |
| args = ap.parse_args(argv) | |
| for kv in args.set: | |
| k, v = kv.split("=", 1) | |
| OPTS[k] = type(OPTS[k])(eval(v, {}, {})) if not isinstance(OPTS[k], str) else v | |
| arc = Archive(args.archive or mirage_archive(args)) | |
| vocab = vocab_of(arc) | |
| charset = charset_of(arc) | |
| t0 = time.time() | |
| model = Mirage(args.weights or mirage_weights(args)) | |
| img = Image.open(args.image).convert("RGB") | |
| enc = model.encode(img) | |
| t1 = time.time() | |
| det = model.detect(enc, conf=args.conf, text_conf=args.text_conf, nms_iou=0.5) | |
| meta = enc["meta"] | |
| print(f"encode {t1 - t0:.2f}s image {meta['W']}x{meta['H']} canvas valid {meta['w']}x{meta['h']}") | |
| for kind in ("object", "text"): | |
| print(f"-- {kind}: {len(det[kind])}") | |
| for k, b in enumerate(det[kind][:20]): | |
| print(f" {k:2d} score {b['score']:.3f} cell {b['cell']} box1000 {model.to_1000(b['box'], meta)}") | |
| ocr = [] | |
| for k, b in enumerate(det["text"]): | |
| bb = model.to_1000(b["box"], meta) | |
| px = (bb[0] / 1000 * meta["W"], bb[1] / 1000 * meta["H"], bb[2] / 1000 * meta["W"], bb[3] / 1000 * meta["H"]) | |
| labels, _ = model.read(img, px) | |
| ocr.append(ctc_text(labels, charset)) | |
| print(f" ocr {k:2d}: {ocr[-1]!r}") | |
| masks = model.masks(enc, det["object"][:8]) | |
| out_dir = args.out or os.path.join(HERE, "runs", os.path.splitext(os.path.basename(args.image))[0]) | |
| os.makedirs(out_dir, exist_ok=True) | |
| vis = img.copy() | |
| dr = ImageDraw.Draw(vis, "RGBA") | |
| for k, (b, m) in enumerate(zip(det["object"][:8], masks)): | |
| mm = (m.numpy() > 0.5).astype(np.uint8) * 255 | |
| sc = CANVAS / meta["r"] / 128 | |
| mimg = Image.fromarray(mm).resize((int(128 * sc), int(128 * sc)), Image.NEAREST).crop((0, 0, meta["W"], meta["H"])) | |
| col = [(230, 25, 75), (60, 180, 75), (0, 130, 200), (245, 130, 48), (145, 30, 180), (70, 240, 240), (240, 50, 230), (210, 245, 60)][k] | |
| vis.paste(Image.new("RGB", vis.size, col), (0, 0), mimg.point(lambda v: 90 if v else 0)) | |
| for kind, color in (("object", (255, 0, 0, 255)), ("text", (255, 205, 0, 255))): | |
| for k, b in enumerate(det[kind]): | |
| bb = model.to_1000(b["box"], meta) | |
| xy = [bb[0] / 1000 * meta["W"], bb[1] / 1000 * meta["H"], bb[2] / 1000 * meta["W"], bb[3] / 1000 * meta["H"]] | |
| dr.rectangle(xy, outline=color, width=2) | |
| dr.text((xy[0] + 2, xy[1] + 2), f"{kind[0]}{k} {b['score']:.2f}", fill=color) | |
| vis.save(os.path.join(out_dir, "annotated.png")) | |
| heat = torch.sigmoid(det["raw"][..., :2]).numpy() | |
| for ch, name in ((0, "heat0"), (1, "heat1")): | |
| Image.fromarray((heat[..., ch] * 255).astype(np.uint8)).resize((256, 256), Image.NEAREST).save(os.path.join(out_dir, f"{name}.png")) | |
| M, mvalid, refs = model.memory(enc, det["object"], det["text"], vocab) | |
| dec = Decoder(model, M, mvalid) | |
| prefix = [] | |
| for item in args.prefix.split(","): | |
| item = item.strip() | |
| if not item: | |
| continue | |
| prefix.append(int(item) if item.isdigit() else vocab.id_of[item]) | |
| banned = torch.zeros(vocab.size, dtype=torch.bool) | |
| for i in vocab.lang_ids.values(): | |
| banned[i] = True | |
| banned[vocab.ref_base:] = True | |
| for r in refs: | |
| banned[r] = False | |
| t2 = time.time() | |
| logits = None | |
| for tid in prefix: | |
| logits = dec.step(tid) | |
| gen = [] | |
| for _ in range(args.max_new): | |
| lg = logits.masked_fill(banned, float("-inf")) | |
| nxt = int(lg.argmax()) | |
| gen.append(nxt) | |
| if nxt == vocab.eos: | |
| break | |
| logits = dec.step(nxt) | |
| t3 = time.time() | |
| text = vocab.decode(gen) | |
| print(f"decoder: prefix {prefix} -> {len(gen)} tokens in {t3 - t2:.1f}s") | |
| print(text) | |
| json.dump(dict(meta=meta, objects=[dict(box=model.to_1000(b['box'], meta), score=b['score']) for b in det['object']], | |
| text=[dict(box=model.to_1000(b['box'], meta), score=b['score'], ocr=o) for b, o in zip(det['text'], ocr)], | |
| prefix=prefix, generated=gen, decoded=text), | |
| open(os.path.join(out_dir, "result.json"), "w", encoding="utf-8"), ensure_ascii=False, indent=1) | |
| # ===================================================================================== commands: reader | |
| class RealCrops(torch.utils.data.IterableDataset): | |
| """Real word crops (parquet rows of image bytes and text) as reader strips with CTC labels; each worker | |
| reads its own share of the row groups. jitter: random box margins and rescaling blur.""" | |
| def __init__(self, files, index, jitter): | |
| super().__init__() | |
| self.files, self.index, self.jitter = files, index, jitter | |
| def __iter__(self): | |
| import pyarrow.parquet as pq | |
| info = torch.utils.data.get_worker_info() | |
| wid, nw = (info.id, info.num_workers) if info else (0, 1) | |
| rng = random.Random(4242 + wid) | |
| units = [(f, g) for f in self.files for g in range(pq.ParquetFile(f).num_row_groups)] | |
| mine = [u for i, u in enumerate(units) if i % nw == wid] | |
| while True: | |
| rng.shuffle(mine) | |
| for f, g in mine: | |
| rows = pq.ParquetFile(f).read_row_group(g).to_pylist() | |
| rng.shuffle(rows) | |
| for row in rows: | |
| text = row["text"] | |
| if not text or text == "###" or any(c not in self.index for c in text): | |
| continue | |
| try: | |
| img = Image.open(io.BytesIO(row["image"]["bytes"])).convert("RGB") | |
| except Exception: | |
| continue | |
| w, h = img.size | |
| if h < 6 or w < 4: | |
| continue | |
| if self.jitter: | |
| a = np.asarray(img, dtype=np.float32) | |
| fill = tuple(int(v) for v in np.concatenate([a[0], a[-1], a[:, 0], a[:, -1]]).mean(0)) | |
| side = [int(rng.uniform(-0.06, 0.3) * h) for _ in range(4)] | |
| canvas = Image.new("RGB", (w + max(0, side[0]) + max(0, side[2]), h + max(0, side[1]) + max(0, side[3])), fill) | |
| canvas.paste(img, (max(0, side[0]), max(0, side[1]))) | |
| x0, y0 = max(0, -side[0]), max(0, -side[1]) | |
| x1, y1 = canvas.width - max(0, -side[2]), canvas.height - max(0, -side[3]) | |
| if x1 - x0 >= 4 and y1 - y0 >= 6: | |
| img = canvas.crop((x0, y0, x1, y1)) | |
| if rng.random() < 0.5: | |
| f = rng.uniform(0.5, 1.0) | |
| img = img.resize((max(1, int(img.width * f)), max(1, int(img.height * f))), Image.BILINEAR) | |
| w, h = img.size | |
| U = max(1, int(round(48 * w / h))) | |
| if U > 320: | |
| continue | |
| labels = [self.index[c] for c in text] | |
| reps = sum(1 for a, b in zip(labels, labels[1:]) if a == b) | |
| if -(-U // 4) < len(labels) + reps: | |
| continue | |
| strip = np.full((48, 320, 3), 127, np.uint8) | |
| strip[:, :U] = np.asarray(img.resize((U, 48), Image.BILINEAR)) | |
| yield strip, U, labels, text | |
| class SyntheticLines(torch.utils.data.IterableDataset): | |
| """Synthetic text lines from Generator, seeded per worker.""" | |
| def __init__(self, charset, dead, words, fonts, style, tight): | |
| super().__init__() | |
| self.charset, self.dead, self.words, self.fonts, self.style, self.tight = charset, dead, words, fonts, style, tight | |
| def __iter__(self): | |
| info = torch.utils.data.get_worker_info() | |
| gen = Generator(self.charset, self.dead, self.words, self.fonts, 9000 + (info.id if info else 0) + random.randint(0, 10**6), | |
| style=self.style, tight=self.tight) | |
| while True: | |
| yield gen.sample() | |
| def collate_strips(batch): | |
| x = torch.from_numpy(np.stack([b[0] for b in batch])).permute(0, 3, 1, 2).float() / 127.5 - 1.0 | |
| lens = torch.tensor([-(-b[1] // 4) for b in batch]) | |
| tgt = torch.tensor([c for b in batch for c in b[2]]) | |
| tl = torch.tensor([len(b[2]) for b in batch]) | |
| return x, lens, tgt, tl | |
| def batches(ds, batch, workers): | |
| return iter(torch.utils.data.DataLoader(ds, batch_size=batch, num_workers=workers, collate_fn=collate_strips, | |
| prefetch_factor=4, persistent_workers=True)) | |
| def cmd_train(argv): | |
| """Quantisation-aware fine-tuning of the CTC reader on real word crops mixed with synthetic lines. | |
| The 2-bit matrices run through the stock quantiser in every forward pass (straight-through gradients), so the | |
| exported 2-bit reader is the trained one. Real crops get CTC plus distillation toward Mirage's 2-bit reader; | |
| synthetic lines carrying the untrained characters get CTC. Checkpoints of the reader's tensors and a log are | |
| written to --out; `checkpoint --replace` puts one into the full weights. | |
| Oasis's reader: --real 'data/ocr/textocr/train-*.parquet' --act-quant --jitter for 8000 steps, then | |
| --init on the result with --synth realistic for 4000, then again with --tight 0.5 for 3000.""" | |
| ap = parser("train", cmd_train.__doc__) | |
| ap.add_argument("--weights", help="Mirage's float32 weights (default: from --mirage)") | |
| ap.add_argument("--init", default=None, help="reader weights to start from instead of the original") | |
| ap.add_argument("--real", required=True, help="glob of parquet files of word crops (image bytes and text)") | |
| ap.add_argument("--evalsets", help="comma list of packed crop archives (default: data/ocr/real_words.npz and synth_val.npz)") | |
| ap.add_argument("--steps", type=int, default=8000) | |
| ap.add_argument("--real-batch", type=int, default=128) | |
| ap.add_argument("--syn-batch", type=int, default=64) | |
| ap.add_argument("--syn-weight", type=float, default=1.0) | |
| ap.add_argument("--real-weight", type=float, default=1.0, help="weight of CTC on real labels (0: distillation only)") | |
| ap.add_argument("--distill", type=float, default=0.5) | |
| ap.add_argument("--lr", type=float, default=1e-4) | |
| ap.add_argument("--act-quant", action="store_true", help="engine-style int8 activations in student and teacher") | |
| ap.add_argument("--jitter", action="store_true", help="random box margins and rescaling blur on real crops") | |
| ap.add_argument("--synth", choices=["uniform", "realistic"], default="uniform", help="synthetic text distribution") | |
| ap.add_argument("--fonts", default=",".join(LINUX_FONTS), help="comma list of font globs for synthetic lines") | |
| ap.add_argument("--tight", type=float, default=0.0, help="share of synthetic lines cut as tightly as a detector box") | |
| ap.add_argument("--freeze", default="", help="comma list of name prefixes to freeze, e.g. w0,dw_0,pw_0") | |
| ap.add_argument("--workers", type=int, default=12) | |
| ap.add_argument("--eval-every", type=int, default=1000) | |
| ap.add_argument("--out", required=True) | |
| args = ap.parse_args(argv) | |
| os.makedirs(args.out, exist_ok=True) | |
| dev = "cuda" if torch.cuda.is_available() else "cpu" | |
| torch.manual_seed(0) | |
| arc = Archive(mirage_archive(args)) | |
| charset = charset_of(arc) | |
| index = {c: i + 1 for i, c in enumerate(charset)} | |
| words = vocab_words(vocab_of(arc)) | |
| base = load_file(args.weights or mirage_weights(args)) | |
| dead_cls = dead_classes(base) | |
| dead_chars = [charset[c - 1] for c in dead_cls] | |
| start = dict(base) | |
| if args.init: | |
| start.update(load_file(args.init)) | |
| quantizer = Quantizer(arc.codebook, dev) | |
| model = Reader(start, quantizer, act_quant=args.act_quant).to(dev) | |
| teacher = Reader(base, quantizer, act_quant=args.act_quant).to(dev).eval() | |
| if not args.init: | |
| # Untrained classes start from the output weights of the letter they are closest to. | |
| K = model.params["reader__out__kernel"] | |
| Bv = model.params["reader__out__bias"] | |
| special = {"–": "-", "¡": "!", "ĸ": "k", "ʼn": "n", "Ŋ": "N", "ŋ": "n", "Ŧ": "T", "ŧ": "t", "Ð": "D", | |
| "Þ": "P", "þ": "p", "Ŀ": "L", "ŀ": "l", "IJ": "I", "ij": "i"} | |
| with torch.no_grad(): | |
| for c in dead_cls: | |
| ch = charset[c - 1] | |
| cands = ([special[ch]] if ch in special else []) + [v for v in charset if v != ch and v not in dead_chars and unicodedata.normalize("NFKD", v)[0] == unicodedata.normalize("NFKD", ch)[0]] + [unicodedata.normalize("NFKD", ch)[0]] | |
| d = next((charset.index(v) + 1 for v in cands if v in charset and v not in dead_chars), None) | |
| if d is not None: | |
| K[:, c] = K[:, d] | |
| Bv[c] = Bv[d] - 0.5 | |
| for q in teacher.parameters(): | |
| q.requires_grad_(False) | |
| frozen = [f for f in args.freeze.split(",") if f] | |
| trainable = [] | |
| for name in model.names: | |
| q = model.params[name.replace("/", "__")] | |
| short = name[len("reader/"):] | |
| keep = not any(short == f or short.startswith(f + "/") for f in frozen) | |
| q.requires_grad_(keep) | |
| if keep: | |
| trainable.append(q) | |
| print(f"training {sum(q.numel() for q in trainable):,} of {sum(q.numel() for q in model.parameters()):,} reader weights", flush=True) | |
| real_files = sorted(glob.glob(args.real)) | |
| print("real files:", real_files, flush=True) | |
| evalsets = args.evalsets.split(",") if args.evalsets else [os.path.join(args.data, "ocr", "real_words.npz"), | |
| os.path.join(args.data, "ocr", "synth_val.npz")] | |
| sets = [load_packed(p) for p in evalsets] | |
| dset = set(dead_chars) | |
| res = evaluate_reader(teacher, sets, charset, dset, dev) | |
| print("original:", json.dumps(res), flush=True) | |
| log = open(os.path.join(args.out, "log.jsonl"), "a") | |
| log.write(json.dumps(dict(step=0, eval=res, args=vars(args))) + "\n") | |
| real_it = batches(RealCrops(real_files, index, args.jitter), args.real_batch, max(1, args.workers * 2 // 3)) | |
| syn_it = batches(SyntheticLines(charset, dead_chars, words, args.fonts.split(","), args.synth, args.tight), | |
| args.syn_batch, max(1, args.workers // 3)) | |
| opt = torch.optim.AdamW(trainable, lr=args.lr, weight_decay=0.0) | |
| sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / 200) * 0.5 * (1 + math.cos(math.pi * min(s, args.steps) / args.steps))) | |
| ctc = torch.nn.CTCLoss(blank=0, zero_infinity=True) | |
| model.train() | |
| t0 = time.time() | |
| for step in range(1, args.steps + 1): | |
| xr, lr_, tr, tlr = next(real_it) | |
| xs, ls_, ts, tls = next(syn_it) | |
| xr, xs = xr.to(dev), xs.to(dev) | |
| lp_r = F.log_softmax(model(xr), -1) | |
| lp_s = F.log_softmax(model(xs), -1) | |
| l_real = ctc(lp_r.transpose(0, 1), tr.to(dev), lr_.to(dev), tlr.to(dev)) | |
| l_syn = ctc(lp_s.transpose(0, 1), ts.to(dev), ls_.to(dev), tls.to(dev)) | |
| with torch.no_grad(): | |
| tp = F.softmax(teacher(xr), -1) | |
| kd = F.kl_div(lp_r, tp, reduction="batchmean") / lp_r.shape[1] | |
| loss = args.real_weight * l_real + args.syn_weight * l_syn + args.distill * kd | |
| opt.zero_grad(set_to_none=True) | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(trainable, 5.0) | |
| opt.step() | |
| sched.step() | |
| if step % 100 == 0: | |
| print(f"step {step} real {l_real.item():.3f} syn {l_syn.item():.3f} kd {kd.item():.4f} lr {sched.get_last_lr()[0]:.2e} {time.time() - t0:.0f}s", flush=True) | |
| if step % args.eval_every == 0 or step == args.steps: | |
| res = evaluate_reader(model, sets, charset, dset, dev) | |
| print(f"eval step {step}: {json.dumps(res)}", flush=True) | |
| log.write(json.dumps(dict(step=step, eval=res)) + "\n") | |
| log.flush() | |
| save_file(model.state(), os.path.join(args.out, f"reader_{step:06d}.safetensors")) | |
| print("done", flush=True) | |
| def reader_name(r): | |
| return "base" if r == "base" else os.path.basename(os.path.dirname(r)) + "/" + os.path.basename(r).replace(".safetensors", "") | |
| def cmd_eval_reader(argv): | |
| """Batched reader evaluation on packed crops, with the engine's window and frame-retention rules. | |
| Reports, per source: word accuracy (lowercase, letters and digits), exact match, CER, and recall of the | |
| characters the original checkpoint never trained. --dead-offset adds a constant to those classes' biases, | |
| which trades their recall against false readings on ordinary text.""" | |
| ap = parser("eval-reader", cmd_eval_reader.__doc__) | |
| ap.add_argument("packed", nargs="+") | |
| ap.add_argument("--reader", action="append", default=[], help="reader weights to score in place of Mirage's; repeatable") | |
| ap.add_argument("--dead-offset", type=float, action="append", default=[]) | |
| ap.add_argument("--weights", help="Mirage's float32 weights (default: from --mirage)") | |
| ap.add_argument("--qat", action="store_true", help="run the 2-bit matrices through the training-time quantiser") | |
| ap.add_argument("--act-quant", action="store_true", help="int8 activations per rotated 128-group, as the engine computes them") | |
| args = ap.parse_args(argv) | |
| dev = "cuda" if torch.cuda.is_available() else "cpu" | |
| arc = Archive(mirage_archive(args)) | |
| charset = charset_of(arc) | |
| base = load_file(args.weights or mirage_weights(args)) | |
| dead_cls = dead_classes(base) | |
| dead_chars = set(charset[c - 1] for c in dead_cls) | |
| sets = [load_packed(f) for f in args.packed] | |
| for r in (args.reader or ["base"]): | |
| w = dict(base) | |
| if r != "base": | |
| w.update(load_file(r)) | |
| net = Reader(w, Quantizer(arc.codebook, dev) if args.qat else None, act_quant=args.act_quant).to(dev).eval() | |
| b0 = net.params["reader__out__bias"].detach().clone() | |
| for off in (args.dead_offset or [0.0]): | |
| with torch.no_grad(): | |
| net.params["reader__out__bias"].copy_(b0) | |
| net.params["reader__out__bias"][dead_cls] += off | |
| for X, owner, keeps, labels, sources in sets: | |
| stats = {} | |
| for i, path in enumerate(reader_paths(net, X, owner, keeps, len(labels), dev)): | |
| pred, gt = ctc_text(path, charset), str(labels[i]) | |
| st = stats.setdefault(str(sources[i]), dict(n=0, acc=0, exact=0, ed=0, len=0, dh=0, dt=0, fp=0)) | |
| st["n"] += 1 | |
| st["acc"] += word_key(pred) == word_key(gt) | |
| st["exact"] += pred == gt | |
| st["ed"] += edit_distance(pred, gt) | |
| st["len"] += len(gt) | |
| for ch in set(gt) & dead_chars: | |
| st["dt"] += gt.count(ch) | |
| st["dh"] += min(gt.count(ch), pred.count(ch)) | |
| st["fp"] += sum(1 for ch in pred if ch in dead_chars and ch not in gt) | |
| name = reader_name(r) | |
| for src, st in sorted(stats.items()): | |
| n = st["n"] | |
| print(f"{src:20s} {name:26s} off={off:+6.2f} n={n} word_acc={st['acc'] / n:.4f} exact={st['exact'] / n:.4f} " | |
| f"cer={st['ed'] / max(st['len'], 1):.4f} dead_recall={st['dh'] / max(st['dt'], 1):.3f} dead_false={st['fp']}", flush=True) | |
| def cmd_eval_detected(argv): | |
| """Score readers on packed detector crops (pack-detected): the crops of a word read and joined left to | |
| right, against the word.""" | |
| ap = parser("eval-detected", cmd_eval_detected.__doc__) | |
| ap.add_argument("packed", nargs="+") | |
| ap.add_argument("--reader", action="append", default=[], help="reader weights to score in place of Mirage's; repeatable") | |
| ap.add_argument("--qat", action="store_true", help="run the 2-bit matrices through the training-time quantiser") | |
| ap.add_argument("--act-quant", action="store_true", help="int8 activations per rotated 128-group, as the engine computes them") | |
| ap.add_argument("--dead-offset", type=float, action="append", default=[]) | |
| ap.add_argument("--weights", help="Mirage's float32 weights (default: from --mirage)") | |
| args = ap.parse_args(argv) | |
| dev = "cuda" if torch.cuda.is_available() else "cpu" | |
| arc = Archive(mirage_archive(args)) | |
| charset = charset_of(arc) | |
| base = load_file(args.weights or mirage_weights(args)) | |
| dead_cls = dead_classes(base) | |
| dead = set(charset[c - 1] for c in dead_cls) | |
| sets = [] | |
| for f in args.packed: | |
| d = dict(np.load(f, allow_pickle=False)) | |
| pix, offs = d["pix"], d["offsets"] | |
| wins, owner, keeps = [], [], [] | |
| for i in range(len(offs) - 1): | |
| ws, ks = reader_windows(pix[:, offs[i]:offs[i + 1]]) | |
| wins += ws | |
| owner += [i] * len(ws) | |
| keeps += ks | |
| X = torch.from_numpy(np.stack(wins)).permute(0, 3, 1, 2) if wins else torch.zeros(0, 3, 48, 320, dtype=torch.uint8) | |
| sets.append((os.path.basename(f), X, owner, keeps, d["item"], d["x0"], d["labels"], d["sources"])) | |
| for r in (args.reader or ["base"]): | |
| w = dict(base) | |
| if r != "base": | |
| w.update(load_file(r)) | |
| q = Quantizer(arc.codebook, dev) if args.qat else None | |
| net = Reader(w, q, act_quant=args.act_quant).to(dev).eval() | |
| b0 = net.params["reader__out__bias"].detach().clone() | |
| for off in (args.dead_offset or [0.0]): | |
| with torch.no_grad(): | |
| net.params["reader__out__bias"].copy_(b0) | |
| net.params["reader__out__bias"][dead_cls] += off | |
| for name, X, owner, keeps, item, x0, labels, sources in sets: | |
| crop_text = [ctc_text(path, charset) for path in reader_paths(net, X, owner, keeps, len(item), dev)] | |
| words = [[] for _ in range(len(labels))] | |
| for c, (wi, xx) in enumerate(zip(item, x0)): | |
| words[wi].append((xx, crop_text[c])) | |
| stats = {} | |
| for wi, gt in enumerate(labels): | |
| pred = " ".join(t for _, t in sorted(words[wi])).strip() | |
| st = stats.setdefault(str(sources[wi]), [0, 0, 0, 0, 0]) | |
| st[0] += 1 | |
| st[1] += word_key(pred) == word_key(str(gt)) | |
| st[2] += edit_distance(pred, str(gt)) | |
| st[3] += len(str(gt)) | |
| st[4] += sum(1 for ch in pred if ch in dead and ch not in str(gt)) | |
| nm = reader_name(r) | |
| for src, (n, a, e, ln, fp) in sorted(stats.items()): | |
| print(f"{src:28s} {nm:26s} off={off:+5.2f} n={n} word_acc={a / n:.4f} cer={e / max(ln, 1):.4f} dead_false={fp}", flush=True) | |
| def cmd_pack_words(argv): | |
| """Pack cropped-word OCR benchmarks (parquet files of image bytes and text) into a numpy archive, each crop | |
| resized to height 48 with its aspect kept: pix uint8 [48, total_width, 3], offsets int64 [N+1], labels and | |
| sources.""" | |
| ap = parser("pack-words", cmd_pack_words.__doc__) | |
| ap.add_argument("out") | |
| ap.add_argument("files", nargs="+") | |
| ap.add_argument("--limit", type=int, default=0, help="the first N rows of each file") | |
| args = ap.parse_args(argv) | |
| import pyarrow.parquet as pq | |
| strips, labels, sources = [], [], [] | |
| for f in args.files: | |
| table = pq.read_table(f) | |
| if args.limit: | |
| table = table.slice(0, args.limit) | |
| for row in table.to_pylist(): | |
| img = Image.open(io.BytesIO(row["image"]["bytes"])).convert("RGB") | |
| w, h = img.size | |
| U = max(1, int(round(48 * w / h))) | |
| strips.append(np.asarray(img.resize((U, 48), Image.BILINEAR), dtype=np.uint8)) | |
| labels.append(row["text"]) | |
| sources.append(os.path.basename(f)) | |
| offsets = np.cumsum([0] + [s.shape[1] for s in strips]).astype(np.int64) | |
| np.savez_compressed(args.out, pix=np.concatenate(strips, 1), offsets=offsets, labels=np.array(labels), sources=np.array(sources)) | |
| print(len(labels), "crops, total width", int(offsets[-1])) | |
| def cmd_pack_detected(argv): | |
| """Pack the crops the Oasis engine's text detector cuts from word images placed in scenes as bench-ocr | |
| places them: each text box is cropped from the scene with the engine's rounding and resized to 48 px height, | |
| so that readers can be scored on exactly the crops the engine reads. Arrays: pix [48, total, 3], offsets, item | |
| (crop -> word index), x0 (left edge, for ordering), labels (per word) and sources (per word).""" | |
| ap = parser("pack-detected", cmd_pack_detected.__doc__) | |
| ap.add_argument("out") | |
| ap.add_argument("files", nargs="+") | |
| ap.add_argument("--oasis-lib", help="the Oasis engine library (default: found as oasis.py finds it)") | |
| ap.add_argument("--oasis-model", required=True) | |
| ap.add_argument("--limit", type=int, default=0) | |
| ap.add_argument("--text-conf", type=float, default=0.2) | |
| args = ap.parse_args(argv) | |
| import pyarrow.parquet as pq | |
| eng = OasisEngine(args.oasis_lib, args.oasis_model) | |
| opts = {"sections": ["text"], "text_conf": args.text_conf, "max_text_lines": 4, "max_objects": 1, "max_tokens": 1, "mask_format": "largest"} | |
| strips, item, xs, labels, sources = [], [], [], [], [] | |
| for f in args.files: | |
| t = pq.read_table(f) | |
| if args.limit: | |
| t = t.slice(0, args.limit) | |
| for r in t.to_pylist(): | |
| scene = word_scene(Image.open(io.BytesIO(r["image"]["bytes"])).convert("RGB")) | |
| rec, _, _ = eng.describe(np.asarray(scene, dtype=np.uint8).tobytes(), 640, 480, opts) | |
| wi = len(labels) | |
| labels.append(r["text"]) | |
| sources.append(os.path.basename(f)) | |
| for tbox in rec.get("text", []) if rec else []: | |
| b = tbox["bbox"] | |
| x0, y0, x1, y1 = b[0] / 1000 * 640, b[1] / 1000 * 480, b[2] / 1000 * 640, b[3] / 1000 * 480 | |
| crop = scene.crop((math.floor(x0), math.floor(y0), math.ceil(x1), math.ceil(y1))) | |
| if crop.width < 1 or crop.height < 1: | |
| continue | |
| U = max(1, int(round(48 * crop.width / crop.height))) | |
| strips.append(np.asarray(crop.resize((U, 48), Image.BILINEAR), dtype=np.uint8)) | |
| item.append(wi) | |
| xs.append(b[0]) | |
| offsets = np.cumsum([0] + [s.shape[1] for s in strips]).astype(np.int64) | |
| np.savez_compressed(args.out, pix=np.concatenate(strips, 1) if strips else np.zeros((48, 0, 3), np.uint8), offsets=offsets, | |
| item=np.array(item, np.int64), x0=np.array(xs, np.float32), labels=np.array(labels), sources=np.array(sources)) | |
| print(f"{len(labels)} words, {len(strips)} detected crops") | |
| def cmd_synth_val(argv): | |
| """Write a fixed synthetic validation set in the packed crop format, with the characters Mirage's reader | |
| never learned listed under "dead".""" | |
| ap = parser("synth-val", cmd_synth_val.__doc__) | |
| ap.add_argument("out") | |
| ap.add_argument("--n", type=int, default=2000) | |
| ap.add_argument("--seed", type=int, default=12345) | |
| ap.add_argument("--style", choices=["uniform", "realistic"], default="uniform") | |
| ap.add_argument("--tight", type=float, default=0.0) | |
| ap.add_argument("--fonts", default=",".join(LINUX_FONTS), help="comma list of font globs") | |
| args = ap.parse_args(argv) | |
| arc = Archive(mirage_archive(args)) | |
| charset = charset_of(arc) | |
| words = vocab_words(vocab_of(arc)) | |
| dead = [charset[c - 1] for c in dead_classes(load_file(mirage_weights(args)))] | |
| gen = Generator(charset, dead, words, args.fonts.split(","), args.seed, style=args.style, tight=args.tight) | |
| strips, labels = [], [] | |
| for _ in range(args.n): | |
| strip, U, _, s = gen.sample() | |
| strips.append(strip[:, :U]) | |
| labels.append(s) | |
| offsets = np.cumsum([0] + [s.shape[1] for s in strips]).astype(np.int64) | |
| np.savez_compressed(args.out, pix=np.concatenate(strips, 1), offsets=offsets, labels=np.array(labels), | |
| sources=np.array(["synthetic"] * len(labels)), dead=np.array(dead)) | |
| print(len(labels), "synthetic crops") | |
| def cmd_det_eval(argv): | |
| """Class-agnostic detection quality of the reference model against COCO val2017 ground truth, for several | |
| weight variants: AP at IoU 0.5 over all proposals (score-ranked, greedy matching), and recall and precision at | |
| a score threshold. Crowd annotations are ignored. A variant is name=path.safetensors, whose tensors replace | |
| Mirage's float32 ones, or name=stock, Mirage's weights through the stock quantiser.""" | |
| ap = parser("det-eval", cmd_det_eval.__doc__) | |
| ap.add_argument("--n", type=int, default=100) | |
| ap.add_argument("--first", type=int, default=0) | |
| ap.add_argument("--variant", action="append", default=[]) | |
| ap.add_argument("--thr", type=float, default=0.4) | |
| args = ap.parse_args(argv) | |
| model = Mirage(mirage_weights(args)) | |
| p32 = {k: v.clone() for k, v in model.p.items()} | |
| arc = Archive(mirage_archive(args)) | |
| plan = cq_plan() | |
| gt = coco_truth(args) | |
| photos = val300(args, args.first, args.n) | |
| variants = [("fp32", None)] + [tuple(v.split("=", 1)) for v in args.variant] | |
| for name, src in variants: | |
| if src is None: | |
| model.p = {k: v.clone() for k, v in p32.items()} | |
| elif src == "stock": | |
| model.p, _ = variant(p32, plan, arc.codebook, "stock") | |
| else: | |
| model.p = dict(p32) | |
| model.p.update({k: v.float() for k, v in load_file(src).items()}) | |
| model.T = model.p["embedding/embedding"] | |
| dets, n_gt = [], 0 | |
| hits_thr = props_thr = matched_thr = 0 | |
| t0 = time.time() | |
| for iid, path in photos: | |
| img = Image.open(path).convert("RGB") | |
| enc = model.encode(img) | |
| det = model.detect(enc, conf=0.05, text_conf=1.1, nms_iou=0.5) | |
| meta = enc["meta"] | |
| g = gt.get(int(iid), []) | |
| n_gt += len(g) | |
| boxes = [] | |
| for b in det["object"]: | |
| bb = model.to_1000(b["box"], meta) | |
| boxes.append(([bb[0] / 1000 * meta["W"], bb[1] / 1000 * meta["H"], bb[2] / 1000 * meta["W"], bb[3] / 1000 * meta["H"]], b["score"])) | |
| used = set() | |
| for box, s in sorted(boxes, key=lambda x: -x[1]): | |
| best, bi = 0.0, -1 | |
| for k, gb in enumerate(g): | |
| if k in used: | |
| continue | |
| v = iou(box, gb) | |
| if v > best: | |
| best, bi = v, k | |
| tp = best >= 0.5 | |
| if tp: | |
| used.add(bi) | |
| dets.append((s, tp)) | |
| sel = [bx for bx, s in boxes if s >= args.thr] | |
| props_thr += len(sel) | |
| matched_thr += sum(any(iou(bx, gb) >= 0.5 for gb in g) for bx in sel) | |
| hits_thr += sum(any(iou(bx, gb) >= 0.5 for bx in sel) for gb in g) | |
| dets.sort(key=lambda x: -x[0]) | |
| tp = np.cumsum([d[1] for d in dets]) | |
| fp = np.cumsum([not d[1] for d in dets]) | |
| rec = tp / max(n_gt, 1) | |
| prec = tp / np.maximum(tp + fp, 1) | |
| ap50 = 0.0 | |
| for t in np.linspace(0, 1, 101): | |
| p = prec[rec >= t] | |
| ap50 += (p.max() if p.size else 0.0) / 101 | |
| print(f"{name:10s} images {len(photos)} gt {n_gt} AP50 {ap50:.4f} max_recall {rec[-1] if len(rec) else 0:.4f} " | |
| f"| at score>={args.thr}: recall {hits_thr / max(n_gt, 1):.4f} precision {matched_thr / max(props_thr, 1):.4f} " | |
| f"proposals/img {props_thr / len(photos):.1f} | {time.time() - t0:.0f}s", flush=True) | |
| # ===================================================================================== entry point | |
| COMMANDS = { | |
| "export": (cmd_export, "write an .oasis weight file for the engine"), | |
| "checkpoint": (cmd_checkpoint, "assemble a float32 checkpoint from Mirage's weights and replacements"), | |
| "run": (cmd_run, "the reference forward pass on one image"), | |
| "det-eval": (cmd_det_eval, "detection quality of weight variants on the reference model"), | |
| "train": (cmd_train, "quantisation-aware training of the text reader"), | |
| "eval-reader": (cmd_eval_reader, "score readers on packed word crops"), | |
| "eval-detected": (cmd_eval_detected, "score readers on the engine's detector crops"), | |
| "pack-words": (cmd_pack_words, "pack word-crop benchmarks into an archive"), | |
| "pack-detected": (cmd_pack_detected, "pack the crops the engine's text detector cuts"), | |
| "synth-val": (cmd_synth_val, "write a synthetic validation set"), | |
| } | |
| def main(argv=None): | |
| argv = sys.argv[1:] if argv is None else argv | |
| if not argv or argv[0] not in COMMANDS: | |
| print(__doc__.rstrip()) | |
| print() | |
| for name, (_, what) in COMMANDS.items(): | |
| print(f" {name:14s} {what}") | |
| sys.exit(0 if not argv or argv[0] in ("-h", "--help") else 2) | |
| COMMANDS[argv[0]][0](argv[1:]) | |
| if __name__ == "__main__": | |
| main() | |