Update app.py

app.py

@@ -2,7 +2,7 @@
 import gradio as gr
 
 import os
-from transformers import Trainer, TrainingArguments, AutoTokenizer, EsmForMaskedLM, AutoModelForMaskedLM, TrainerCallback
+from transformers import Trainer, TrainingArguments, AutoTokenizer, EsmForMaskedLM, AutoModelForMaskedLM, TrainerCallback, EsmForProteinFolding
 import torch
 from torch.utils.data import DataLoader, Dataset, RandomSampler
 from torch.optim import AdamW
@@ -141,8 +141,27 @@ def compute_pseudo_perplexity(model, tokenizer, protein_seq, binder_seq):
     # Loss is already averaged by the model
     avg_loss = loss.item()
     pseudo_perplexity = np.exp(avg_loss)
+    
     return pseudo_perplexity
 
+def compute_avg_plddt(protein_seq, binder_seq):
+    # always the ESMFold model
+    model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1")
+    tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
+    
+    sequence = protein_seq + binder_seq
+    inputs = tokenizer(sequence, return_tensors='pt', add_special_tokens=False)
+
+    # Get model predictions 
+    with torch.no_grad():
+        outputs = model(**inputs)
+        plddt = outputs.plddt
+        print("159: plddt:",plddt)
+
+    avg_plddt = plddt.mean()
+    
+    return avg_plddt
+    
 def generate_peptide_for_single_sequence(model, tokenizer, protein_seq, peptide_length = 15, top_k = 3, num_binders = 4):
 
     peptide_length = int(peptide_length)
@@ -171,8 +190,13 @@ def generate_peptide_for_single_sequence(model, tokenizer, protein_seq, peptide_
         predicted_indices = Categorical(probabilities).sample()
         predicted_token_ids = top_k_indices.gather(-1, predicted_indices.unsqueeze(-1)).squeeze(-1)
         generated_binder = tokenizer.decode(predicted_token_ids, skip_special_tokens=True).replace(' ', '')
+        
         # Compute PPL for the generated binder
         ppl_value = compute_pseudo_perplexity(model, tokenizer, protein_seq, generated_binder)
+
+        # Get PLDDT from ESMFold model
+        plddt_value = compute_avg_plddt(protein_seq, generated_binder)
+        
         
         # Add the generated binder and its PPL to the results list
         binders_with_ppl.append([generated_binder, ppl_value])