Update utils.py

utils.py

@@ -1,5 +1,6 @@
 from transformers import AutoTokenizer
 import torch
+import torch.nn.functional as F
 
 def validate_sequence(sequence):
     valid_amino_acids = set("ACDEFGHIKLMNPQRSTVWY")  # 20 standard amino acids
@@ -15,4 +16,7 @@ def predict(model, sequence):
     tokenizer = AutoTokenizer.from_pretrained('facebook/esm2_t6_8M_UR50D')
     tokenized_input = tokenizer(sequence, return_tensors="pt", truncation=True, padding=True)
     output = model(**tokenized_input)
-    return output.item()
+    logits = output.logits  # Extract logits
+    probabilities = F.softmax(logits, dim=-1)  # Apply softmax to convert logits to probabilities
+    predicted_label = torch.argmax(probabilities, dim=-1)  # Get the predicted label
+    return predicted_label.item()  # Return the label as a Python integer