File size: 7,609 Bytes
aa45001 | 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 171 172 173 174 175 | """Direct Metalorian sampling."""
from metalorian.sampling import Sampler, FLAGS, METAL_LABELS, TOKEN_DROPOUT_SCALE
from metalorian.cli import parser, collect
import torch
import torch.nn.functional as F
CANVAS = 1024
class FixedGradSampler(Sampler):
def __init__(self, cfg, device='cuda:0'):
super().__init__(cfg, device)
assert float(self.aa_bias.abs().max()) == 0.0, "aa_bias must be zero (no chemistry)"
print(f"guide_on={cfg.guide_on} accumulate={not cfg.no_accum} "
f"grad_steps={cfg.grad_steps} scale={cfg.guidance_scale}")
def x0_hat(self, x_t, t, cond):
eps = self.net_sampler.model(x_t, t, cond)
return self.net_sampler.predict_xstart_from_eps(x_t, t, eps=eps)
def target_logit(self, lat, attn_len):
cfg = self.cfg
logits = self.model_con.decode_embeddings(lat[:, :attn_len])
seg = logits + self.vocab_mask[None, None, :]
soft = F.softmax(seg / cfg.tau, dim=-1)
emb = soft @ self.W
B = emb.shape[0]
cls_e = self.W[self.cls_id].view(1, 1, -1).expand(B, 1, -1)
eos_e = self.W[self.eos_id].view(1, 1, -1).expand(B, 1, -1)
inp = torch.cat([cls_e, emb, eos_e], dim=1) * TOKEN_DROPOUT_SCALE
am = torch.ones(B, attn_len + 2, device=self.device)
x_pred, _, _, _ = self.predictor.forward_from_embeds(inp, am)
tgt = x_pred[:, cfg.target_idx]
oth = x_pred[:, self.others].max(dim=1).values
return tgt - cfg.penalty * oth, soft
def init_latent(self, B, lmask):
return torch.randn(B, CANVAS, 1280, device=self.device) * lmask
@torch.no_grad()
def _encode_perpos(self, seqs):
"""Encode sequences as residue latents of shape (N, S, 1280)."""
S = self.cfg.seq_len
m = self.model_con.base_model
enc = self.tokenizer(seqs, return_tensors='pt', padding='max_length',
max_length=S + 2, truncation=True)
ids = enc['input_ids'].to(self.device); am = enc['attention_mask'].to(self.device)
emb = m.esm_model(input_ids=ids, attention_mask=am).last_hidden_state
x_pool, x_attns = m.attn_head(emb, am)
x_pool = m.attn_ln(x_pool + m.attn_skip(x_pool))
for lin in m.linear_layers:
r = x_pool; x_pool = F.silu(lin(x_pool)); x_pool = x_pool + r
x_w = torch.einsum('bhlk,bld->bhld', x_attns, x_pool).mean(dim=1)
return m.clf_ln(x_w)[:, 1:1 + S, :]
@torch.no_grad()
def _best_decode_one(self, lat_S, K, temp):
"""Select among K decodes by mean positional cosine similarity."""
S = self.cfg.seq_len
lg = self.model_con.decode_embeddings(lat_S.unsqueeze(0))[0] + self.vocab_mask[None, :]
probs = F.softmax(lg / temp, dim=-1)
toks = torch.multinomial(probs, K, replacement=True).T
cand = [self.tokenizer.decode(t, skip_special_tokens=True).replace(' ', '') for t in toks]
cand = [c for c in cand if len(c) == S]
if not cand:
return self.tokenizer.decode(lg.argmax(-1), skip_special_tokens=True).replace(' ', '')
xc = F.normalize(self._encode_perpos(cand), dim=-1)
Ln = F.normalize(lat_S, dim=-1)
cos = (xc * Ln[None]).sum(-1).mean(-1)
return cand[int(cos.argmax())]
def decode(self, x, attn_len):
if getattr(self.cfg, 'best_decode', 0) and self.cfg.best_decode > 0:
return [self._best_decode_one(x[i, :attn_len], self.cfg.best_decode,
self.cfg.decode_temp) for i in range(x.shape[0])]
return self.hard_decode(x, attn_len)
def run_round(self):
cfg, dev = self.cfg, self.device
B, S = cfg.batch, cfg.seq_len
lens = torch.full((B,), S, device=dev, dtype=torch.long)
attn = torch.arange(CANVAS, device=dev)[None, :] < lens[:, None]
lmask = attn.unsqueeze(-1).float()
x = self.init_latent(B, lmask)
tgt1h = torch.zeros(B, 15, device=dev)
tgt1h[:, cfg.target_idx] = 1.0
log_dis = torch.log(tgt1h + 1e-10)
best_score = [-1.0] * B
best_seq = [None] * B
cys_trace = []
for step in reversed(range(FLAGS.T)):
t = torch.full((B,), step, device=dev, dtype=torch.long)
for k in range(cfg.grad_steps):
x = x.detach().requires_grad_(True)
if cfg.guide_on == 'x0':
lat = self.x0_hat(x, t, log_dis)
else:
lat = x
obj, soft = self.target_logit(lat, S)
g, = torch.autograd.grad(obj.sum(), x)
with torch.no_grad():
gn = g.flatten(1).norm(dim=1).clamp_min(1e-8).view(-1, 1, 1)
stepv = cfg.guidance_scale * (g / gn) * lmask
if cfg.no_accum:
mean, _ = self.net_sampler.p_mean_variance(
x_t=x.detach(), t=t, cond=log_dis, trans=None)
x = (mean * lmask + stepv).clamp(-1, 1)
else:
x = (x.detach() + stepv).clamp(-1, 1)
x = x * lmask
with torch.no_grad():
mean, log_var = self.net_sampler.p_mean_variance(
x_t=x, t=t, cond=log_dis, trans=None)
mean = (mean * lmask).float()
noise = torch.randn_like(x) if step > 0 else torch.zeros_like(x)
x = (mean + torch.exp(0.5 * log_var.float()) * noise * lmask).clamp(-1, 1)
x = x * lmask
log_dis = self.trainer_dis.p_sample(log_dis, t, x)
if step % 2 == 0:
c = cfg.label_clamp
log_dis = (1 - c) * log_dis + c * torch.log(tgt1h + 1e-10)
if cfg.track_best:
with torch.no_grad():
cur = self.hard_decode(x, S)
ok = [i for i, q in enumerate(cur) if len(q) == S]
if ok:
pr, _ = self.score([cur[i] for i in ok])
sc = pr[:, cfg.target_idx]
for j, i in enumerate(ok):
if sc[j].item() > best_score[i]:
best_score[i] = sc[j].item()
best_seq[i] = cur[i]
if step % 10 == 0:
with torch.no_grad():
seqs = self.hard_decode(x, S)
cys_trace.append(sum(q.count('C') for q in seqs) / max(len(seqs), 1))
if cfg.track_best:
final = self.decode(x, S)
out = [best_seq[i] if best_seq[i] is not None else final[i] for i in range(B)]
return out, cys_trace
return self.decode(x, S), cys_trace
def accept(self, seqs, _unused):
cfg = self.cfg
probs, fires = self.score(seqs)
out = []
for i, s in enumerate(seqs):
if len(s) != cfg.seq_len:
continue
sc = probs[i, cfg.target_idx].item()
if cfg.accept_all:
out.append({'sequence': s, 'score': sc, 'n_cys': s.count('C'),
'n_his': s.count('H')})
continue
if sc >= self.threshold and fires[i, cfg.target_idx].item() >= 0.5:
out.append({'sequence': s, 'score': sc, 'n_cys': s.count('C'),
'n_his': s.count('H')})
return out, probs
if __name__ == '__main__':
collect(parser('direct').parse_args(), FixedGradSampler, 'direct')
|