| import json |
| import math |
| import random |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| torch.set_grad_enabled(False) |
|
|
| SEED = 3410 |
| MAXLEN = 800 |
| AA = 'ACDEFGHIKLMNPQRSTVWY' |
| CFG = {'dim': 128, 'expand': 2, 'kernel': 9, 'dilations': [1, 2, 4, 8, 1, 2], 'emb_dim': 128} |
|
|
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
|
|
| def seed_all(seed=SEED): |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
|
|
|
|
| LUT = np.full(256, 21, dtype=np.int64) |
| for i, c in enumerate(AA): |
| LUT[ord(c)] = i + 1 |
|
|
|
|
| NOISE_P = 0.0 |
| MINKEEP = 1.0 |
|
|
|
|
| def encode(seq, train=False): |
| codes = LUT[np.frombuffer(seq.encode('ascii', 'ignore'), dtype=np.uint8)] |
| n = len(codes) |
| if train: |
| keep = min(MAXLEN, int(n * np.random.uniform(MINKEEP, 1.0))) |
| start = np.random.randint(0, n - keep + 1) |
| codes = codes[start:start + keep] |
| codes = np.where(np.random.rand(len(codes)) < NOISE_P, 21, codes) |
| elif n > MAXLEN: |
| start = (n - MAXLEN) // 2 |
| codes = codes[start:start + MAXLEN] |
| return codes |
|
|
|
|
| def make_batch(seqs, train, dev): |
| enc = [encode(s, train) for s in seqs] |
| x = np.zeros((len(enc), max(len(e) for e in enc)), dtype=np.int64) |
| for i, e in enumerate(enc): |
| x[i, :len(e)] = e |
| return torch.from_numpy(x).to(dev) |
|
|
|
|
| def make_batches(lens, bs, shuffle): |
| n = len(lens) |
| idx = np.random.permutation(n) if shuffle else np.arange(n) |
| batches = [] |
| for i in range(0, n, bs * 50): |
| ch = idx[i:i + bs * 50] |
| ch = ch[np.argsort([lens[j] for j in ch], kind='stable')] |
| batches += [ch[j:j + bs] for j in range(0, len(ch), bs)] |
| if shuffle: |
| random.shuffle(batches) |
| return batches |
|
|
|
|
| def acc(G, k, g): |
| G[k] = G[k] + g if k in G else g |
|
|
|
|
| def linear_fwd(x, W, b): |
| return x @ W + b, x |
|
|
|
|
| def linear_bwd(dy, W, name, x, G): |
| acc(G, name + '.w', x.reshape(-1, x.shape[-1]).T @ dy.reshape(-1, dy.shape[-1])) |
| acc(G, name + '.b', dy.reshape(-1, dy.shape[-1]).sum(0)) |
| return dy @ W.T |
|
|
|
|
| def ln_fwd(x, g, b, eps=1e-6): |
| mu = x.mean(-1, keepdim=True) |
| var = x.var(-1, unbiased=False, keepdim=True) |
| rstd = 1.0 / torch.sqrt(var + eps) |
| xhat = (x - mu) * rstd |
| return g * xhat + b, (xhat, rstd, g) |
|
|
|
|
| def ln_bwd(dy, name, cache, G): |
| xhat, rstd, g = cache |
| red = tuple(range(dy.dim() - 1)) |
| acc(G, name + '.g', (dy * xhat).sum(red)) |
| acc(G, name + '.b', dy.sum(red)) |
| dxhat = dy * g |
| return rstd * (dxhat - dxhat.mean(-1, keepdim=True) - xhat * (dxhat * xhat).mean(-1, keepdim=True)) |
|
|
|
|
| def gelu_fwd(x): |
| return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0))), x |
|
|
|
|
| def gelu_bwd(dy, x): |
| cdf = 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0))) |
| pdf = torch.exp(-0.5 * x * x) / math.sqrt(2.0 * math.pi) |
| return dy * (cdf + x * pdf) |
|
|
|
|
| def dwconv_fwd(x, w, b, d): |
| xc = x.transpose(1, 2).contiguous() |
| y = F.conv1d(xc, w, b, padding=d * (w.shape[-1] // 2), dilation=d, groups=xc.shape[1]) |
| return y.transpose(1, 2).contiguous(), (xc, d) |
|
|
|
|
| def dwconv_bwd(dy, w, name, cache, G): |
| xc, d = cache |
| k, L = w.shape[-1], xc.shape[2] |
| p = d * (k // 2) |
| dyc = dy.transpose(1, 2).contiguous() |
| dx = F.conv1d(dyc, w.flip(2), padding=p, dilation=d, groups=xc.shape[1]) |
| xpad = F.pad(xc, (p, p)) |
| dw = torch.empty_like(w) |
| for u in range(k): |
| dw[:, 0, u] = torch.einsum('ncl,ncl->c', dyc, xpad[:, :, u * d:u * d + L]) |
| acc(G, name + '.w', dw) |
| acc(G, name + '.b', dyc.sum((0, 2))) |
| return dx.transpose(1, 2).contiguous() |
|
|
|
|
| def pool_fwd(h, m): |
| n = m.sum(1).clamp(min=1.0) |
| mean = h.sum(1) / n |
| mx, idx = (h - (1.0 - m) * 1e4).max(1) |
| return torch.cat([mx, mean], 1), (idx, m, n) |
|
|
|
|
| def pool_bwd(dpool, cache): |
| idx, m, n = cache |
| C = dpool.shape[1] // 2 |
| dh = (dpool[:, C:] / n).unsqueeze(1) * m |
| dh.scatter_add_(1, idx.unsqueeze(1), dpool[:, :C].unsqueeze(1)) |
| return dh |
|
|
|
|
| def block_fwd(x, m, P, pre, d): |
| c = {} |
| h, c['dw'] = dwconv_fwd(x, P[pre + 'dw.w'], P[pre + 'dw.b'], d) |
| h = h * m |
| h, c['ln'] = ln_fwd(h, P[pre + 'ln.g'], P[pre + 'ln.b']) |
| h, c['pw1'] = linear_fwd(h, P[pre + 'pw1.w'], P[pre + 'pw1.b']) |
| h, c['act'] = gelu_fwd(h) |
| h, c['pw2'] = linear_fwd(h, P[pre + 'pw2.w'], P[pre + 'pw2.b']) |
| return (x + h) * m, c |
|
|
|
|
| def block_bwd(dout, m, P, pre, c, G): |
| dout = dout * m |
| dh = linear_bwd(dout, P[pre + 'pw2.w'], pre + 'pw2', c['pw2'], G) |
| dh = gelu_bwd(dh, c['act']) |
| dh = linear_bwd(dh, P[pre + 'pw1.w'], pre + 'pw1', c['pw1'], G) |
| dh = ln_bwd(dh, pre + 'ln', c['ln'], G) * m |
| dh = dwconv_bwd(dh, P[pre + 'dw.w'], pre + 'dw', c['dw'], G) |
| return dout + dh |
|
|
|
|
| def net_fwd(tokens, P, cfg, drop_p=0.0): |
| m = (tokens > 0).unsqueeze(-1).to(P['embed.w'].dtype) |
| x = P['embed.w'][tokens] * m |
| C = {'tokens': tokens, 'm': m, 'blocks': []} |
| for i, d in enumerate(cfg['dilations']): |
| x, c = block_fwd(x, m, P, f'b{i}.', d) |
| C['blocks'].append(c) |
| pooled, C['pool'] = pool_fwd(x, m) |
| z, C['ln'] = ln_fwd(pooled, P['head.ln.g'], P['head.ln.b']) |
| e, C['proj'] = linear_fwd(z, P['head.proj.w'], P['head.proj.b']) |
| a, C['act'] = gelu_fwd(e) |
| C['drop'] = None |
| if drop_p > 0: |
| C['drop'] = (torch.rand_like(a) > drop_p).to(a.dtype) / (1.0 - drop_p) |
| a = a * C['drop'] |
| logits, C['out'] = linear_fwd(a, P['head.out.w'], P['head.out.b']) |
| return logits, e, C |
|
|
|
|
| def net_bwd(dlogits, P, cfg, C): |
| G = {} |
| da = linear_bwd(dlogits, P['head.out.w'], 'head.out', C['out'], G) |
| if C['drop'] is not None: |
| da = da * C['drop'] |
| de = gelu_bwd(da, C['act']) |
| dz = linear_bwd(de, P['head.proj.w'], 'head.proj', C['proj'], G) |
| dx = pool_bwd(ln_bwd(dz, 'head.ln', C['ln'], G), C['pool']) |
| for i in reversed(range(len(cfg['dilations']))): |
| dx = block_bwd(dx, C['m'], P, f'b{i}.', C['blocks'][i], G) |
| dx = dx * C['m'] |
| dE = torch.zeros_like(P['embed.w']) |
| dE.index_add_(0, C['tokens'].reshape(-1), dx.reshape(-1, dx.shape[-1])) |
| G['embed.w'] = dE |
| return G |
|
|
|
|
| def init_params(cfg, n_classes, dev): |
| C, E, K, D = cfg['dim'], cfg['expand'], cfg['kernel'], cfg['emb_dim'] |
| P = {'embed.w': torch.randn(22, C, device=dev)} |
| for i in range(len(cfg['dilations'])): |
| pre = f'b{i}.' |
| P[pre + 'dw.w'] = torch.randn(C, 1, K, device=dev) * (1.0 / K) ** 0.5 |
| P[pre + 'dw.b'] = torch.zeros(C, device=dev) |
| P[pre + 'ln.g'] = torch.ones(C, device=dev) |
| P[pre + 'ln.b'] = torch.zeros(C, device=dev) |
| P[pre + 'pw1.w'] = torch.randn(C, E * C, device=dev) * (2.0 / C) ** 0.5 |
| P[pre + 'pw1.b'] = torch.zeros(E * C, device=dev) |
| P[pre + 'pw2.w'] = torch.randn(E * C, C, device=dev) * 0.5 * (1.0 / (E * C)) ** 0.5 |
| P[pre + 'pw2.b'] = torch.zeros(C, device=dev) |
| P['head.ln.g'] = torch.ones(2 * C, device=dev) |
| P['head.ln.b'] = torch.zeros(2 * C, device=dev) |
| P['head.proj.w'] = torch.randn(2 * C, D, device=dev) * (1.0 / (2 * C)) ** 0.5 |
| P['head.proj.b'] = torch.zeros(D, device=dev) |
| P['head.out.w'] = torch.randn(D, n_classes, device=dev) * 0.02 |
| P['head.out.b'] = torch.zeros(n_classes, device=dev) |
| return P |
|
|
|
|
| def softmax(logits): |
| e = torch.exp(logits - logits.max(1, keepdim=True).values) |
| return e / e.sum(1, keepdim=True) |
|
|
|
|
| def cross_entropy_soft_loss_and_grad(logits, targets): |
| probs = softmax(logits) |
| loss = -(targets * torch.log(probs + 1e-12)).sum(1).mean() |
| return loss, (probs - targets) / logits.shape[0] |
|
|
|
|
| def adam_init(P): |
| return {'m': {k: torch.zeros_like(v) for k, v in P.items()}, |
| 'v': {k: torch.zeros_like(v) for k, v in P.items()}, 't': 0} |
|
|
|
|
| def adam_step(P, G, S, lr, clip, wd, b1=0.9, b2=0.999, eps=1e-8): |
| norm = torch.sqrt(sum((g * g).sum() for g in G.values())).item() |
| scale = min(1.0, clip / (norm + 1e-12)) |
| S['t'] += 1 |
| for k in P: |
| g = G[k] * scale |
| S['m'][k].mul_(b1).add_((1 - b1) * g) |
| S['v'][k].mul_(b2).add_((1 - b2) * g * g) |
| mh = S['m'][k] / (1 - b1 ** S['t']) |
| vh = S['v'][k] / (1 - b2 ** S['t']) |
| if k.endswith('.w') and k != 'embed.w': |
| P[k].mul_(1 - lr * wd) |
| P[k].sub_(lr * mh / (torch.sqrt(vh) + eps)) |
| return norm |
|
|
|
|
| def lr_at(step, total, base): |
| warm = max(1, int(0.05 * total)) |
| if step < warm: |
| return base * (step + 1) / warm |
| t = (step - warm) / max(1, total - warm) |
| return base * 0.5 * (1 + math.cos(math.pi * t)) |
|
|
|
|
| def run(P, cfg, seqs, dev, bs=64, want_probs=True): |
| order = np.argsort([len(s) for s in seqs], kind='stable') |
| emb = np.zeros((len(seqs), cfg['emb_dim']), np.float32) |
| probs = np.zeros((len(seqs), P['head.out.b'].shape[0]), np.float32) if want_probs else None |
| for s in range(0, len(seqs), bs): |
| ids = order[s:s + bs] |
| logits, e, _ = net_fwd(make_batch([seqs[i] for i in ids], False, dev), P, cfg) |
| emb[ids] = e.cpu().numpy() |
| if want_probs: |
| probs[ids] = softmax(logits).cpu().numpy() |
| return emb, probs |
|
|
|
|
| def make_bank(emb, labels, ec_list): |
| return {'emb': F.normalize(torch.from_numpy(emb.astype(np.float32)), dim=1), |
| 'labels': labels, 'ec_list': ec_list} |
|
|
|
|
| def knn_votes(bank, emb, topk=5, temp=0.05): |
| q = F.normalize(torch.from_numpy(emb.astype(np.float32)), dim=1) |
| votes = np.zeros((len(q), len(bank['ec_list'])), np.float32) |
| for s in range(0, len(q), 256): |
| v, idx = (q[s:s + 256] @ bank['emb'].T).topk(topk, dim=1) |
| w = softmax(v / temp).numpy() |
| idx = idx.numpy() |
| for r in range(len(idx)): |
| for j in range(topk): |
| labs = bank['labels'][idx[r, j]] |
| votes[s + r, labs] += w[r, j] / len(labs) |
| return votes |
|
|
|
|
| def pick(probs, votes, ec_list, alpha): |
| score = alpha * votes + (1 - alpha) * probs |
| return [ec_list[i] for i in score.argmax(axis=1)] |
|
|
|
|
| def level_scores(pred_ecs, truths): |
| res = {} |
| for level in (4, 3, 2, 1): |
| total = 0.0 |
| for p, recorded in zip(pred_ecs, truths): |
| key = p.split('.')[:level] |
| total += sum(e.split('.')[:level] == key for e in recorded) / len(recorded) |
| res[level] = total / len(pred_ecs) |
| return res |
|
|
|
|
| def save(P, cfg, bank_emb, labels, ec_list, folder, alpha): |
| folder = Path(folder) |
| folder.mkdir(parents=True, exist_ok=True) |
| torch.save({k: v.cpu() for k, v in P.items()}, folder / 'model.pt') |
| np.save(folder / 'bank.npy', bank_emb.astype(np.float16)) |
| (folder / 'meta.json').write_text(json.dumps({'ec_list': ec_list, 'labels': labels, 'cfg': cfg, 'alpha': alpha})) |
|
|
|
|
| def load(folder): |
| folder = Path(folder) |
| meta = json.loads((folder / 'meta.json').read_text()) |
| P = torch.load(folder / 'model.pt', map_location='cpu') |
| bank = make_bank(np.load(folder / 'bank.npy'), meta['labels'], meta['ec_list']) |
| return P, meta['cfg'], bank, meta['alpha'] |
|
|
