Update utils.py

utils.py

@@ -1,3 +1,4 @@
+from transformers import AutoTokenizer
 import torch
 
 def validate_sequence(sequence):
@@ -5,13 +6,13 @@ def validate_sequence(sequence):
     return all(aa in valid_amino_acids for aa in sequence) and len(sequence) <= 200
 
 def load_model():
-    # Assuming the model is a simple PyTorch model, adjust the path as needed
+    # Load your model as before
     model = torch.load('solubility_model.pth', map_location=torch.device('cpu'))
     model.eval()
     return model
 
 def predict(model, sequence):
-    # Dummy tensor conversion, replace with your actual model's input handling
-    tensor = torch.tensor([ord(char) for char in sequence], dtype=torch.float32)
-    output = model(tensor)
-    return output.item()
+    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()