Update README.md

README.md

@@ -41,7 +41,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
 # Define the amino acid sequence
-sequence = "MKFLIVFVALFAMAVARPNLAEIVRQVSDVEPEKWSSDVETSDGTSIKQEGVLKNAGTDNEAAVVHGSFTWVDEKTGEKFTITYVADENGYQPQGAHLPVAPVA"
+sequence = "MNSLLMITACLALVGTVWAKEGYLVNSYTGCKFECFKLGDNDYCLRECRQQYGKGSGGYCYAFGCWCTHLYEQAVVWPLPNKTCNGK"
 
 # Tokenize the sequence
 inputs = tokenizer(sequence, return_tensors="pt")
@@ -52,9 +52,11 @@ with torch.no_grad():
     logits = outputs.logits
     predicted_class_id = logits.argmax().item()
 
-# Map the predicted class id to the class label
-labels = ["SP", "MT", "CH", "TH", "Other"]
-predicted_label = labels[predicted_class_id]
+# Define the ID to Label mapping
+id2label = {0: 'CH', 1: 'MT', 2: 'Other', 3: 'SP', 4: 'TH'}
+
+# Get the predicted label
+predicted_label = id2label[predicted_class_id]
 
 print(f"The predicted class for the sequence is: {predicted_label}")