"""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')