Views
No views yet
0.12358924768426839.1import torch
2from transformers import AutoTokenizer, EsmForTokenClassification
3
4# Define the class mapping
5class_mapping = {
6 0: 'Not Binding Site',
7 1: 'Binding Site',
8}
9
10# Load the trained model and tokenizer
11model = EsmForTokenClassification.from_pretrained("AmelieSchreiber/esm2_t12_35M_UR50D_rna_binding_site_predictor")
12tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t12_35M_UR50D")
13
14# Define the new sequences
15new_sequences = [
16 'VLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTK',
17 'SQETFSDLWKLLPENNVLSPLPSQAMDDLMLSPDDIEQWF',
18 # ... add more sequences here ...
19]
20
21# Iterate over the new sequences
22for seq in new_sequences:
23 # Convert sequence to input IDs
24 inputs = tokenizer(seq, truncation=True, padding='max_length', max_length=1290, return_tensors="pt")["input_ids"]
25
26 # Apply the model to get the logits
27 with torch.no_grad():
28 outputs = model(inputs)
29
30 # Get the predictions by picking the label (class) with the highest logit
31 predictions = torch.argmax(outputs.logits, dim=-1)
32
33 # Convert the tensor to a list of integers
34 prediction_list = predictions.tolist()[0]
35
36 # Convert the predicted class indices to class names
37 predicted_labels = [class_mapping[pred] for pred in prediction_list]
38
39 # Create a list that matches each amino acid in the sequence to its predicted class label
40 residue_to_label = list(zip(list(seq), predicted_labels))
41
42 # Print out the list
43 for i, (residue, predicted_label) in enumerate(residue_to_label):
44 print(f"Position {i+1} - {residue}: {predicted_label}")