| """MLX port of the NCM multiscale-residual VQ-VAE codec (XSum/WMT14 variant). Encode + decode only.""" |
| import json, math |
| import numpy as np |
| import mlx.core as mx |
|
|
| GN_EPS, LN_EPS = 1e-5, 1e-6 |
|
|
|
|
| def area_matrix(src, dst, dt=np.float32): |
| m = np.zeros((dst, src), dt) |
| for i in range(dst): |
| s, e = (i * src) // dst, math.ceil((i + 1) * src / dst) |
| m[i, s:e] = 1.0 / (e - s) |
| return m |
|
|
|
|
| def linear_matrix(src, dst, dt=np.float32): |
| m = np.zeros((dst, src), dt) |
| if src == 1: |
| m[:, 0] = 1 |
| return m |
| sc = src / dst |
| for j in range(dst): |
| x = max((j + .5) * sc - .5, 0.0) |
| i0 = min(math.floor(x), src - 1); i1 = min(i0 + 1, src - 1); w = x - i0 |
| m[j, i0] += 1 - w; m[j, i1] += w |
| return m |
|
|
|
|
| def plan_strides(total, base=4): |
| S, out = int(total), [] |
| while S > 1: |
| f = min(base, S) |
| while f > 1 and S % f: f -= 1 |
| if f <= 1: out.append(S); S = 1 |
| else: out.append(f); S //= f |
| return out or [1] |
|
|
|
|
| def phi_index(level, K): |
| t = np.linspace(1 / 3 / K, 1 - 1 / 3 / K, K) if K == 4 else np.linspace(1 / 2 / K, 1 - 1 / 2 / K, K) |
| return int(np.argmin(np.abs(t - level / max(K - 1, 1)))) |
|
|
|
|
| def cascade_ops(total, base, inter_norm, final_norm): |
| """Op list for level-0 down/up. ('d'|'u', stride, seq_idx) / ('n',). seq_idx = index inside upstream nn.Sequential.""" |
| st = plan_strides(total, base) |
| down, i = [], 0 |
| for k, s in enumerate(st): |
| if inter_norm and k > 0: down.append(("n",)); i += 1 |
| down.append(("d", s, i)); i += 1 |
| if final_norm: down.append(("n",)) |
| up, i = [], 0 |
| for k, s in enumerate(reversed(st)): |
| if inter_norm and k > 0: up.append(("n",)); i += 1 |
| up.append(("u", s, i)); i += 1 |
| return down, up |
|
|
|
|
| def group_norm(x, G): |
| """x (B,L,C): torch GroupNorm(G, C, affine=False) on (B,C,L) -> stats over (L, C/G) per group.""" |
| B, L, C = x.shape |
| g = x.reshape(B, L, G, C // G) |
| mu = g.mean(axis=(1, 3), keepdims=True) |
| var = ((g - mu) ** 2).mean(axis=(1, 3), keepdims=True) |
| return ((g - mu) * mx.rsqrt(var + GN_EPS)).reshape(B, L, C) |
|
|
|
|
| def layer_norm(x, w, b): |
| mu = x.mean(-1, keepdims=True); v = ((x - mu) ** 2).mean(-1, keepdims=True) |
| return (x - mu) * mx.rsqrt(v + LN_EPS) * w + b |
|
|
|
|
| def rms(x, w): |
| return x * mx.rsqrt((x * x).mean(-1, keepdims=True) + LN_EPS) * w |
|
|
|
|
| class Codec: |
| def __init__(self, W: dict, cfg: dict, dtype=mx.float32): |
| self.cfg, self.dt = cfg, dtype |
| self.W = {k: v.astype(dtype) for k, v in W.items()} |
| self.levels = cfg["levels"]; self.K = len(self.levels); self.L = self.levels[-1] |
| self.G = cfg["groupnorm"]; self.H = cfg["n_head"] |
| self.ratio = cfg["phi_ratio"] |
| self.down_ops, self.up_ops = cascade_ops(self.L // self.levels[0], cfg["cascade_base"], cfg["cascade_norm"], cfg["cascade_norm_final"]) |
| nd = np.float64 if dtype == mx.float64 else np.float32 |
| self.area = [mx.array(area_matrix(self.L, n, nd), dtype=dtype) for n in self.levels] |
| self.lin = [mx.array(linear_matrix(n, self.L, nd), dtype=dtype) for n in self.levels] |
| self.n_extra = cfg.get("dec_layers", 1) - 1 |
|
|
| |
| def embed(self, tok): |
| z = self.W["wte"][tok] + self.W["wpe"][: tok.shape[1]] |
| return group_norm(z, self.G) |
|
|
| def _down0(self, x): |
| for op in self.down_ops: |
| if op[0] == "n": x = group_norm(x, self.G); continue |
| s, i = op[1], op[2] |
| B, L, C = x.shape |
| x = x.reshape(B, L // s, s * C) @ self.W[f"down.{i}.w"] + self.W[f"down.{i}.b"] |
| return x |
|
|
| def _up0(self, x): |
| for op in self.up_ops: |
| if op[0] == "n": x = group_norm(x, self.G); continue |
| s, i = op[1], op[2] |
| B, L, _ = x.shape |
| x = (x @ self.W[f"up.{i}.w"]).reshape(B, L * s, -1) + self.W[f"up.{i}.b"] |
| return x |
|
|
| def contrib(self, codes, k): |
| q = self.W[f"cb.{k}"][codes] |
| if k == 0: q = self._up0(q) |
| elif k != self.K - 1: q = self.lin[k] @ q |
| j = phi_index(k, self.K) |
| c = mx.conv1d(q, self.W[f"phi.{j}.w"], padding=1) + self.W[f"phi.{j}.b"] |
| return q * (1 - self.ratio) + c * self.ratio |
|
|
| def encode(self, tok): |
| rest = self.embed(tok); codes = [] |
| for k in range(self.K): |
| cb = self.W[f"cb.{k}"] |
| if k == 0: x = self._down0(rest) |
| elif k != self.K - 1: x = self.area[k] @ rest |
| else: x = rest |
| d = (x * x).sum(-1, keepdims=True) + (cb * cb).sum(-1) - 2 * (x @ cb.T) |
| idx = mx.argmin(d, axis=-1) |
| codes.append(idx) |
| rest = rest - self.contrib(idx, k) |
| return codes |
|
|
| def cumulative(self, codes): |
| out, z = [], None |
| for k, c in enumerate(codes): |
| q = self.contrib(c, k) |
| z = q if z is None else z + q |
| out.append(z) |
| return out |
|
|
| |
| def _block(self, x, p): |
| W, H = self.W, self.H |
| B, T, C = x.shape; hd = C // H |
| h = layer_norm(x, W[f"{p}.ln1.w"], W[f"{p}.ln1.b"]) |
| qkv = h @ W[f"{p}.attn.w"].T + W[f"{p}.attn.b"] |
| q, k, v = [t.reshape(B, T, H, hd).transpose(0, 2, 1, 3) for t in mx.split(qkv, 3, axis=-1)] |
| q, k = rms(q, W[f"{p}.qn"]), rms(k, W[f"{p}.kn"]) |
| y = mx.softmax((q @ k.transpose(0, 1, 3, 2)) * hd ** -0.5, axis=-1) @ v |
| y = y.transpose(0, 2, 1, 3).reshape(B, T, C) @ W[f"{p}.proj.w"].T + W[f"{p}.proj.b"] |
| x = x + y |
| h = layer_norm(x, W[f"{p}.ln2.w"], W[f"{p}.ln2.b"]) |
| h = h @ W[f"{p}.fc.w"].T + W[f"{p}.fc.b"] |
| h = 0.5 * h * (1 + mx.erf(h / math.sqrt(2))) |
| return x + h @ W[f"{p}.fcp.w"].T + W[f"{p}.fcp.b"] |
|
|
| def decode_logits(self, z): |
| x = self._block(z, "blk0") |
| for i in range(self.n_extra): x = self._block(x, f"blk{i+1}") |
| x = layer_norm(x, self.W["lnf.w"], self.W["lnf.b"]) |
| return x @ self.W["head"].T |
|
|
| def decode_codes(self, codes): |
| return mx.argmax(self.decode_logits(self.cumulative(codes)[-1]), -1) |
|
|
|
|
| def load(d, dtype=mx.float32): |
| cfg = json.load(open(f"{d}/spec.json")) |
| return Codec(mx.load(f"{d}/codec.safetensors"), cfg, dtype) |
|
|
