from .paths import ROOT, BACKBONE, PREDICTOR, DIFFUSION, TRAIN, configure configure() import argparse import numpy as np import torch from absl import flags from transformers import AutoTokenizer, EsmModel from datasets import load_from_disk from .continuous import GaussianDiffusionSampler, DiffusionProteinModel from .discrete import MultinomialDiffusion from .unet import tabularUnet from .predictor import MultitaskProteinModel FLAGS = flags.FLAGS METAL_LABELS = ['Ag','Cd','Co','Cu','Fe','Hg','Mn','Mo','Ni','Pb','Pt','V','W','Zn','NA'] TOKEN_DROPOUT_SCALE = 0.88 CODI_CKPT, PRED_CKPT, TRAIN_DS, ESM_NAME = map(str, (DIFFUSION,PREDICTOR,TRAIN,BACKBONE)) def build_hparams(): ns = argparse.Namespace() ns.date = 'test' ns.attempts = '1' ns.output_model_path = str(ROOT/'runs') ns.num_workers = 8 ns.num_linear_layers = 3 ns.batch_size = 5 ns.devices = 1 ns.special_prefix = None ns.focal_gamma = 1 ns.lr = 1e-3 ns.eps = 0.4 ns.n_splits = 2 ns.epochs = 100 ns.save_top_k = 3 return ns def define_flags(): if not FLAGS.is_parsed(): FLAGS([""]) d = { 'encoder_dim_con': ('string', "64,128,256"), 'encoder_dim_dis': ('string', "64,128,256"), 'nf_con': ('int', 16), 'nf_dis': ('int', 64), 'input_size': ('int', None), 'cond_size': ('int', None), 'output_size': ('int', None), 'activation': ('string', 'relu'), 'nf': ('int', None), 'hidden_size': ('int', 768), 'num_layers': ('int', 12), 'num_heads': ('int', 12), 'intermediate_size': ('int', 4 * 768), 'max_seq_length': ('int', 1280), 'T': ('int', 50), 'beta_1': ('float', 0.00001), 'beta_T': ('float', 0.02), } for name, (kind, default) in d.items(): if not hasattr(FLAGS, name): {'string': flags.DEFINE_string, 'int': flags.DEFINE_integer, 'float': flags.DEFINE_float}[kind](name, default, help=name) if not hasattr(FLAGS, 'encoder_dim'): flags.DEFINE_multi_integer('encoder_dim', None, help='encoder_dim') if not hasattr(FLAGS, 'mean_type'): flags.DEFINE_enum('mean_type', 'epsilon', ['xprev', 'xstart', 'epsilon'], help='mean type') if not hasattr(FLAGS, 'var_type'): flags.DEFINE_enum('var_type', 'fixedsmall', ['fixedlarge', 'fixedsmall'], help='var type') def load_diffusion_models(checkpoint_path, base_model, device): ckpt = torch.load(checkpoint_path, map_location=device) model_con_state = ckpt['model_con'] base_keys = [k for k in model_con_state if 'base_model' in k] if base_keys: base_state = {k.replace('base_model.', ''): model_con_state[k] for k in base_keys} base_model.load_state_dict(base_state, strict=False) print("Loaded base model weights from CoDi checkpoint") def strip(sd): return {k.replace('module.', ''): v for k, v in sd.items()} if 'module.' in next(iter(ckpt['model_con'])): ckpt['model_con'] = strip(ckpt['model_con']) if 'module.' in next(iter(ckpt['model_dis'])): ckpt['model_dis'] = strip(ckpt['model_dis']) FLAGS.input_size, FLAGS.cond_size, FLAGS.output_size = 1280, 15, 1280 FLAGS.encoder_dim = [int(x) for x in FLAGS.encoder_dim_con.split(',')] FLAGS.nf = FLAGS.nf_con FLAGS.hidden_size, FLAGS.num_layers, FLAGS.num_heads = 768, 12, 12 FLAGS.intermediate_size, FLAGS.max_seq_length = 4 * 768, 1280 model_con = DiffusionProteinModel(base_model, FLAGS).to(device) filtered = {} for k, v in ckpt['model_con'].items(): if not k.startswith('base_model.'): filtered[k] = v else: bk = k.replace('base_model.', '') if bk in base_model.state_dict() and v.shape == base_model.state_dict()[bk].shape: filtered[k] = v model_con.load_state_dict(filtered, strict=False) FLAGS.input_size, FLAGS.cond_size, FLAGS.output_size = 15, 1280, 15 FLAGS.encoder_dim = [int(x) for x in FLAGS.encoder_dim_dis.split(',')] FLAGS.nf = FLAGS.nf_dis model_dis = tabularUnet(FLAGS).to(device) model_dis.load_state_dict(ckpt['model_dis']) trainer_dis = MultinomialDiffusion( num_classes=np.array([15]), shape=torch.Size([10, 15]), denoise_fn=model_dis, FLAGS=FLAGS, timesteps=FLAGS.T, loss_type='vb_stochastic').to(device) net_sampler = GaussianDiffusionSampler( model_con, FLAGS.beta_1, FLAGS.beta_T, FLAGS.T, mean_type=FLAGS.mean_type, var_type=FLAGS.var_type).to(device) return model_con, model_dis, trainer_dis, net_sampler class Sampler: def __init__(self, cfg, device='cuda:0'): self.cfg = cfg self.device = torch.device(device) define_flags() self.tokenizer = AutoTokenizer.from_pretrained(ESM_NAME) self.cys_id = self.tokenizer.convert_tokens_to_ids('C') self.cls_id = self.tokenizer.cls_token_id self.eos_id = self.tokenizer.eos_token_id ds = load_from_disk(TRAIN_DS) labels = np.stack([l for b in ds['anchor_labels'] for l in b]) cw = 1.0 / (labels.sum(axis=0) + 1e-6) cw = torch.tensor(cw / cw.sum() * labels.shape[1], dtype=torch.float32).to(self.device) hp = build_hparams() esm = EsmModel.from_pretrained(ESM_NAME).cpu() base_model = MultitaskProteinModel.load_from_checkpoint( PRED_CKPT, hparams=hp, esm_model=esm, num_layers_to_finetune=2, input_dim=1280, output_dim=15, class_weights=cw).cpu() esm2 = EsmModel.from_pretrained(ESM_NAME).cpu() self.predictor = MultitaskProteinModel.load_from_checkpoint( PRED_CKPT, hparams=hp, esm_model=esm2, num_layers_to_finetune=2, input_dim=1280, output_dim=15, class_weights=cw).to(self.device).eval() self.model_con, self.model_dis, self.trainer_dis, self.net_sampler = \ load_diffusion_models(CODI_CKPT, base_model, self.device) self.model_con.eval() self.model_dis.eval() # Gradients are needed only with respect to the latent. for m in (self.predictor, self.model_con, self.model_dis): for p in m.parameters(): p.requires_grad_(False) self.W = self.predictor.esm_model.embeddings.word_embeddings.weight self.threshold = self.predictor.classification_thresholds[cfg.target_idx].item() self.vocab_mask = torch.full((33,), float('-inf'), device=self.device) self.vocab_mask[4:24] = 0.0 self.aa_bias = torch.zeros(33, device=self.device) self.aa_bias[self.cys_id] = cfg.cys_logit_bias self.others = [i for i in range(15) if i != cfg.target_idx] print(f"target={METAL_LABELS[cfg.target_idx]} idx={cfg.target_idx} " f"threshold={self.threshold:.4f} cys_token={self.cys_id}") @torch.no_grad() def hard_decode(self, x, attn_len): logits = self.model_con.decode_embeddings(x[:, :attn_len]) biased = logits + self.vocab_mask[None, None, :] + self.aa_bias[None, None, :] toks = torch.argmax(biased, dim=-1) return [self.tokenizer.decode(t, skip_special_tokens=True).replace(' ', '') for t in toks] @torch.no_grad() def score(self, seqs): enc = self.tokenizer(seqs, return_tensors='pt', padding=True) enc = {k: v.to(self.device) for k, v in enc.items()} probs, fires, _ = self.predictor.predict(input_ids=enc['input_ids'], attention_masks=enc['attention_mask']) return probs, fires