File size: 4,166 Bytes
48b5986 | 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 | # Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import logging
from typing import Tuple, Dict
import torch
from torch import nn
from esm.data import Alphabet
from utils.tensor import add_eos_bos, is_1hot_tensor
from torch import nn
from torch.nn import functional as F
from abc import ABC, abstractmethod
import hydra
import os
from utils.linear_projection import LinearProjectionDistogramModel
logger = logging.getLogger(__name__) # Hydra configured
PDB_LOADER_PARAM_REGISTRY = {
'LinearProjectionDist-1A': {
"DMIN": 2.5,
"DMAX": 20.0,
"DBINS": 17,
"ABINS": 17,
"PHI_BINS": 7,
"WMIN": 0.8,
"LMIN": 150,
"LMAX": 400,
"EPOCHS": 10,
"NCPU": 8,
"SLICE": "CONT",
"contact_bin_cutoff": (0, 5) # both inclusive
},
}
class WrapStruct(nn.Module, ABC):
"""
Wraps structure models used by Designer.
Companion of WrapLM.
"""
def __init__(self, model, vocab, has_eos=False):
super().__init__()
self.model = model
self.vocab = vocab
self.has_eos = has_eos
@abstractmethod
def forward(self, seq1h, pos_idx=None):
raise NotImplementedError
@property
def msa_based(self):
return False
def load(vocab: Alphabet) -> Tuple[nn.Module, Dict]:
pdb_loader_params = PDB_LOADER_PARAM_REGISTRY['LinearProjectionDist-1A']
model = load_model()
has_eos = True
return WrapStructModel(model, vocab, has_eos), pdb_loader_params
def load_model():
model_url_path = 'https://dl.fbaipublicfiles.com/fair-esm/examples/lm_design/linear_projection_model.pt'
local_model_path = hydra.utils.to_absolute_path('./linear_projection_model.pt')
if not os.path.exists(local_model_path):
logger.info(f'Downloading linear projection model from {model_url_path} to {local_model_path}')
os.system(f'wget {model_url_path} -O {local_model_path}')
# load model_path
state = torch.load(local_model_path, map_location='cpu')
# chkpt_args = state['cfg']
model = LinearProjectionDistogramModel()
model_state = state['model']
model.load_state_dict(model_state)
from esm.pretrained import esm2_t33_650M_UR50D
base_model, alphabet = esm2_t33_650M_UR50D()
model.base_model = base_model
return model
class WrapStructModel(WrapStruct):
def forward(self, seq1h, pos_idx=None):
assert is_1hot_tensor(seq1h)
B, L, K = seq1h.shape
# Prepend bos/cls only.
eos_idx = self.vocab.eos_idx if self.has_eos else None
seq1h = add_eos_bos(seq1h, bos_idx=self.vocab.cls_idx, eos_idx=eos_idx)
# In designer, all seqs have same length.
src_lengths = torch.full((B,), L + 1 + int(self.has_eos)).to(seq1h).long()
# In some cases, a vocab padded to 64 positions
# with dummy character was used.
# As a workaround, pad K dimensions up to 64 (embed_dim),
# then remove this portion after to get K tokens back.
embed_dim = self.model.base_model.embed_tokens.weight.size(0)
if embed_dim != K:
seq1h = torch.cat([
seq1h,
torch.zeros(B, L+1, embed_dim - K).to(seq1h)
], -1)
# Call forward. bos/eos will be stripped inside.
out_dict = self.model(seq1h, src_lengths=src_lengths)
return self._process_output(out_dict)
@property
def lm(self):
return self.model.base_model
@staticmethod
def _process_output(out_dict):
for k in ('logits', 'theta_logits', 'phi_logits', 'omega_logits'):
out_dict[k] = out_dict[k].permute(0, 2, 3, 1).contiguous()
if 'logits' in out_dict:
out_dict['dist_logits'] = out_dict['logits']
del out_dict['logits']
for k in [k for k in out_dict if k.endswith('_logits')]:
targetname = k[:-len('_logits')]
out_dict[f'p_{targetname}'] = F.softmax(out_dict[k], -1)
return out_dict
|