pocket_binding_prediction_challenge/create_esm_embeddings.py at main · BinaryChrisEntropy/pocket_binding_prediction_challenge · GitHub

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
from Bio.PDB import PDBParser
import esm
from Bio.SeqUtils import seq1
import os
import torch
import pickle
from tqdm import tqdm


test_dir = r"./test"

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

output_dir = "./esm_test"
esm_model, esm_alphabet = esm.pretrained.esm2_t33_650M_UR50D()
esm_model = esm_model.to(device)
esm_model.eval()
esm_batch_converter = esm_alphabet.get_batch_converter()


AMINO_ACID_PROPERTIES = {
    'ALA': [1.8,  0, 0],
    'ARG': [-4.5, +1, 1],
    'ASN': [-3.5,  0, 1],
    'ASP': [-3.5, -1, 1],
    'CYS': [2.5,  0, 1],
    'GLN': [-3.5,  0, 1],
    'GLU': [-3.5, -1, 1],
    'GLY': [-0.4, 0, 0],
    'HIS': [-3.2, +1, 1],
    'ILE': [4.5,  0, 0],
    'LEU': [3.8,  0, 0],
    'LYS': [-3.9, +1, 1],
    'MET': [1.9,  0, 0],
    'PHE': [2.8,  0, 0],
    'PRO': [-1.6, 0, 0],
    'SER': [-0.8, 0, 1],
    'THR': [-0.7, 0, 1],
    'TRP': [-0.9, 0, 1],
    'TYR': [-1.3, 0, 1],
    'VAL': [4.2,  0, 0],
}


def calc_embeddings(pdb_path):
    pdb_name = os.path.basename(pdb_path)
    base_name = os.path.splitext(pdb_name)[0]
    output_file = os.path.join(output_dir, base_name + ".pkl")


    # jump over if already processed
    if os.path.exists(output_file):
        return

    try:
        parser = PDBParser(QUIET=True)
        structure = parser.get_structure("protein", pdb_path)
        model = structure[0]

        resnames = []
        for chain in model:
            for residue in chain:
                if residue.id[0] != " ":
                    continue
                resnames.append([residue.get_resname()])

        try:
            sequence = ''.join([seq1(res[0]) if res[0] in AMINO_ACID_PROPERTIES else 'X' for res in resnames])
        except KeyError:
            sequence = ''.join([seq1(res[0], custom_map={'UNK': 'X'}) if isinstance(res[0], str) else 'X' for res in resnames])

        data = [("protein", sequence)]
        _, _, batch_tokens = esm_batch_converter(data)
        batch_tokens = batch_tokens.to(device)
        with torch.no_grad():
            results = esm_model(batch_tokens, repr_layers=[33])
            embedding = results["representations"][33][0, 1:len(sequence)+1]

        esm_features = embedding.cpu()

        # Save
        with open(output_file, 'wb') as f:
            pickle.dump(esm_features, f)

    except Exception as e:
        print(f"Error for {pdb_name}: {e}")


if __name__ == "__main__":
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    pdb_files = [os.path.join(test_dir, f) for f in os.listdir(test_dir) if f.endswith('.pdb')]
    for f in tqdm(pdb_files, desc="Processing"):
        calc_embeddings(f)