FlashABB/example.py at main · oxpig/FlashABB · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from flash_abb import pretrained
import torch

flabb = pretrained(device='cuda')

# Sequences in heavy|light format
seqs = [
    'EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTTVTVSS|DIQLTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK',
    'EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTTVTVS|DIQLTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK',
]

with torch.no_grad():
    result = flabb(seqs)

print(result.coords.shape)          # (2, n_residues, 37, 3)
print(result.bb_coords.shape)       # (2, n_residues, 4, 3)

result.to_pdbs(['ab1', 'ab2'], pdb_dir='sample_preds')