| 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() |
|
|
| |
| 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 |
|
|
