Blame for dlmodel.py

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']