Update app.py

app.py

@@ -63,53 +63,63 @@ demo = gr.Interface(
 demo.launch()
 '''
 
-#--hydrophobic classification
-#--define hydrophobicity classification
+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]
+    return ["nonpolar" if aa in nonpolar else "polar" if aa in polar else "other" for aa in seq]
 
 def compute_cosine_heatmap(seq):
-    #--tokenize
+    # Tokenize
     inputs = tokenizer(seq, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
-        embedding = outputs.last_hidden_state[0]
+        embedding = outputs.last_hidden_state[0]  # shape (L, 1280)
 
-    #--remove [CLS] and [EOS] if present
+    # Remove [CLS] and [EOS] if present
     L = len(seq)
     embedding = embedding[1:L+1]
 
-    #--cosine similarity matrix
+    # Cosine similarity matrix
     sim_matrix = cosine_similarity(embedding.detach().cpu().numpy())
 
-    #--residue classification
+    # Residue classification
     residue_classes = classify_residues(seq)
     class_colors = {
-        "nonpolor": "magenta",
-        "polar": "indigo",
-        "other": "steelblue"
+        "nonpolar": "magenta",  
+        "polar": "indigo",     
+        "other": "steelblue"      
     }
     row_colors = [class_colors[c] for c in residue_classes]
 
-    #--plot heatmap
+    # Plot heatmap
     fig, ax = plt.subplots(figsize=(8, 6))
     im = ax.imshow(sim_matrix, cmap="viridis")
-    fig.colors(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")
+    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
+    # 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
+    # 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):
@@ -120,7 +130,7 @@ def compute_cosine_heatmap(seq):
     fig.tight_layout()
     return fig
 
-#--Gradio UI
+# Gradio UI
 demo = gr.Interface(
     fn=compute_cosine_heatmap,
     inputs=gr.Textbox(label="Input Protein Sequence (1-letter code)"),