Update app.py

app.py

@@ -177,7 +177,7 @@ def compute_plddt_iptm(protein_seq, binder_seq):
     
     return avg_plddt, iPTM
     
-def generate_peptide_for_single_sequence(model, tokenizer, protein_seq, peptide_length = 15, top_k = 3, num_binders = 4):
+def generate_peptide_for_single_sequence(model, tokenizer, protein_seq, peptide_length = 15, top_k = 3, num_binders = 5):
     start = time.time()
     
     peptide_length = int(peptide_length)
@@ -211,7 +211,7 @@ def generate_peptide_for_single_sequence(model, tokenizer, protein_seq, peptide_
         ppl_value = compute_pseudo_perplexity(model, tokenizer, protein_seq, generated_binder)
 
         # Get PLDDT from ESMFold model
-        #plddt_value, iPTM_value = compute_plddt_iptm(protein_seq, generated_binder)
+        #plddt_value, iPTM_value = compute_plddt_iptm(protein_seq, generated_binder) #too time-consuming
         plddt_value, iPTM_value = [0, 0]
         
         # Add the generated binder and its PPL to the results list
@@ -254,7 +254,7 @@ def predict_peptide(base_model_path, finetuned_model_path, input_seqs, peptide_l
     
     return results_df, PPC
 
-def predict_peptide_from_file(base_model_path, finetuned_model_path, file_obj, max_peptide_length=15, num_binders=4, top_k=3):
+def predict_peptide_from_file(base_model_path, finetuned_model_path, file_obj, max_peptide_length=15, num_binders=5, top_k=3):
     # Load the model
     loaded_model = AutoModelForMaskedLM.from_pretrained(finetuned_model_path)
 
@@ -275,7 +275,8 @@ def predict_peptide_from_file(base_model_path, finetuned_model_path, file_obj, m
 
         #get metrics for ground truth peptide
         ppl = compute_pseudo_perplexity(loaded_model, tokenizer, protein_seq, peptide_seq)
-        plddt, iptm = compute_plddt_iptm(protein_seq, peptide_seq)
+        #plddt, iptm = compute_plddt_iptm(protein_seq, peptide_seq) #too time-consuming
+        plddt_value, iPTM_value = [0, 0]
 
         results.append([protein_seq, peptide_seq, ppl, plddt, iptm, 1]) # flag 1 for ground truth peptide
 
@@ -339,7 +340,7 @@ with demo:
                 )
                 with gr.Row():
                     peptide_length=gr.Slider(minimum=10, maximum=100, step=1, label="Peptide Maximum Length", value=15)
-                    num_pred_peptides=gr.Slider(minimum=1, maximum=10, step=1, label="Number of Predicted Peptides", value=4)
+                    num_pred_peptides=gr.Slider(minimum=1, maximum=10, step=1, label="Number of Predicted Peptides", value=5)
             with gr.Column(scale=5, variant="compact"): 
                     name = gr.Dropdown(
                         label="Choose a Sample Protein",