Make webfiles compatible with ZeroGPU processor

app.py

@@ -13,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().to("cuda")
+model_t5 = model_t5.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).to("cuda")
+esm_model = AutoModel.from_pretrained(model_name)
 
 def extract_prott5_embedding(sequence):
     sequence = sequence.replace(" ", "")
@@ -42,14 +42,12 @@ def extract_esm_embedding(sequence):
     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)
-#    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"
-    
+def estimate_duration(sequence):
+    # Estimate duration based on sequence length
+    base_time = 30  # Base time in seconds
+    time_per_residue = 0.5  # Estimated time per residue
+    estimated_time = base_time + len(sequence) * time_per_residue
+    return min(int(estimated_time), 300)  # Cap at 300 seconds
 
 @spaces.GPU(duration=120)
 def classify(sequence):
@@ -65,4 +63,4 @@ demo = gr.Interface(fn=classify,
                     outputs=gr.Label(label="Prediction"))
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

requirements.txt

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