Views
No views yet
| Task | EMFP | Random Forest | ESM+MLP | ProtBERT+MLP |
|---|---|---|---|---|
| Authenticity | 0.967 | 0.718 | 0.892 | 0.856 |
| Canonical Protein Function | 0.932 | 0.505 | 0.827 | 0.791 |
1import torch
2import esm
3
4model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
5checkpoint = torch.load("best_model.pt")
6
7class ESMClassifier(torch.nn.Module):
8 def __init__(self, esm_model, num_labels=2, hidden_dim=1280, dropout=0.1):
9 super().__init__()
10 self.esm = esm_model
11 self.classifier = torch.nn.Sequential(
12 torch.nn.Dropout(dropout),
13 torch.nn.Linear(hidden_dim, hidden_dim // 2),
14 torch.nn.ReLU(),
15 torch.nn.Dropout(dropout),
16 torch.nn.Linear(hidden_dim // 2, num_labels)
17 )
18
19 def forward(self, tokens):
20 results = self.esm(tokens, repr_layers=[33], return_contacts=False)
21 return self.classifier(results["representations"][33][:, 0, :])
22
23classifier = ESMClassifier(model)
24classifier.load_state_dict(checkpoint['model_state_dict'])
25classifier.eval()
26
27# Predict
28batch_converter = alphabet.get_batch_converter()
29data = [("protein1", "MKTAYIAKQRQISFVKSHFSRQLEERLG")]
30labels, strs, tokens = batch_converter(data)
31
32with torch.no_grad():
33 logits = classifier(tokens)
34 probs = torch.softmax(logits, dim=1)
35 print(f"Probability of encoding canonical functional protein: {probs[0, 1].item():.4f}")esm2_t33_650M_UR50D)huggingface-cli download huangruihua/EMFP best_model.pt --local-dir ./1@software{emfp_2026,
2 title={EMFP: ESM-2 Micropeptide Predictor for Canonical Functional Proteins},
3 author={Huang, Rui-Hua},
4 year={2026},
5 url={https://github.com/huangruihua/EMFP}
6}