Update README.md

README.md

@@ -35,30 +35,29 @@ pip install transformers
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 
-# Load the fine-tuned model and tokenizer
-model_name = "sihuapeng/ESM2-finetuned-PPSL"
+# Load the fine-tuned model and tokenizer from Hugging Face
+model_name = "sihuapeng/TarPepSubLoc-ESM2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
-# Protein sequence for prediction the subcelluar localization
-protein_sequence = "MSKKVLITGGAGYIGSVLTPILLEKGYEVCVIDNLMFDQISLLSCFHNKNFTFINGDAMDENLIRQEVAKADIIIPLAALVGAPLCKRNPKLAKMINYEAVKMISDFASPSQIFIYPNTNSGYGIGEKDAMCTEESPLRPISEYGIDKVHAEQYLLDKGNCVTFRLATVFGISPRMRLDLLVNDFTYRAYRDKFIVLFEEHFRRNYIHVRDVVKGFIHGIENYDKMKGQAYNMGLSSANLTKRQLAETIKKYIPDFYIHSANIGEDPDKRDYLVSNTKLEATGWKPDNTLEDGIKELLRAFKMMKVNRFANFN"
+# Define the amino acid sequence
+sequence = "MKFLIVFVALFAMAVARPNLAEIVRQVSDVEPEKWSSDVETSDGTSIKQEGVLKNAGTDNEAAVVHGSFTWVDEKTGEKFTITYVADENGYQPQGAHLPVAPVA"
 
-# Encode the sequence as model input
-inputs = tokenizer(protein_sequence, return_tensors="pt")
+# Tokenize the sequence
+inputs = tokenizer(sequence, return_tensors="pt")
 
-# Perform inference using the model
+# Make the prediction
 with torch.no_grad():
     outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class_id = logits.argmax().item()
 
-# Get the prediction results
-logits = outputs.logits
-predicted_class_id = torch.argmax(logits, dim=-1).item()
+# Map the predicted class id to the class label
+labels = ["SP", "MT", "CH", "TH", "Other"]
+predicted_label = labels[predicted_class_id]
+
+print(f"The predicted class for the sequence is: {predicted_label}")
 
-# Output the predicted class
-print ("===========================================================================================================================================")
-print ("ID to Label mapping: {0: 'CYtoplasmicMembrane', 1: 'Cellwall', 2: 'Cytoplasmic', 3: 'Extracellular', 4: 'OuterMembrane', 5: 'Periplasmic'}")
-print(f"Predicted class ID: {predicted_class_id}")
-print ("===========================================================================================================================================")
 ```
 
 ## Funding