Views
No views yet
esm2_t6_8M_UR50D) for Token Classification0: other, 1: alpha helix, 2: beta strand. It was trained with
this notebook and achieves
78.13824286786025 % accuracy.1from transformers import AutoTokenizer, AutoModelForTokenClassification
2import numpy as np
3
4# 1. Prepare the Model and Tokenizer
5# Replace with the path where your trained model is saved if you're training a new model
6model_dir = "AmelieSchreiber/esm2_t6_8M_UR50D-finetuned-secondary-structure"
7
8model = AutoModelForTokenClassification.from_pretrained(model_dir)
9tokenizer = AutoTokenizer.from_pretrained(model_dir)
10
11# Define a mapping from label IDs to their string representations
12label_map = {0: "Other", 1: "Helix", 2: "Strand"}
13
14# 2. Tokenize the New Protein Sequence
15new_protein_sequence = "MAVPETRPNHTIYINNLNEKIKKDELKKSLHAIFSRFGQILDILVSRSLKMRGQAFVIFKEVSSATNALRSMQGFPFYDKPMRIQYAKTDSDIIAKMKGT" # Replace with your protein sequence
16tokens = tokenizer.tokenize(new_protein_sequence)
17inputs = tokenizer.encode(new_protein_sequence, return_tensors="pt")
18
19# 3. Predict with the Model
20with torch.no_grad():
21 outputs = model(inputs).logits
22 predictions = np.argmax(outputs[0].numpy(), axis=1)
23
24# 4. Decode the Predictions
25predicted_labels = [label_map[label_id] for label_id in predictions]
26
27# Print the tokens along with their predicted labels
28for token, label in zip(tokens, predicted_labels):
29 print(f"{token}: {label}")