This is the 650M parameter version of ESM-2, finetuned with QLoRA to predict binding sites of proteins based on single sequences alone.
No multiple sequence alignment or structure is required. The embeddings from this model can also be used in structural models. The model is trained on
approximately 12M protein sequences from UniProt, with an 80/20 train/test split.
1'eval_loss': 0.05597764626145363,
2'eval_accuracy': 0.9829392036087405,
3'eval_precision': 0.5626191259397847,
4'eval_recall': 0.9488112528941492,
5'eval_f1': 0.7063763773187873,
6'eval_auc': 0.9662524626230765,
7'eval_mcc': 0.7235838533979579
Due to the size of the dataset we had to get the test metrics in chunks and aggregate. To see the metrics for each chunk,
refer to this text file.
1'eval_loss': 0.16281947493553162,
2'eval_accuracy': 0.9569658774883986,
3'eval_precision': 0.3209956738348438,
4'eval_recall': 0.7883697002335764,
5'eval_f1': 0.4562306866120791,
6'eval_auc': 0.8746433990040084,
7'eval_mcc': 0.48648765699020435
The metrics for the earlier checkpoints are not reported here yet.
1from transformers import AutoModelForTokenClassification, AutoTokenizer
2from peft import PeftModel
3import torch
4
5# Path to the saved LoRA model
6model_path = "AmelieSchreiber/esm2_t33_650M_qlora_binding_12M"
7# ESM2 base model
8base_model_path = "facebook/esm2_t33_650M_UR50D"
9
10# Load the model
11base_model = AutoModelForTokenClassification.from_pretrained(base_model_path)
12loaded_model = PeftModel.from_pretrained(base_model, model_path)
13
14# Ensure the model is in evaluation mode
15loaded_model.eval()
16
17# Load the tokenizer
18loaded_tokenizer = AutoTokenizer.from_pretrained(base_model_path)
19
20# Protein sequence for inference
21protein_sequence = "MAVPETRPNHTIYINNLNEKIKKDELKKSLHAIFSRFGQILDILVSRSLKMRGQAFVIFKEVSSATNALRSMQGFPFYDKPMRIQYAKTDSDIIAKMKGT" # Replace with your actual sequence
22
23# Tokenize the sequence
24inputs = loaded_tokenizer(protein_sequence, return_tensors="pt", truncation=True, max_length=1024, padding='max_length')
25
26# Run the model
27with torch.no_grad():
28 logits = loaded_model(**inputs).logits
29
30# Get predictions
31tokens = loaded_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]) # Convert input ids back to tokens
32predictions = torch.argmax(logits, dim=2)
33
34# Define labels
35id2label = {
36 0: "No binding site",
37 1: "Binding site"
38}
39
40# Print the predicted labels for each token
41for token, prediction in zip(tokens, predictions[0].numpy()):
42 if token not in ['<pad>', '<cls>', '<eos>']:
43 print((token, id2label[prediction]))