Views
No views yet
pytorch_model.bin: Trained model weights.alphabet.bin: ESM2 alphabet (tokenizer).config.json: Model configuration.README.md: This file.pip install torch esm biopython huggingface_hub1import torch
2import torch.nn as nn
3import esm
4from huggingface_hub import hf_hub_download
5import json
6
7# Define the classifier architecture (must match training)
8class ProteinClassifier(nn.Module):
9 def __init__(self, esm_model, embedding_dim, num_classes):
10 super(AntiFP2Classifier, self).__init__()
11 self.esm_model = esm_model
12 self.fc = nn.Linear(embedding_dim, num_classes)
13 def forward(self, tokens):
14 with torch.no_grad():
15 results = self.esm_model(tokens, repr_layers=[36])
16 embeddings = results["representations"][36].mean(1)
17 output = self.fc(embeddings)
18 return output
19
20# Download model files from Hugging Face Hub
21repo_id = "raghavagps-group/antifp2"
22model_weights_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")
23alphabet_path = hf_hub_download(repo_id=repo_id, filename="alphabet.bin")
24config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
25
26# Load ESM2 backbone model
27esm_model, alphabet = esm.pretrained.esm2_t36_3B_UR50D()
28
29# Load configuration
30with open(config_path, 'r') as f:
31 config = json.load(f)
32
33# Initialize classifier
34classifier = ProteinClassifier(esm_model, embedding_dim=config['embedding_dim'], num_classes=config['num_classes'])
35
36# Load weights
37classifier.load_state_dict(torch.load(model_weights_path))
38classifier.eval()
39
40# Load alphabet tokenizer
41alphabet = torch.load(alphabet_path)
42batch_converter = alphabet.get_batch_converter()
43
44# Move model to GPU if available
45device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
46classifier = classifier.to(device)