Update app.py

app.py

@@ -5,35 +5,61 @@ import gradio as gr
 import matplotlib.pyplot as plt
 from transformers import AutoTokenizer, EsmModel
 from sklearn.metrics.pairwise import cosine_similarity
-from Bio.PDB import PDBParser, PDBIO
-from Bio.PDB.StructureBuilder import StructureBuilder
+from Bio.PDB import PDBParser, PDBIO, DSSP
+from Bio.PDB.Polypeptide import PPBuilder
 import tempfile
 import os
 
-#----------------------------------------------------Analysis--------------------------------------------------------------------------------------------------------
+#-------------------------------------------------Analysis---------------------------------------------------------------------------------------------------------------
 # Load ESM-1b model and tokenizer
 model = EsmModel.from_pretrained("facebook/esm1b_t33_650M_UR50S", output_hidden_states=True)
 tokenizer = AutoTokenizer.from_pretrained("facebook/esm1b_t33_650M_UR50S")
 
-def compute_residue_scores(seq):
+# Compute per-residue cosine similarity scores (ASA-aware)
+def compute_asa_filtered_scores(seq, pdb_path):
     inputs = tokenizer(seq, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
-        embedding = outputs.last_hidden_state[0]  # shape (L+2, d)
+        embedding = outputs.last_hidden_state[0]
 
     L = len(seq)
-    embedding = embedding[1:L+1]  # shape (L, d)
+    embedding = embedding[1:L+1]  # Remove CLS/EOS
 
+    # Parse structure and compute DSSP
+    parser = PDBParser(QUIET=True)
+    structure = parser.get_structure("prot", pdb_path)
+    model_struct = next(structure.get_models())
+    dssp = DSSP(model_struct, pdb_path)
+
+    # Extract rASA and match to sequence indices
+    rASA = []
+    for key in list(dssp.keys())[:L]:
+        asa = dssp[key][3]  # absolute ASA
+        max_acc = dssp.residue_max_acc[dssp[key][1]]
+        rASA.append(asa / max_acc if max_acc > 0 else 0.0)
+    rASA = np.array(rASA)
+
+    # Bin into buried (<= 0.25) and exposed (> 0.25)
+    buried_idx = np.where(rASA <= 0.25)[0]
+    exposed_idx = np.where(rASA > 0.25)[0]
+
+    # Compute cosine similarity matrix
     sim_matrix = cosine_similarity(embedding.detach().cpu().numpy())
-    residue_scores = np.sum(sim_matrix, axis=1)
-    norm_scores = 100 * (residue_scores - np.min(residue_scores)) / (np.max(residue_scores) - np.min(residue_scores))
 
+    # Sum similarities only within ASA bins
+    filtered_scores = np.zeros(L)
+    for i in range(L):
+        group = buried_idx if i in buried_idx else exposed_idx
+        filtered_scores[i] = np.sum(sim_matrix[i, group])
+
+    # Normalize
+    norm_scores = 100 * (filtered_scores - np.min(filtered_scores)) / (np.max(filtered_scores) - np.min(filtered_scores))
     return norm_scores
 
+# Inject scores into B-factor column
 def inject_bfactors_into_pdb(pdb_file, scores):
     parser = PDBParser(QUIET=True)
     structure = parser.get_structure("prot", pdb_file.name)
-
     i = 0
     for model in structure:
         for chain in model:
@@ -43,16 +69,15 @@ def inject_bfactors_into_pdb(pdb_file, scores):
                 for atom in residue:
                     atom.bfactor = float(scores[i])
                 i += 1
-
     out_path = tempfile.NamedTemporaryFile(delete=False, suffix=".pdb").name
     io = PDBIO()
     io.set_structure(structure)
     io.save(out_path)
-
     return out_path
 
+# Combined Gradio interface
 def process(seq, pdb_file):
-    scores = compute_residue_scores(seq)
+    scores = compute_asa_filtered_scores(seq, pdb_file.name)
     pdb_with_scores = inject_bfactors_into_pdb(pdb_file, scores)
     return pdb_with_scores
 
@@ -63,8 +88,8 @@ demo = gr.Interface(
         gr.Textbox(label="Input Protein Sequence (1-letter code)"),
         gr.File(label="Upload PDB File", file_types=[".pdb"])
     ],
-    outputs=gr.File(label="Modified PDB with Scores in B-factor Column"),
-    title="ESM-1b Residue Scoring: B-factor Injection for Structural Visualization"
+    outputs=gr.File(label="Modified PDB with ASA-filtered Embedding Scores in B-factor Column"),
+    title="ESM-1b ASA-Aware Residue Scoring for Structural Visualization"
 )
 
 demo.launch()