Update README.md

README.md

@@ -26,44 +26,38 @@ import torch
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 from joblib import load
 
-model_path = "littleworth/esm2_t6_8M_UR50D_pep2rec_cppp"
-model = AutoModelForSequenceClassification.from_pretrained(model_path)
-tokenizer = AutoTokenizer.from_pretrained(model_path)
+MODEL_PATH = "littleworth/esm2_t6_8M_UR50D_pep2rec_cppp"
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
 
-# Load the label encoder
-label_encoder = load(f"{model_path}/label_encoder.joblib")
+LABEL_ENCODER_PATH = f"{MODEL_PATH}/label_encoder.joblib"
+label_encoder = load(LABEL_ENCODER_PATH)
 
-# Define the input peptide sequence
-input_sequence = "GNLIVVGRVIMS"  # Example peptide sequence
 
-# Tokenize the input sequence
+input_sequence = "GNLIVVGRVIMS"
+
 inputs = tokenizer(input_sequence, return_tensors="pt", truncation=True, padding=True)
 
-# Make the prediction
 with torch.no_grad():
     outputs = model(**inputs)
-    logits = outputs.logits
-    probabilities = torch.softmax(logits, dim=1)
-    predicted_class_idx = torch.argmax(probabilities, dim=1).item()
+    probabilities = torch.softmax(outputs.logits, dim=1)
+    predicted_class_idx = probabilities.argmax(dim=1).item()
 
-# Decode the predicted class index to the original label
 predicted_class = label_encoder.inverse_transform([predicted_class_idx])[0]
 
-# Get the probabilities for each class
 class_probabilities = probabilities.squeeze().tolist()
 class_labels = label_encoder.inverse_transform(range(len(class_probabilities)))
 
-# Sort the class probabilities and labels in descending order
-sorted_indices = sorted(range(len(class_probabilities)), key=lambda i: class_probabilities[i], reverse=True)
-sorted_class_labels = [class_labels[i] for i in sorted_indices]
-sorted_class_probabilities = [class_probabilities[i] for i in sorted_indices]
+sorted_indices = torch.argsort(probabilities, descending=True).squeeze()
+sorted_class_labels = [class_labels[i] for i in sorted_indices.tolist()]
+sorted_class_probabilities = probabilities.squeeze()[sorted_indices].tolist()
 
-# Print the predicted class and probabilities
 print(f"Predicted Receptor Class: {predicted_class}")
 print("Top 10 Class Probabilities:")
 for label, prob in zip(sorted_class_labels[:10], sorted_class_probabilities[:10]):
     print(f"{label}: {prob:.4f}")
 
+
 ```
 
 Which gives this output: