Views
No views yet
1Train:
2({'accuracy': 0.9406146072672105,
3 'precision': 0.2947122459102886,
4 'recall': 0.952624323712029,
5 'f1': 0.4501592605994876,
6 'auc': 0.9464622170085311,
7 'mcc': 0.5118390407598565},
8Test:
9 {'accuracy': 0.9266827008067329,
10 'precision': 0.22378953253253775,
11 'recall': 0.7790246675002842,
12 'f1': 0.3476966444342296,
13 'auc': 0.8547531675185658,
14 'mcc': 0.3930283737012391})1from transformers import AutoModelForTokenClassification, AutoTokenizer
2from peft import PeftModel
3import torch
4
5# Path to the saved LoRA model
6model_path = "AmelieSchreiber/esm2_t12_35M_lora_binding_sites_cp1"
7# ESM2 base model
8base_model_path = "facebook/esm2_t12_35M_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]))1from datasets import Dataset
2from transformers import AutoTokenizer
3import pickle
4
5# Load tokenizer
6tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t12_35M_UR50D")
7
8# Function to truncate labels
9def truncate_labels(labels, max_length):
10 """Truncate labels to the specified max_length."""
11 return [label[:max_length] for label in labels]
12
13# Set the maximum sequence length
14max_sequence_length = 1000
15
16# Load the data from pickle files
17with open("train_sequences_chunked_by_family.pkl", "rb") as f:
18 train_sequences = pickle.load(f)
19with open("test_sequences_chunked_by_family.pkl", "rb") as f:
20 test_sequences = pickle.load(f)
21with open("train_labels_chunked_by_family.pkl", "rb") as f:
22 train_labels = pickle.load(f)
23with open("test_labels_chunked_by_family.pkl", "rb") as f:
24 test_labels = pickle.load(f)
25
26# Tokenize the sequences
27train_tokenized = tokenizer(train_sequences, padding=True, truncation=True, max_length=max_sequence_length, return_tensors="pt", is_split_into_words=False)
28test_tokenized = tokenizer(test_sequences, padding=True, truncation=True, max_length=max_sequence_length, return_tensors="pt", is_split_into_words=False)
29
30# Truncate the labels to match the tokenized sequence lengths
31train_labels = truncate_labels(train_labels, max_sequence_length)
32test_labels = truncate_labels(test_labels, max_sequence_length)
33
34# Create train and test datasets
35train_dataset = Dataset.from_dict({k: v for k, v in train_tokenized.items()}).add_column("labels", train_labels)
36test_dataset = Dataset.from_dict({k: v for k, v in test_tokenized.items()}).add_column("labels", test_labels)1from sklearn.metrics import(
2 matthews_corrcoef,
3 accuracy_score,
4 precision_recall_fscore_support,
5 roc_auc_score
6)
7from peft import PeftModel
8from transformers import DataCollatorForTokenClassification, AutoModelForTokenClassification
9from transformers import Trainer
10from accelerate import Accelerator
11
12# Instantiate the accelerator
13accelerator = Accelerator()
14
15# Define paths to the LoRA and base models
16base_model_path = "facebook/esm2_t12_35M_UR50D"
17lora_model_path = "AmelieSchreiber/esm2_t12_35M_lora_binding_sites_cp1" # "path/to/your/lora/model" Replace with the correct path to your LoRA model
18
19# Load the base model
20base_model = AutoModelForTokenClassification.from_pretrained(base_model_path)
21
22# Load the LoRA model
23model = PeftModel.from_pretrained(base_model, lora_model_path)
24model = accelerator.prepare(model) # Prepare the model using the accelerator
25
26# Define label mappings
27id2label = {0: "No binding site", 1: "Binding site"}
28label2id = {v: k for k, v in id2label.items()}
29
30# Create a data collator
31data_collator = DataCollatorForTokenClassification(tokenizer)
32
33# Define a function to compute the metrics
34def compute_metrics(dataset):
35 # Get the predictions using the trained model
36 trainer = Trainer(model=model, data_collator=data_collator)
37 predictions, labels, _ = trainer.predict(test_dataset=dataset)
38
39 # Remove padding and special tokens
40 mask = labels != -100
41 true_labels = labels[mask].flatten()
42 flat_predictions = np.argmax(predictions, axis=2)[mask].flatten().tolist()
43
44 # Compute the metrics
45 accuracy = accuracy_score(true_labels, flat_predictions)
46 precision, recall, f1, _ = precision_recall_fscore_support(true_labels, flat_predictions, average='binary')
47 auc = roc_auc_score(true_labels, flat_predictions)
48 mcc = matthews_corrcoef(true_labels, flat_predictions) # Compute the MCC
49
50 return {"accuracy": accuracy, "precision": precision, "recall": recall, "f1": f1, "auc": auc, "mcc": mcc} # Include the MCC in the returned dictionary
51
52# Get the metrics for the training and test datasets
53train_metrics = compute_metrics(train_dataset)
54test_metrics = compute_metrics(test_dataset)
55
56train_metrics, test_metrics