LSTM-SA/prom_shap_generation.py at main · PEESEgroup/LSTM-SA · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import pandas as pd
from generate_shap import PeptideSHAPGenerator

if __name__ == "__main__":
    #Load the model paths
    promiscuous_model_paths = {
        'PP': 'ScoreModel_Results_onehot/exp_onehot_lstm_PP_l_2_h_512/model_type-lstm_epochs-500_learning_rate-0.001_batch_size-32_n_samples-346789_early_stopping_patience-10/Trained_Models/',
        'PE': 'ScoreModel_Results_onehot/exp_onehot_lstm_PE_l_2_h_512/model_type-lstm_epochs-500_learning_rate-0.001_batch_size-32_n_samples-572406_early_stopping_patience-10/Trained_Models/',
        'PET': 'ScoreModel_Results_onehot/exp_onehot_lstm_PET_l_2_h_512/model_type-lstm_epochs-500_learning_rate-0.001_batch_size-32_n_samples-353581_early_stopping_patience-10/Trained_Models/',
        'PVC': 'ScoreModel_Results_onehot/exp_onehot_lstm_PVC_l_2_h_512/model_type-lstm_epochs-500_learning_rate-0.001_batch_size-32_n_samples-166886_early_stopping_patience-10/Trained_Models/',
        'Nylon': 'ScoreModel_Results_onehot/exp_onehot_lstm_Nylon_l_2_h_512/model_type-lstm_epochs-500_learning_rate-0.001_batch_size-32_n_samples-114091_early_stopping_patience-10/Trained_Models/'
    }

    #Load the sequences to perform SHAP analysis on
    promiscuous_seq = pd.read_csv('Results/top1000_promiscuous_designs.csv')

    promiscuous_seq = list(promiscuous_seq.iloc[:, 0])

    prom_shap_generator = PeptideSHAPGenerator(plastic_type="promiscuous", device="cuda:0", explainer_type="gradient")
    prom_shap_generator.initialize_explainers(promiscuous_model_paths, n_samples=200)
    result = prom_shap_generator.generate_shap_analysis(promiscuous_seq)