File size: 1,183 Bytes
1fdc49a | 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 | from einops import rearrange
from igfold import IgFoldInput
from igfold.utils.folding import get_sequence_dict, process_template
def embed(
antiberty,
model,
fasta_file=None,
sequences=None,
template_pdb=None,
ignore_cdrs=None,
ignore_chain=None,
mask=None,
):
seq_dict = get_sequence_dict(
sequences,
fasta_file,
)
embeddings, attentions = antiberty.embed(
seq_dict.values(),
return_attention=True,
)
embeddings = [e[1:-1].unsqueeze(0) for e in embeddings]
attentions = [a[:, :, 1:-1, 1:-1].unsqueeze(0) for a in attentions]
temp_coords, temp_mask = process_template(
template_pdb,
fasta_file,
ignore_cdrs=ignore_cdrs,
ignore_chain=ignore_chain,
)
model_in = IgFoldInput(
embeddings=embeddings,
attentions=attentions,
template_coords=temp_coords,
template_mask=temp_mask,
return_embeddings=True,
batch_mask=mask,
)
model_out = model(model_in)
prmsd = rearrange(
model_out.prmsd,
"b (l a) -> b l a",
a=4,
)
model_out.prmsd = prmsd
return model_out |