| from DeepPD.model import MyModel,DeepPD | |
| import torch | |
| import torch.nn as nn | |
| from DeepPD.config import ArgsConfig | |
| args = ArgsConfig() | |
| softmax = nn.Softmax(1) | |
| def predict(seqs,data,model_path,threshold=0.5, device=args.device): | |
| with torch.no_grad(): | |
| model = DeepPD(vocab_size=21,embedding_size=args.embedding_size,esm_path=args.ems_path,layer_idx=args.esm_layer_idx,seq_len=args.max_len,dropout=args.dropout, | |
| fan_layer_num=1,num_heads=8,encoder_layer_num=1,Contrastive_Learning=False,info_bottleneck=args.info_bottleneck).to(args.device) | |
| model.eval() | |
| state_dict = torch.load(model_path, map_location=device) | |
| model.load_state_dict(state_dict,strict=False) | |
| model.to(device) | |
| seqs = seqs.to(device) | |
| out,_,_ = model(seqs) | |
| prob = softmax(out)[:,1] | |
| final_out = [] | |
| for i, j in zip(data, prob): | |
| temp = [i[0], i[1], f"{j:.3f}", 'Peptide' if j >threshold else 'Non-Peptide'] | |
| final_out.append(temp) | |
| return final_out |