Update model code for Hugging Face ZeroGPU compatibility

app.py

@@ -5,6 +5,7 @@ from transformers import T5Tokenizer, T5EncoderModel
 import esm
 from inference import load_models, predict_ensemble
 from transformers import AutoTokenizer, AutoModel
+import spaces
 
 # Load trained models
 model_protT5, model_cat = load_models()
@@ -12,12 +13,12 @@ model_protT5, model_cat = load_models()
 # Load ProtT5 model
 tokenizer_t5 = T5Tokenizer.from_pretrained("Rostlab/prot_t5_xl_uniref50", do_lower_case=False)
 model_t5 = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_uniref50")
-model_t5 = model_t5.eval()
+model_t5 = model_t5.eval().to("cuda")
 
 # 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)
+esm_model = AutoModel.from_pretrained(model_name).to("cuda")
 
 def extract_prott5_embedding(sequence):
     sequence = sequence.replace(" ", "")
@@ -42,6 +43,15 @@ def extract_esm_embedding(sequence):
     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)
+#    concat = torch.cat((esm_emb, protT5_emb), dim=1)
+#    pred = predict_ensemble(protT5_emb, concat, model_protT5, model_cat)
+#    return "Potential Allergen" if pred.item() == 1 else "Non-Allergen"
+    
+
+@spaces.GPU(duration=120)
 def classify(sequence):
     protT5_emb = extract_prott5_embedding(sequence)
     esm_emb = extract_esm_embedding(sequence)

inference.py

@@ -25,16 +25,20 @@ class NeuralNet(nn.Module):
 def load_models():
     model_protT5 = NeuralNet(1024, 200, 100, 50, 2)
     model_protT5.load_state_dict(torch.load("checkpoints/model17-protT5.pt", map_location=torch.device("cpu")))
-    model_protT5.eval()
+    model_protT5.eval().to("cuda")
 
     model_cat = NeuralNet(2304, 200, 100, 100, 2)
     model_cat.load_state_dict(torch.load("checkpoints/model-esm-protT5-5.pt", map_location=torch.device("cpu")))
-    model_cat.eval()
+    model_cat.eval().to("cuda")
 
     return model_protT5, model_cat
 
 
 def predict_ensemble(X_protT5, X_concat, model_protT5, model_cat, weight1=0.60, weight2=0.30):
+    device = next(model_protT5.parameters()).device
+    X_protT5 = X_protT5.to(device)
+    X_concat = X_concat.to(device)
+    
     with torch.no_grad():
         outputs1 = model_cat(X_concat)
         outputs2 = model_protT5(X_protT5)

requirements.txt

@@ -5,4 +5,6 @@ esm
 fair-esm  # if esm isn't installed via pip
 sentencepiece
 h5py
+spaces
 git+https://github.com/facebookresearch/esm.git
+git+https://github.com/agemagician/ProtTrans.git