Update app.py

app.py

@@ -137,12 +137,10 @@ def visualize_attention(model, feats, drug_idx: Optional[int] = None) -> str:
     """
     Render a Protein → Drug cross-attention heat-map and, optionally, a
     Top-30 protein-residue table for a chosen drug-token index.
-
     The token index shown on the x-axis (and accepted via *drug_idx*) is **the
     position of that token in the *original* drug sequence**, *after* the
     tokeniser but *before* any pruning or truncation (1-based in the labels,
     0-based for the function argument).
-
     Returns
     -------
     html : str
@@ -178,6 +176,7 @@ def visualize_attention(model, feats, drug_idx: Optional[int] = None) -> str:
         d_indices_full = d_indices_full[: attn.size(1)]
         attn           = attn[: len(p_tokens_full), : len(d_tokens_full)]
 
+        orig_attn = attn.clone()
         # ── adaptive sparsity pruning ───────────────────────────────────────────
         thr = attn.max().item() * 0.05
         row_keep = (attn.max(dim=1).values > thr)
@@ -237,7 +236,7 @@ def visualize_attention(model, feats, drug_idx: Optional[int] = None) -> str:
 
         # ───────────────────── Top-30 tabel ─────────────────────
         table_html = ""                  
-        if drug_idx is not None:
+        if drug_idx is not None and 0 <= drug_idx < orig_attn.size(1):
             # map original 0-based drug_idx → current column position
             if (drug_idx + 1) in d_indices:
                 col_pos = d_indices.index(drug_idx + 1)
@@ -254,8 +253,8 @@ def visualize_attention(model, feats, drug_idx: Optional[int] = None) -> str:
                 res_row  = "".join(f"<td>{p_tokens[i]}</td>" for i in topk)
                 pos_row  = "".join(f"<td>{p_indices[i]}</td>"for i in topk)
 
-                drug_tok_text = d_tokens[col_pos]
-                orig_idx      = d_indices[col_pos]
+                drug_tok_text = d_tokens_full[col_pos]
+                orig_idx      = d_indices_full[col_pos]
                 
                 # 1) build the header row: leading “Rank”, then 1…30
                 header_cells = (
@@ -274,7 +273,7 @@ def visualize_attention(model, feats, drug_idx: Optional[int] = None) -> str:
                     "background:#f7f7f7; text-align:center;'>Residue</th>"
                     + "".join(
                         f"<td style='border:1px solid #ccc; padding:6px; "
-                        f"text-align:center'>{p_tokens[i]}</td>"
+                        f"text-align:center'>{p_tokens_full[i]}</td>"
                         for i in topk
                     )
                 )
@@ -285,7 +284,7 @@ def visualize_attention(model, feats, drug_idx: Optional[int] = None) -> str:
                     "background:#f7f7f7; text-align:center;'>Position</th>"
                     + "".join(
                         f"<td style='border:1px solid #ccc; padding:6px; "
-                        f"text-align:center'>{p_indices[i]}</td>"
+                        f"text-align:center'>{p_indices_full[i]}</td>"
                         for i in topk
                     )
                 )
@@ -397,29 +396,24 @@ button, .gr-button { font-family: Inter,sans-serif; font-weight: 600; }
 .link-btn.project{background:linear-gradient(to right,#10b981,#059669);}
 .link-btn.arxiv  {background:linear-gradient(to right,#ef4444,#dc2626);}
 .link-btn.github {background:linear-gradient(to right,#3b82f6,#2563eb);}
-
 /* make *all* gradio buttons a bit taller */
 .gr-button { min-height: 10px !important; }
-
 /* now target just our two big action buttons */
 #extract-btn, #inference-btn {
     width: 5px !important;
     min-height: 36px !important;
     margin-top: 12px !important;
 }
-
 /* and make clear button full width but shorter */
 #clear-btn {
     width: 10px  !important;
     min-height: 36px !important;
     margin-top: 12px !important;
 }
-
 #input-card label {
     font-weight: 600 !important;    /* make the text bold */
     color: var(--text) !important;  /* use your standard text color */
 }
-
 .card {
   background: var(--card);
   border: 1px solid var(--border);
@@ -429,7 +423,6 @@ button, .gr-button { font-family: Inter,sans-serif; font-weight: 600; }
   margin: 0 auto 32px;
   box-shadow: 0 2px 6px rgba(0,0,0,0.05);
 }
-
 #guidelines-card h2 {
   font-size: 1.4rem;
   margin-bottom: 16px;
@@ -453,7 +446,9 @@ button, .gr-button { font-family: Inter,sans-serif; font-weight: 600; }
 
 with gr.Blocks(css=css) as demo:
     # ───────────── Title ─────────────
-    gr.Markdown("<h1>Token-level Visualiser for Drug-Target Interaction</h1>")
+    gr.Markdown(
+    "<h1 style='text-align: center;'>Token-level Visualiser for Drug-Target Interaction</h1>"
+)
 
     # ───────────── Project Links ─────────────
     gr.Markdown("""
@@ -477,22 +472,18 @@ with gr.Blocks(css=css) as demo:
                 based on 3D structures from
                 <a href="https://alphafold.ebi.ac.uk" target="_blank">AlphaFold&nbsp;DB</a> or the
                 <a href="https://www.rcsb.org" target="_blank">Protein Data Bank (PDB)</a>.</li>
-
             <li><strong>If you only have an amino acid sequence or a UniProt ID,</strong>
                 you must first visit the
                 <a href="https://www.rcsb.org" target="_blank">Protein Data Bank (PDB)</a>
                 or <a href="https://alphafold.ebi.ac.uk" target="_blank">AlphaFold&nbsp;DB</a>
                 to search and download the corresponding <code>.cif</code> or <code>.pdb</code> file.</li>
-
             <li><strong>Drug input supports both SELFIES and SMILES:</strong><br>
                 You can enter a SELFIES string directly, or paste a SMILES string.
                 SMILES will be automatically converted to SELFIES using
                 <a href="https://github.com/aspuru-guzik-group/selfies" target="_blank">SELFIES encoder</a>.
                 If conversion fails, a red error message will be displayed.</li>
-
             <li>Optionally enter a <strong>1-based</strong> drug atom or substructure index
-                to highlight the Top-10 interacting protein residues.</li>
-
+                to highlight the Top-30 interacting protein residues.</li>
             <li>After inference, you can use the
                 “Download PDF” link to export a high-resolution vector version.</li>
           </ul>
@@ -565,4 +556,4 @@ with gr.Blocks(css=css) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)