Update README.md

README.md

@@ -36,11 +36,11 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 
 # Load the fine-tuned model and tokenizer
-model_name = "sihuapeng/ESM2-finetuned-PPSL"
+model_name = "sihuapeng/PPPSL-ESM2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
-# Protein sequence for prediction the subcelluar localization
+# Sample protein sequence
 protein_sequence = "MSKKVLITGGAGYIGSVLTPILLEKGYEVCVIDNLMFDQISLLSCFHNKNFTFINGDAMDENLIRQEVAKADIIIPLAALVGAPLCKRNPKLAKMINYEAVKMISDFASPSQIFIYPNTNSGYGIGEKDAMCTEESPLRPISEYGIDKVHAEQYLLDKGNCVTFRLATVFGISPRMRLDLLVNDFTYRAYRDKFIVLFEEHFRRNYIHVRDVVKGFIHGIENYDKMKGQAYNMGLSSANLTKRQLAETIKKYIPDFYIHSANIGEDPDKRDYLVSNTKLEATGWKPDNTLEDGIKELLRAFKMMKVNRFANFN"
 
 # Encode the sequence as model input
@@ -50,15 +50,17 @@ inputs = tokenizer(protein_sequence, return_tensors="pt")
 with torch.no_grad():
     outputs = model(**inputs)
 
-# Get the prediction results
+# Get the prediction result
 logits = outputs.logits
 predicted_class_id = torch.argmax(logits, dim=-1).item()
+id2label = {0: 'CYtoplasmicMembrane', 1: 'Cellwall', 2: 'Cytoplasmic', 3: 'Extracellular', 4: 'OuterMembrane', 5: 'Periplasmic'}
+predicted_label = id2label[predicted_class_id]
 
 # 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(f"Predicted class Label: {predicted_label}")
 print ("===========================================================================================================================================")
+
 ```
 
 ## Funding