"""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("", 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"" if k < N_OBJECT_REFS else f"" 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)} @dataclass 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("> (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) @staticmethod 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 @torch.inference_mode() 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 @torch.inference_mode() 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"]) @torch.inference_mode() 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 @torch.inference_mode() 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) @torch.inference_mode() 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) @torch.inference_mode() 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: # trace # S key=value; key=value # E name @position * {key=value; key=value} # T * # R predicate # 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"(.*?)", text, flags=re.S) if m: trace = m.group(1).strip() body = text[m.end():] body = body.replace("", "") 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 @torch.inference_mode() 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 @torch.no_grad() 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 @torch.no_grad() 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(" 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()