Update app.py

app.py

@@ -159,6 +159,7 @@ def compute_pseudo_perplexity(model, tokenizer, protein_seq, binder_seq):
     
     return pseudo_perplexity
 
+# compute pLLDT and iPMT from ESMFOLD model directly, very slow
 def compute_plddt_iptm(protein_seq, binder_seq):
     start = time.time()
     # always the ESMFold model
@@ -193,6 +194,50 @@ def compute_plddt_iptm(protein_seq, binder_seq):
     #print(f'compute_plddt_iptm time: {elapsed:.4f} seconds')
     
     return avg_plddt, ptm
+
+# call API of ESMFOLD to get pLLDT
+def get_plddt(protein_seq, binder_seq):
+    start = time.time()
+    sequence = protein_seq + binder_seq
+    
+    retries = 0
+    pdb_str = None
+    url = "https://api.esmatlas.com/foldSequence/v1/pdb/"
+    while retries < 3 and pdb_str is None:
+        response = requests.post(url, data=sequence, verify=False)
+        pdb_str = response.text
+        if pdb_str == "INTERNAL SERVER ERROR":
+            retries += 1
+            time.sleep(0.1)
+            pdb = None #pdb = str = None
+
+    #save a pdb format file
+    name = sequence[:3] + sequence[-3:]  #combine the firt and last 3 AAs of sequence as a filename.
+    outpath = (
+        Path.cwd() / f"PDB-{name}.pdb")
+    with open(outpath.name, "w") as f:
+        f.write(pdb_str)
+    outpath_str = str(outpath)
+
+    #get pdb column values
+    p = Bio.PDB.PDBParser()
+    structure = p.get_structure('myStructureName', pdb_filename)
+    ids = [a.get_id() for a in structure.get_atoms()]
+    pLDDTs = [a.get_bfactor() for a in structure.get_atoms()]
+
+    print("pLDDTs in line228:",pLDDTs)
+    print("pLDDTs shape:",plDDTs.shape())       
+
+    avg_plddt = pLDDTs.mean()
+    #avg_plddt = plddt[0,:,1].mean().item()
+    #iPTM = ptm
+    #print("170: iPTM:",iPTM)
+    ptm = 0
+    end = time.time()
+    elapsed = end - start
+    print(f'get_plddt time: {elapsed:.4f} seconds')
+    
+    return avg_plddt, ptm
     
 def generate_peptide_for_single_sequence(model, tokenizer, protein_seq, peptide_length = 15, top_k = 3, num_binders = 5, plddt_iptm_yes="no"):
     start = time.time()
@@ -230,8 +275,9 @@ def generate_peptide_for_single_sequence(model, tokenizer, protein_seq, peptide_
         ppl = compute_pseudo_perplexity(model, tokenizer, protein_seq, generated_binder)
 
         # Get PLDDT from ESMFold model
-        if  plddt_iptm_yes=="yes":                              
-            plddt, iptm = compute_plddt_iptm(protein_seq, generated_binder) #too time-consuming
+        if  plddt_iptm_yes=="yes": 
+            plddt, iptm = get_plddt(protein_seq, binder_seq)
+            #plddt, iptm = compute_plddt_iptm(protein_seq, generated_binder) #too time-consuming
         else:
             plddt, iptm = [0, 0]