Update app.py

app.py

@@ -218,15 +218,15 @@ def predict_peptide(base_model_path, finetuned_model_path, input_seqs, peptide_l
 
     if isinstance(input_seqs, str):  # Single sequence
         binders = generate_peptide_for_single_sequence(loaded_model, tokenizer, input_seqs, peptide_length, top_k, num_binders)
-        results_df = pd.DataFrame(binders, columns=['Binder', 'Pseudo Perplexity'])
+        results_df = pd.DataFrame(binders, columns=['Binder', 'Pseudo Perplexity', 'pLDDT'])
 
     elif isinstance(input_seqs, list):  # List of sequences
         results = []
         for seq in input_seqs:
             binders = generate_peptide_for_single_sequence(loaded_model, tokenizer, seq, peptide_length, top_k, num_binders)
-            for binder, ppl in binders:
-                results.append([seq, binder, ppl])
-        results_df = pd.DataFrame(results, columns=['Input Sequence', 'Binder', 'Pseudo Perplexity'])
+            for binder, ppl, plddt in binders:
+                results.append([seq, binder, ppl, plddt])
+        results_df = pd.DataFrame(results, columns=['Input Sequence', 'Binder', 'Pseudo Perplexity', 'pLDDT'])
     print(results_df)
 
     #combine target protein and predicted peptide with 20 G amino acids.
@@ -254,9 +254,9 @@ def predict_peptide_from_file(base_model_path, finetuned_model_path, file_obj, p
         seq = row['Receptor Sequence']
         binders = generate_peptide_for_single_sequence(loaded_model, tokenizer, seq, peptide_length, top_k, num_binders)
         for binder, ppl in binders:
-            results.append([seq, binder, ppl])
+            results.append([seq, binder, ppl, plddt])
         print("Predicted sample: binder, ppl", i, binder, ppl )
-    results_df = pd.DataFrame(results, columns=['Input Sequence', 'Binder', 'Pseudo Perplexity'])
+    results_df = pd.DataFrame(results, columns=['Input Sequence', 'Binder', 'Pseudo Perplexity', 'pLDDT'])
 
     outpath = (
         Path.cwd() / "predicted_peptides.csv"