HMD-AMP/prediction.py at main · ml4bio/HMD-AMP · 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
100
101
102
103
104
from sklearn.utils import shuffle
import esm
# import joblib
from Bio import SeqIO
import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
from deepforest import CascadeForestClassifier
from src.utils import *
import pandas as pd


device = "cuda" if torch.cuda.is_available() else "cpu"

'''
AMP/non-AMP prediction
specify AMP/non-AMP model path
'''
ftmodel_save_path = ''
clsmodel_save_path = ''


'''
prepare your data (in fasta format)
and labels (if any)
'''
sequences_file_path = ''
labels = []


seqs = []
for record in SeqIO.parse(sequences_file_path, "fasta"):
    seqs.append(str(record.seq))

seqs = np.array(seqs)
# generate sequence features
seq_embeddings, _, seq_ids = amp_feature_extraction(ftmodel_save_path, device, sequences_file_path)
# start prediction
cls_model = CascadeForestClassifier()
cls_model.load(clsmodel_save_path)

# binary classification
binary_pred = cls_model.predict(seq_embeddings)

# filter predicted amps
amp_index = np.where(binary_pred == 1)[0]
predicted_amps_seqs = seqs[amp_index]
predicted_amps_ids = np.array(seq_ids)[amp_index]
# predicted_amps_labels = seq_labels[amp_index]


non_amp_index = np.where(binary_pred == 0)[0]
predicted_nonamps_seqs = seqs[non_amp_index]
predicted_nonamps_ids = np.array(seq_ids)[non_amp_index]


# save AMP/non-AMP prediction results
result_AMP_df = pd.DataFrame(
    {
        'ID': predicted_amps_ids,
        'Sequence': predicted_amps_seqs,
    }
)
result_AMP_df['AMP'] = 'Yes'

result_nonAMP_df = pd.DataFrame(
    {
        'ID': predicted_nonamps_ids,
        'Sequence': predicted_nonamps_seqs,
    }
)

result_nonAMP_df['AMP'] = 'No'

result_df = pd.concat([result_AMP_df, result_nonAMP_df], axis=0).reset_index(drop=True)


# AMP targets prediction
for target in ['Gram+', 'Gram-', 'Mammalian_Cell', 'Virus', 'Fungus', 'Cancer']:
    '''
    specify target model path
    both fine-tuned feature extractor and trained classifier
    '''
    target_ftmodel_save_path = ''
    target_clsmodel_save_path = ''


    # generate target features
    target_ids, target_embeddings = targets_feature_extraction(target_ftmodel_save_path, device, predicted_amps_ids, predicted_amps_seqs)

    # start prediction
    target_cls_model = CascadeForestClassifier()
    target_cls_model.load(target_clsmodel_save_path)

    # AMP target annotation
    target_pred = target_cls_model.predict(target_embeddings)
    target_pred = pd.Series(target_pred).map({1: 'Yes', 0: 'No'})
    result_df[target] = target_pred.reindex(result_df.index, fill_value='N/A')


# show all target predictions
print(result_df)
# result_df.to_csv('prediction_result.csv', index=False)