File size: 6,228 Bytes
b5731a6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """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
# ---- encoder
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
# ---- decoder
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)
|