Update app.py

app.py

@@ -1,10 +1,11 @@
 #-------------------------------------------------------libraries------------------------------------------------------------------------------------
 import torch
-import numpy
+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
 
 #----------------------------------------------------Analysis------------------------------------------------------------------------------------
 #--load model and tokenizer
@@ -19,8 +20,8 @@ print("Transformers version:", transformers.__version__)
 #import torch
 print("Torch NumPy test:", torch.ones(1).numpy())
 
-
-#--task to execute
+'''
+#--principal component plot
 def extract_and_plot(seq, layer=-1):
     #--preprocess sequence
     inputs = tokenizer(seq, return_tensors="pt")
@@ -59,4 +60,71 @@ demo = gr.Interface(
     outputs=gr.Plot()
 )
 
+demo.launch()
+'''
+
+#--hydrophobic classification
+#--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, reuturn_tenors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+        embedding = outputs.last_hidden_state[0]
+
+    #--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 = {
+        "nonpolor": "magenta",
+        "polar": "indigo",
+        "other": "steelblue"
+    }
+    row_colors = [class_colors[c] for c in residues_classes]
+
+    #--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")
+
+    #--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()