Update esm embedder function

app.py

@@ -4,6 +4,7 @@ import numpy as np
 from transformers import T5Tokenizer, T5EncoderModel
 import esm
 from inference import load_models, predict_ensemble
+from transformers import AutoTokenizer, AutoModel
 
 # Load trained models
 model_protT5, model_cat = load_models()
@@ -13,11 +14,10 @@ tokenizer_t5 = T5Tokenizer.from_pretrained("Rostlab/prot_t5_xl_uniref50", do_low
 model_t5 = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_uniref50")
 model_t5 = model_t5.eval()
 
-# Load ESM model
-esm_model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
-batch_converter = alphabet.get_batch_converter()
-esm_model.eval()
-
+# Load the tokenizer and model
+model_name = "facebook/esm2_t33_650M_UR50D"
+tokenizer_esm = AutoTokenizer.from_pretrained(model_name)
+esm_model = AutoModel.from_pretrained(model_name)
 
 def extract_prott5_embedding(sequence):
     sequence = sequence.replace(" ", "")
@@ -28,14 +28,21 @@ def extract_prott5_embedding(sequence):
     return torch.mean(embedding, dim=1)
 
 
+# Extract ESM2 embedding
 def extract_esm_embedding(sequence):
-    batch_labels, batch_strs, batch_tokens = batch_converter([("protein1", sequence)])
+    # Tokenize the sequence
+    inputs = tokenizer_esm(sequence, return_tensors="pt", padding=True, truncation=True)
+    
+    # Forward pass through the model
     with torch.no_grad():
-        results = esm_model(batch_tokens, repr_layers=[33], return_contacts=False)
-    token_representations = results["representations"][33]
+        outputs = esm_model(**inputs)
+    
+    # Extract the embeddings from the 33rd layer (ESM2 layer)
+    token_representations = outputs.last_hidden_state  # This is the default layer
     return torch.mean(token_representations[0, 1:len(sequence)+1], dim=0).unsqueeze(0)
 
 
+
 def classify(sequence):
     protT5_emb = extract_prott5_embedding(sequence)
     esm_emb = extract_esm_embedding(sequence)