Update app.py

app.py

@@ -244,7 +244,7 @@ def generate_peptide_for_single_sequence(model, tokenizer, protein_seq, peptide_
     return binders_with_ppl_plddt
 
 # Predict peptide binder with finetuned model
-def predict_peptide(base_model_path, finetuned_model_path, input_seqs, peptide_length=15,  num_binders=4, top_k=3, plddt_iptm_yes="no"):
+def predict_peptide(base_model_path, finetuned_model_path, input_seqs, peptide_length=15,  num_binders=4, plddt_iptm_yes="no"):
     # Load the model
     loaded_model = AutoModelForMaskedLM.from_pretrained(finetuned_model_path)   #.to(device) inference use cpu
 
@@ -254,6 +254,8 @@ def predict_peptide(base_model_path, finetuned_model_path, input_seqs, peptide_l
     # Tokenization
     tokenizer = AutoTokenizer.from_pretrained(base_model_path)
 
+    # set top_k mutations for each AA position
+    top_k=3
     if isinstance(input_seqs, str):  # Single sequence
         binders = generate_peptide_for_single_sequence(loaded_model, tokenizer, input_seqs, peptide_length, top_k, num_binders, plddt_iptm_yes)
         results_df = pd.DataFrame(binders, columns=['Binder', 'PPL', 'pLDDT', 'iPTM'])