|
|
from .model import EsmModel |
|
|
from .utils import load_hub_workaround |
|
|
import torch |
|
|
|
|
|
MODEL_URL = "https://zenodo.org/record/7042286/files/model.pth" |
|
|
|
|
|
def predict(peptide_list, device='cpu'): |
|
|
with torch.no_grad(): |
|
|
neuroPred_model = EsmModel() |
|
|
neuroPred_model.eval() |
|
|
state_dict = load_hub_workaround(MODEL_URL) |
|
|
|
|
|
neuroPred_model.load_state_dict(state_dict) |
|
|
neuroPred_model = neuroPred_model.to(device) |
|
|
prob, att = neuroPred_model(peptide_list, device) |
|
|
pred = torch.argmax(prob, dim=-1).cpu().tolist() |
|
|
att = att.cpu().numpy() |
|
|
out = {i[0]:[j,m[:, :len(i[1])]] for i, j, m in zip(peptide_list, pred, att)} |
|
|
return out |
|
|
|
