Update app.py

app.py

@@ -1,140 +1,70 @@
-#-------------------------------------------------------libraries------------------------------------------------------------------------------------
+#-------------------------------------------------libraries---------------------------------------------------------------------------------------------------------------
 import torch
 import numpy as np
 import gradio as gr
 import matplotlib.pyplot as plt
 from transformers import AutoTokenizer, EsmModel
-from sklearn.decomposition import PCA
 from sklearn.metrics.pairwise import cosine_similarity
+from Bio.PDB import PDBParser, PDBIO
+from Bio.PDB.StructureBuilder import StructureBuilder
+import tempfile
+import os
 
-#----------------------------------------------------Analysis------------------------------------------------------------------------------------
-#--load model and tokenizer
+#----------------------------------------------------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")
 
-#--confirm proper installation
-import torch, transformers
-print("Torch version:", torch.__version__)
-print("Transformers version:", transformers.__version__)
-
-#import torch
-print("Torch NumPy test:", torch.ones(1).numpy())
-
-'''
-#--principal component plot
-def extract_and_plot(seq, layer=-1):
-    #--preprocess sequence
+def compute_residue_scores(seq):
     inputs = tokenizer(seq, return_tensors="pt")
-
-    #--forward pass
     with torch.no_grad():
         outputs = model(**inputs)
-        hidden_states = outputs.hidden_states   #--> tuple: (layer0, ..., layer_final)
-
-    #--select hidden state from specified layer
-    if layer == 1:
-        embedding = hidden_states[-1][0]    #--> (seq_len, hidden_dim)
-    else:
-        embedding = hidden_states[layer][0]
-
-    #--PCA
-    pca = PCA(n_components=2)
-    coords = pca.fit_transform(embedding.numpy())
+        embedding = outputs.last_hidden_state[0]  # shape (L+2, d)
 
-    #--plot
-    plt.figure(figsize=(6, 4))
-    plt.scatter(coords[:, 0], coords[:, 1])
-    plt.title(f"PCA of esm1b embeddings (layer {layer})")
-    plt.xlabel("PCA1")
-    plt.ylabel("PCA2")
-    plt.tight_layout()
-
-    return plt
+    L = len(seq)
+    embedding = embedding[1:L+1]  # shape (L, d)
 
+    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))
+
+    return norm_scores
+
+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:
+            for residue in chain:
+                if i >= len(scores):
+                    break
+                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
+
+def process(seq, pdb_file):
+    scores = compute_residue_scores(seq)
+    pdb_with_scores = inject_bfactors_into_pdb(pdb_file, scores)
+    return pdb_with_scores
+
+# Gradio Interface
 demo = gr.Interface(
-    fn=extract_and_plot,
+    fn=process,
     inputs=[
-        gr.Textbox(label="Protein Sequence"),
-        gr.Slider(minimum=0, maximum=33, step=1, value=33, label="Layer (-1 = final)")
+        gr.Textbox(label="Input Protein Sequence (1-letter code)"),
+        gr.File(label="Upload PDB File", file_types=[".pdb"])
     ],
-    outputs=gr.Plot()
+    outputs=gr.File(label="Modified PDB with Scores in B-factor Column"),
+    title="ESM-1b Residue Scoring: B-factor Injection for Structural Visualization"
 )
 
 demo.launch()
-'''
-
-import torch
-import gradio as gr
-import matplotlib.pyplot as plt
-import numpy as np
-from sklearn.metrics.pairwise import cosine_similarity
-from transformers import AutoTokenizer, EsmModel
-
-# Load model
-model = EsmModel.from_pretrained("facebook/esm1b_t33_650M_UR50S", output_hidden_states=True)
-tokenizer = AutoTokenizer.from_pretrained("facebook/esm1b_t33_650M_UR50S")
-
-# Define hydrophobicity classification
-nonpolar = set("AFLIVMYW")
-polar = set("QERSDHKNT")
-
-def classify_residues(seq):
-    return ["nonpolar" if aa in nonpolar else "polar" if aa in polar else "other" for aa in seq]
-
-def compute_cosine_heatmap(seq):
-    # Tokenize
-    inputs = tokenizer(seq, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-        embedding = outputs.last_hidden_state[0]  # shape (L, 1280)
-
-    # Remove [CLS] and [EOS] if present
-    L = len(seq)
-    embedding = embedding[1:L+1]
-
-    # Cosine similarity matrix
-    sim_matrix = cosine_similarity(embedding.detach().cpu().numpy())
-
-    # Residue classification
-    residue_classes = classify_residues(seq)
-    class_colors = {
-        "nonpolar": "magenta",  
-        "polar": "indigo",     
-        "other": "steelblue"      
-    }
-    row_colors = [class_colors[c] for c in residue_classes]
-
-    # Plot heatmap
-    fig, ax = plt.subplots(figsize=(8, 6))
-    im = ax.imshow(sim_matrix, cmap="viridis")
-    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
-    ax.set_title("Residue–Residue Cosine Similarity")
-    ax.set_xlabel("Residue Index")
-    ax.set_ylabel("Residue Index")
-
-    # Add colored ticks for class annotation
-    for spine in ax.spines.values():
-        spine.set_visible(False)
-    ax.set_xticks(range(L))
-    ax.set_yticks(range(L))
-    ax.tick_params(length=0)
-
-    # Color-code labels
-    ax.set_xticklabels(residue_classes, rotation=90, fontsize=6)
-    ax.set_yticklabels(residue_classes, fontsize=6)
-    for label, color in zip(ax.get_xticklabels(), row_colors):
-        label.set_color(color)
-    for label, color in zip(ax.get_yticklabels(), row_colors):
-        label.set_color(color)
-
-    fig.tight_layout()
-    return fig
-
-# Gradio UI
-demo = gr.Interface(
-    fn=compute_cosine_heatmap,
-    inputs=gr.Textbox(label="Input Protein Sequence (1-letter code)"),
-    outputs=gr.Plot()
-)
-
-demo.launch()