Update app.py

app.py

@@ -2,8 +2,8 @@
 import gradio_client.utils as _gc_utils
 
 # back up originals
-_orig_get_type              = _gc_utils.get_type
-_orig_json2py              = _gc_utils._json_schema_to_python_type
+_orig_get_type   = _gc_utils.get_type
+_orig_json2py    = _gc_utils._json_schema_to_python_type
 
 def _patched_get_type(schema):
     # treat any boolean schema as if it were an empty dict
@@ -54,6 +54,7 @@ from utils.foldseek_util import get_struc_seq
 
 three2one = {k.upper(): v for k, v in IUPACData.protein_letters_3to1.items()}
 three2one.update({"MSE": "M", "SEC": "C", "PYL": "K"})
+
 def simple_seq_from_structure(path: str) -> str:
     parser = MMCIFParser(QUIET=True) if path.endswith(".cif") else PDBParser(QUIET=True)
     structure = parser.get_structure("P", path)
@@ -69,7 +70,7 @@ def smiles_to_selfies(smiles: str) -> Optional[str]:
         if mol is None:
             return None
         return selfies.encoder(smiles)
-    except:
+    except Exception:
         return None
 
 def parse_config():
@@ -132,20 +133,23 @@ def get_case_feature(model, loader):
                      p_ids.cpu(), d_ids.cpu(),
                      p_mask.cpu(), d_mask.cpu(), None)]
 
-        
 # ─────────────── visualisation ───────────────────────────────────────────
+def _safe_is_special(tokenizer, tok: str) -> bool:
+    # Some tokenisers expose different special token sets; fall back conservatively.
+    special_sets = []
+    if hasattr(tokenizer, "all_special_tokens"):
+        special_sets.append(set(tokenizer.all_special_tokens))
+    if hasattr(tokenizer, "special_tokens_map"):
+        special_sets.extend(set(v) if isinstance(v, list) else {v}
+                            for v in tokenizer.special_tokens_map.values())
+    for s in special_sets:
+        if tok in s:
+            return True
+    return False
+
 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
-        Base64-embedded PNG heat-map (+ optional HTML table).
+    Render a Protein → Drug cross-attention heat-map and optional Top-30 residue table.
     """
     model.eval()
     with torch.no_grad():
@@ -156,154 +160,75 @@ def visualize_attention(model, feats, drug_idx: Optional[int] = None) -> str:
 
         # ── forward pass: Protein → Drug attention (B, n_p, n_d) ───────────────
         _, att_pd = model(p_emb, d_emb, p_mask, d_mask)
-        attn = att_pd.squeeze(0).cpu()                                  # (n_p, n_d)
+        attn = att_pd.squeeze(0).cpu()  # (n_p, n_d)
 
         # ── decode tokens (skip special symbols) ────────────────────────────────
         def clean_ids(ids, tokenizer):
             toks = tokenizer.convert_ids_to_tokens(ids.tolist())
-            return [t for t in toks if t not in tokenizer.all_special_tokens]
+            return [t for t in toks if not _safe_is_special(tokenizer, t)]
 
-        # ── decode full sequences + record 1-based indices ──────────────────
         p_tokens_full  = clean_ids(p_ids[0],  prot_tokenizer)
         p_indices_full = list(range(1, len(p_tokens_full)  + 1))
-
         d_tokens_full  = clean_ids(d_ids[0],  drug_tokenizer)
         d_indices_full = list(range(1, len(d_tokens_full)  + 1))
 
-        # ── safety cut-off to match attn mat size ───────────────────────────────
-        p_tokens       = p_tokens_full[: attn.size(0)]
-        p_indices_full = p_indices_full[: attn.size(0)]
-        d_tokens_full  = d_tokens_full[: attn.size(1)]
-        d_indices_full = d_indices_full[: attn.size(1)]
-        attn           = attn[: len(p_tokens_full), : len(d_tokens_full)]
+        # ── safety cut-off to match attn mat size ──────────────────────────────
+        p_tokens = p_tokens_full[: attn.size(0)]
+        p_indices = p_indices_full[: attn.size(0)]
+        d_tokens = d_tokens_full[: attn.size(1)]
+        d_indices = d_indices_full[: attn.size(1)]
+        attn = attn[: len(p_tokens), : len(d_tokens)]
 
         orig_attn = attn.clone()
+
         # ── adaptive sparsity pruning ───────────────────────────────────────────
-        thr = attn.max().item() * 0.05
-        row_keep = (attn.max(dim=1).values > thr)
-        col_keep = (attn.max(dim=0).values > thr)
+        thr = attn.max().item() * 0.05 if attn.numel() > 0 else 0.0
+        row_keep = (attn.max(dim=1).values > thr) if attn.size(0) else torch.tensor([], dtype=torch.bool)
+        col_keep = (attn.max(dim=0).values > thr) if attn.size(1) else torch.tensor([], dtype=torch.bool)
 
-        if row_keep.sum() < 3:
-            row_keep[:] = True
-        if col_keep.sum() < 3:
-            col_keep[:] = True
+        if row_keep.sum().item() < 3 and attn.size(0) > 0:
+            row_keep = torch.ones(attn.size(0), dtype=torch.bool)
+        if col_keep.sum().item() < 3 and attn.size(1) > 0:
+            col_keep = torch.ones(attn.size(1), dtype=torch.bool)
 
-        attn       = attn[row_keep][:, col_keep]
-        p_tokens   = [tok for keep, tok in zip(row_keep, p_tokens)        if keep]
-        p_indices  = [idx for keep, idx in zip(row_keep, p_indices_full)  if keep]
-        d_tokens   = [tok for keep, tok in zip(col_keep, d_tokens_full)   if keep]
-        d_indices  = [idx for keep, idx in zip(col_keep, d_indices_full)  if keep]
+        attn      = attn[row_keep][:, col_keep]
+        p_tokens  = [tok for keep, tok in zip(row_keep.tolist(), p_tokens)       if keep]
+        p_indices = [idx for keep, idx in zip(row_keep.tolist(), p_indices)      if keep]
+        d_tokens  = [tok for keep, tok in zip(col_keep.tolist(), d_tokens)       if keep]
+        d_indices = [idx for keep, idx in zip(col_keep.tolist(), d_indices)      if keep]
 
         # ── cap column count at 150 for readability ─────────────────────────────
         if attn.size(1) > 150:
-            topc       = torch.topk(attn.sum(0), k=150).indices
-            attn       = attn[:, topc]
-            d_tokens   = [d_tokens [i] for i in topc]
-            d_indices  = [d_indices[i] for i in topc]
+            topc = torch.topk(attn.sum(0), k=150).indices
+            attn = attn[:, topc]
+            d_tokens  = [d_tokens[i]  for i in topc]
+            d_indices = [d_indices[i] for i in topc]
 
-        # ── draw heat-map ───────────────────────────────────────────────────────
+        # ── draw heat-map ──────────────────────────────────────────────────────
         x_labels = [f"{idx}:{tok}" for idx, tok in zip(d_indices, d_tokens)]
         y_labels = [f"{idx}:{tok}" for idx, tok in zip(p_indices, p_tokens)]
 
-
-        fig_w = min(22, max(8, len(x_labels) * 0.6))    # ~0.6″ per column
-        fig_h = min(24, max(6, len(p_tokens) * 0.8))  
+        fig_w = min(22, max(8, len(x_labels) * 0.6))
+        fig_h = min(24, max(6, len(y_labels) * 0.8))
 
         fig, ax = plt.subplots(figsize=(fig_w, fig_h))
-        im = ax.imshow(attn.numpy(), aspect="auto",
-                       cmap=cm.viridis, interpolation="nearest")
-
-        ax.set_title("Protein → Drug Attention", pad=8, fontsize=10)
+        im = ax.imshow(attn.numpy(), aspect="auto", cmap=cm.viridis, interpolation="nearest")
 
+        ax.set_title("Protein → Drug Attention", pad=8, fontsize=11)
         ax.set_xticks(range(len(x_labels)))
-        ax.set_xticklabels(x_labels, rotation=90, fontsize=8,
-                           ha="center", va="center")
-        ax.tick_params(axis="x", top=True, bottom=False,
-                       labeltop=True, labelbottom=False, pad=27)
+        ax.set_xticklabels(x_labels, rotation=90, fontsize=8, ha="center", va="center")
+        ax.tick_params(axis="x", top=True, bottom=False, labeltop=True, labelbottom=False, pad=27)
 
         ax.set_yticks(range(len(y_labels)))
         ax.set_yticklabels(y_labels, fontsize=7)
-        ax.tick_params(axis="y", top=True, bottom=False,
-                    labeltop=True, labelbottom=False,
-                    pad=10)
+        ax.tick_params(axis="y", top=True, bottom=False, labeltop=True, labelbottom=False, pad=10)
 
         fig.colorbar(im, fraction=0.026, pad=0.01)
         fig.tight_layout()
 
-        buf = io.BytesIO()
-        fig.savefig(buf, format="png", dpi=140)
-        plt.close(fig)
-        html = f'<img src="data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}" />'
-
-        # ───────────────────── Top-30 tabel ─────────────────────
-        table_html = ""                  
-        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)
-            elif 0 <= drug_idx < len(d_tokens):
-                col_pos = drug_idx
-            else:
-                col_pos = None
-
-            if col_pos is not None:
-                col_vec = attn[:, col_pos]
-                topk    = torch.topk(col_vec, k=min(30, len(col_vec))).indices.tolist()
-
-                rank_hdr = "".join(f"<th>{r+1}</th>"         for r in range(len(topk)))
-                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_full[col_pos]
-                orig_idx      = d_indices_full[col_pos]
-                
-                # 1) build the header row: leading “Rank”, then 1…30
-                header_cells = (
-                    "<th style='border:1px solid #ccc; padding:6px; "
-                    "background:#f7f7f7; text-align:center;'>Rank</th>"
-                    + "".join(
-                        f"<th style='border:1px solid #ccc; padding:6px; "
-                        f"background:#f7f7f7; text-align:center'>{r+1}</th>"
-                        for r in range(len(topk))
-                    )
-                )
-
-                # 2) build the residue row: leading “Residue”, then the residue tokens
-                residue_cells = (
-                    "<th style='border:1px solid #ccc; padding:6px; "
-                    "background:#f7f7f7; text-align:center;'>Residue</th>"
-                    + "".join(
-                        f"<td style='border:1px solid #ccc; padding:6px; "
-                        f"text-align:center'>{p_tokens_full[i]}</td>"
-                        for i in topk
-                    )
-                )
-
-                # 3) build the position row: leading “Position”, then the residue positions
-                position_cells = (
-                    "<th style='border:1px solid #ccc; padding:6px; "
-                    "background:#f7f7f7; text-align:center;'>Position</th>"
-                    + "".join(
-                        f"<td style='border:1px solid #ccc; padding:6px; "
-                        f"text-align:center'>{p_indices_full[i]}</td>"
-                        for i in topk
-                    )
-                )
-
-                # 4) assemble your table_html
-                table_html = (
-                    f"<h4 style='margin-bottom:12px'>"
-                      f"Drug atom #{orig_idx} <code>{drug_tok_text}</code> → Top-30 Protein residues"
-                    f"</h4>"
-                    f"<table style='border-collapse:collapse; margin:0 auto 24px;'>"
-                      f"<tr>{header_cells}</tr>"
-                      f"<tr>{residue_cells}</tr>"
-                      f"<tr>{position_cells}</tr>"
-                    f"</table>"
-                )
-
+        # build PNG / PDF
         buf_png = io.BytesIO()
-        fig.savefig(buf_png, format="png", dpi=140)   
+        fig.savefig(buf_png, format="png", dpi=140)
         buf_png.seek(0)
 
         buf_pdf = io.BytesIO()
@@ -315,24 +240,73 @@ def visualize_attention(model, feats, drug_idx: Optional[int] = None) -> str:
         pdf_b64 = base64.b64encode(buf_pdf.getvalue()).decode()
 
         html_heat = (
-            f"<div style='position: relative; width: 100%;'>"
-              # the PDF button, absolutely positioned
+            f"<div class='heatmap-card' style='position: relative; width: 100%;'>"
               f"<a href='data:application/pdf;base64,{pdf_b64}' download='attention_heatmap.pdf' "
-                 "style='position: absolute; top: 12px; right: 12px; "
-                        "background: var(--primary); color: #fff; "
-                        "padding: 8px 16px; border-radius: 6px; "
-                        "font-size: 0.9rem; font-weight: 500; "
-                        "text-decoration: none;'>"
-                "Download PDF"
-              "</a>"
-              # the clickable heat‐map image
+                 "style='position:absolute; top:12px; right:12px; "
+                 "background: var(--primary); color:#fff; padding:8px 16px; border-radius:8px; "
+                 "font-size:.92rem; font-weight:600; text-decoration:none;'>Download PDF</a>"
               f"<a href='data:image/png;base64,{png_b64}' target='_blank' title='Click to enlarge'>"
                 f"<img src='data:image/png;base64,{png_b64}' "
-                     "style='display: block; width: 100%; height: auto; cursor: zoom-in;'/>"
+                   "style='display:block; width:100%; height:auto; cursor:zoom-in;'/>"
               "</a>"
             "</div>"
         )
 
+        # ───────────────────── Top-30 table (optional) ─────────────────────
+        table_html = ""
+        if drug_idx is not None and orig_attn.size(1) > 0 and 0 <= drug_idx < orig_attn.size(1):
+            # map original 0-based drug_idx → pruned column
+            col_pos = None
+            if (drug_idx + 1) in d_indices:
+                col_pos = d_indices.index(drug_idx + 1)
+            elif 0 <= drug_idx < len(d_tokens):
+                col_pos = drug_idx
+
+            if col_pos is not None:
+                col_vec = attn[:, col_pos]
+                k = min(30, len(col_vec))
+                if k > 0:
+                    topk = torch.topk(col_vec, k=k).indices.tolist()
+
+                    # header cells
+                    header_cells = (
+                        "<th style='border:1px solid #e5e7eb; padding:6px; background:#f8fafc; text-align:center;'>Rank</th>"
+                        + "".join(
+                            f"<th style='border:1px solid #e5e7eb; padding:6px; background:#f8fafc; text-align:center'>{r+1}</th>"
+                            for r in range(len(topk))
+                        )
+                    )
+                    residue_cells = (
+                        "<th style='border:1px solid #e5e7eb; padding:6px; background:#f8fafc; text-align:center;'>Residue</th>"
+                        + "".join(
+                            f"<td style='border:1px solid #e5e7eb; padding:6px; text-align:center'>{p_tokens[i]}</td>"
+                            for i in topk
+                        )
+                    )
+                    position_cells = (
+                        "<th style='border:1px solid #e5e7eb; padding:6px; background:#f8fafc; text-align:center;'>Position</th>"
+                        + "".join(
+                            f"<td style='border:1px solid #e5e7eb; padding:6px; text-align:center'>{p_indices[i]}</td>"
+                            for i in topk
+                        )
+                    )
+
+                    drug_tok_text = d_tokens[col_pos]
+                    orig_idx_disp = d_indices[col_pos]
+
+                    table_html = (
+                        f"<div class='card' style='margin-top:18px'>"
+                        f"<h4 style='margin:0 0 12px; font-size:1rem;'>"
+                        f"Drug atom #{orig_idx_disp} <code>{drug_tok_text}</code> → Top-30 Protein residues"
+                        f"</h4>"
+                        f"<table style='border-collapse:collapse; margin:0 auto 4px; font-size:.95rem'>"
+                        f"<tr>{header_cells}</tr>"
+                        f"<tr>{residue_cells}</tr>"
+                        f"<tr>{position_cells}</tr>"
+                        f"</table>"
+                        f"</div>"
+                    )
+
         return table_html + html_heat
 
 # ───── Gradio Callbacks ─────────────────────────────────────────
@@ -367,116 +341,127 @@ def inference_cb(prot_seq, drug_seq, atom_idx):
     return visualize_attention(model, feats, int(atom_idx)-1 if atom_idx else None)
 
 def clear_cb():
-    return None, "", "", None, ""
+    return "", "", None, "", None
 
-# ───── Gradio Interface Definition ───────────────────────────────
+# ───── Theme & CSS ─────────────────────────────────────────────
 
 css = """
 :root {
-  --bg: #f3f4f6;
-  --card: #ffffff;
-  --border: #e5e7eb;
-  --primary: #6366f1;
-  --primary-dark: #4f46e5;
-  --text: #111827;
+  --bg:#f7f7fb;
+  --card:#ffffff;
+  --border:#e6e7eb;
+  --primary:#4f46e5;
+  --primary-dark:#4338ca;
+  --text:#0f172a;
+  --muted:#6b7280;
+  --radius:14px;
+  --shadow:0 10px 30px rgba(15,23,42,.06);
 }
-* { box-sizing: border-box; margin: 0; padding: 0; }
-body { background: var(--bg); color: var(--text); font-family: Inter,system-ui,Arial,sans-serif; }
-h1 { font-family: Poppins,Inter,sans-serif; font-weight: 600; font-size: 2rem; text-align: center; margin: 24px 0; }
-button, .gr-button { font-family: Inter,sans-serif; font-weight: 600; }
-#project-links { text-align: center; margin-bottom: 32px; }
-#project-links .gr-button { margin: 0 8px; min-width: 160px; }
-#project-links .gr-button:nth-child(1) { background: #10b981; }
-#project-links .gr-button:nth-child(2) { background: #ef4444; }
-#project-links .gr-button:nth-child(3) { background: #3b82f6; }
-#project-links .gr-button:hover { opacity: 0.9; }
-.link-btn{display:inline-block;margin:0 8px;padding:10px 20px;border-radius:8px;
-         color:white;font-weight:600;text-decoration:none;box-shadow:0 2px 6px rgba(0,0,0,0.12);
-         transition:all .2s ease-in-out;}
-.link-btn:hover{opacity:.9;}
-.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;
+*{box-sizing:border-box}
+html,body{background:var(--bg)!important;color:var(--text)!important;font-family:Inter,system-ui,Arial,sans-serif}
+h1{font-weight:700;font-size:32px;margin:22px 0 10px;text-align:center;letter-spacing:.2px}
+p,li,button,.gr-button,label,.gr-text{font-size:14px}
+
+/* Cards */
+.card{
+  background:var(--card); border:1px solid var(--border); border-radius:var(--radius);
+  box-shadow:var(--shadow); padding:24px; max-width:1100px; margin:0 auto 28px;
 }
-/* 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);
-  border-radius: 12px;
-  padding: 24px;
-  max-width: 1000px;
-  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;
-  text-align: center;
+
+/* Project links */
+.link-btn{
+  display:inline-flex;               /* icon + text centred vertically */
+  align-items:center;
+  justify-content:center;
+  margin:0 8px;
+  padding:10px 18px;
+  border-radius:10px;
+  color:#fff;
+  font-weight:650;
+  text-decoration:none;
+  box-shadow:0 6px 18px rgba(79,70,229,.18);
+  transition:transform .12s ease,filter .12s ease;
 }
-#guidelines-card ol {
-  margin-left: 20px;
-  line-height: 1.6;
-  font-size: 1rem;
+.link-btn:hover{transform:translateY(-1px);filter:brightness(1.03)}
+.link-btn svg{margin-right:6px;vertical-align:middle}
+.link-btn.project{background:linear-gradient(135deg,#10b981,#059669)}
+.link-btn.arxiv  {background:linear-gradient(135deg,#ef4444,#dc2626)}
+.link-btn.github {background:linear-gradient(135deg,#3b82f6,#2563eb)}
+
+/* Labels & inputs */
+#input-card label{font-weight:650!important;color:var(--text)!important}
+textarea, input, .gr-textbox, .gr-number{
+  border-radius:12px!important; border:1px solid var(--border)!important;
 }
-#input-card .gr-row, #input-card .gr-cols {
-  gap: 16px;
+#input-card .gr-row, #input-card .gr-cols{gap:16px}
+
+/* Buttons */
+.gr-button{min-height:42px!important; padding:0 18px!important; border-radius:12px!important; font-weight:700!important}
+.gr-button.primary, .gr-button-primary{
+  background:var(--primary)!important; border-color:var(--primary)!important; color:#fff!important
 }
-#input-card .gr-button {
-  flex: 1;
+.gr-button.primary:hover, .gr-button-primary:hover{background:var(--primary-dark)!important;border-color:var(--primary-dark)!important}
+
+/* Action buttons row */
+#action-buttons{gap:12px}
+#extract-btn, #inference-btn{flex:1 1 260px!important; min-width:180px!important}
+#clear-btn{width:100%!important}
+
+/* Output */
+#output-card{padding-top:0}
+#result-html{padding:0; margin:0}
+#result-html .heatmap-card{
+  background:var(--card); border:1px solid var(--border); border-radius:12px; padding:12px; box-shadow:var(--shadow)
 }
-#output-card {
-  padding-top: 0;
+
+/* Guidance */
+#guidelines-card h2{font-size:18px;margin-bottom:14px;text-align:center}
+#guidelines-card ul{margin-left:18px;line-height:1.6}
+
+/* Small screens */
+@media (max-width: 900px){
+  .card{margin:0 12px 24px}
 }
 """
 
-# ───── Gradio Interface ─────────────────────────────────────
+# ───── Gradio Interface Definition ───────────────────────────────
 with gr.Blocks(theme=Soft(primary_hue="indigo", neutral_hue="slate"), css=css) as demo:
-    gr.Markdown("<h1>Token-level Visualiser for Drug-Target Interaction</h1>")
+    # ───────────── Title ─────────────
+    gr.Markdown("<h1 style='text-align: center;'>Token-level Visualiser for Drug-Target Interaction</h1>")
 
-    # Project links with SVG icons
+    # ───────────── Project Links (SVG icons) ─────────────
     gr.HTML("""
-    <div style="text-align:center;margin-bottom:32px;">
-      <a class="link-btn project" href="https://zhaohanm.github.io/FusionDTI.github.io/" target="_blank">
-        <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="currentColor"
-             viewBox="0 0 16 16"><path d="M8 0a8 8 0 1 0 8 8A8.009 8.009 0 0 0 8 0ZM4.5 8a3.5 3.5 0 1 1 3.5 3.5A3.504 3.504 0 0 1 4.5 8Z"/></svg>
-        Project Page
-      </a>
-      <a class="link-btn arxiv" href="https://arxiv.org/abs/2406.01651" target="_blank">
-        <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="currentColor"
-             viewBox="0 0 16 16"><path d="M4 1.5a.5.5 0 0 0-.5.5V3h9V2a.5.5 0 0 0-.5-.5h-8ZM2 4v9a2 2 0 0 0 2 2h8a2 2 0 0 0 2-2V4H2Z"/></svg>
-        ArXiv: 2406.01651
-      </a>
-      <a class="link-btn github" href="https://github.com/ZhaohanM/FusionDTI" target="_blank">
-        <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="currentColor"
-             viewBox="0 0 16 16"><path d="M8 .198a8 8 0 0 0-2.53 15.598c.4.074.547-.174.547-.385v-1.352c-2.23.484-2.7-1.073-2.7-1.073a2.132 2.132 0 0 0-.9-1.184c-.735-.503.056-.493.056-.493a1.688 1.688 0 0 1 1.232.83 1.707 1.707 0 0 0 2.34.667 1.706 1.706 0 0 1 .509-1.073c-1.78-.202-3.644-.89-3.644-3.962A3.106 3.106 0 0 1 5.066 5.47a2.882 2.882 0 0 1 .078-2.13s.672-.215 2.2.82a7.634 7.634 0 0 1 4.004 0c1.528-1.035 2.2-.82 2.2-.82a2.882 2.882 0 0 1 .078 2.13 3.106 3.106 0 0 1 .966 2.152c0 3.08-1.866 3.756-3.648 3.954a1.918 1.918 0 0 1 .547 1.482v2.197c0 .214.146.462.55.383A8 8 0 0 0 8 .198Z"/></svg>
-        GitHub Repo
-      </a>
-    </div>
-    """)
+        <div style="text-align:center;margin-bottom:32px;">
+          <a class="link-btn project" href="https://zhaohanm.github.io/FusionDTI.github.io/" target="_blank" rel="noopener noreferrer" aria-label="Project Page">
+            <!-- globe icon -->
+            <svg xmlns="http://www.w3.org/2000/svg" width="18" height="18" viewBox="0 0 24 24" fill="currentColor" aria-hidden="true">
+              <path d="M12 2a10 10 0 1 0 10 10A10.012 10.012 0 0 0 12 2Zm7.93 9h-3.18a15.84 15.84 0 0 0-1.19-5.02A8.02 8.02 0 0 1 19.93 11ZM12 4c.86 0 2.25 1.86 3.01 6H8.99C9.75 5.86 11.14 4 12 4ZM4.07 13h3.18c.2 1.79.66 3.47 1.19 5.02A8.02 8.02 0 0 1 4.07 13Zm3.18-2H4.07A8.02 8.02 0 0 1 8.44 5.98 15.84 15.84 0 0 0 7.25 11Zm1.37 2h6.76c-.76 4.14-2.15 6-3.01 6s-2.25-1.86-3.01-6Zm9.05 0h3.18a8.02 8.02 0 0 1-4.37 5.02 15.84 15.84 0 0 0 1.19-5.02Z"/>
+            </svg>
+            Project Page
+          </a>
+          <a class="link-btn arxiv" href="https://arxiv.org/abs/2406.01651" target="_blank" rel="noopener noreferrer" aria-label="ArXiv: 2406.01651">
+            <!-- arXiv-like paper icon -->
+            <svg xmlns="http://www.w3.org/2000/svg" width="18" height="18" viewBox="0 0 24 24" fill="currentColor" aria-hidden="true">
+              <path d="M6 2h9l5 5v13a2 2 0 0 1-2 2H6a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2Zm8 1.5V8h4.5L14 3.5ZM7 12h10v2H7v-2Zm0 4h10v2H7v-2Zm0-8h6v2H7V8Z"/>
+            </svg>
+            ArXiv: 2406.01651
+          </a>
+          <a class="link-btn github" href="https://github.com/ZhaohanM/FusionDTI" target="_blank" rel="noopener noreferrer" aria-label="GitHub Repo">
+            <!-- GitHub mark -->
+            <svg xmlns="http://www.w3.org/2000/svg" width="18" height="18" viewBox="0 0 24 24" fill="currentColor" aria-hidden="true">
+              <path d="M12 .5A12 12 0 0 0 0 12.76c0 5.4 3.44 9.98 8.2 11.6.6.12.82-.28.82-.6v-2.3c-3.34.74-4.04-1.44-4.04-1.44-.54-1.38-1.32-1.74-1.32-1.74-1.08-.76.08-.74.08-.74 1.2.08 1.84 1.26 1.84 1.26 1.06 1.86 2.78 1.32 3.46 1.02.1-.8.42-1.32.76-1.62-2.66-.32-5.46-1.36-5.46-6.02 0-1.34.46-2.44 1.22-3.3-.12-.32-.54-1.64.12-3.42 0 0 1-.34 3.32 1.26.96-.28 1.98-.42 3-.42s2.04.14 3 .42c2.32-1.6 3.32-1.26 3.32-1.26.66 1.78.24 3.1.12 3.42.76.86 1.22 1.96 1.22 3.3 0 4.68-2.8 5.68-5.48 6 .44.38.84 1.12.84 2.28v3.38c0 .32.22.74.84.6A12.02 12.02 0 0 0 24 12.76 12 12 0 0 0 12 .5Z"/>
+            </svg>
+            GitHub Repo
+          </a>
+        </div>
+        """)
 
     # ───────────── Guidelines Card ─────────────
-    
     gr.HTML(
         """
-        <div class="card" style="margin-bottom:24px">
-          <h2 style="font-size:1.2rem;margin-bottom:14px">Guidelines for User</h2>
-          <ul style="font-size:1rem; margin-left:18px;line-height:1.55;list-style:decimal;">
+        <div class="card" id="guidelines-card" style="margin-bottom:24px">
+          <h2>Guidelines for Users</h2>
+          <ul style="list-style:decimal;">
             <li><strong>Convert protein structure into a structure-aware sequence:</strong>
                 Upload a <code>.pdb</code> or <code>.cif</code> file. A structure-aware
                 sequence will be generated using
@@ -485,26 +470,25 @@ with gr.Blocks(theme=Soft(primary_hue="indigo", neutral_hue="slate"), css=css) a
                 <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
+                please 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
+                to download the corresponding <code>.cif</code> or <code>.pdb</code> file.</li>
+            <li><strong>Drug input supports both SELFIES and SMILES:</strong>
+                Enter a SELFIES string directly, or paste a SMILES string. SMILES will
+                be converted to SELFIES using the
                 <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
+            <li>Optionally enter a <strong>1-based</strong> drug atom/substructure index
                 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>
+            <li>After inference, use “Download PDF” to export a high-resolution vector figure.</li>
           </ul>
         </div>
-        """)
-    
+        """
+    )
+
     # ───────────── Input Card ─────────────
     with gr.Column(elem_id="input-card", elem_classes="card"):
-
         protein_seq = gr.Textbox(
             label="Protein Structure-aware Sequence",
             lines=3,
@@ -529,25 +513,12 @@ with gr.Blocks(theme=Soft(primary_hue="indigo", neutral_hue="slate"), css=css) a
 
     # ───────────── Action Buttons ─────────────
     with gr.Row(elem_id="action-buttons", equal_height=True):
-        btn_extract = gr.Button(
-            "Extract sequence",
-            variant="primary",
-            elem_id="extract-btn"
-        )
-        btn_infer = gr.Button(
-            "Inference",
-            variant="primary",
-            elem_id="inference-btn"
-        )
+        btn_extract = gr.Button("Extract sequence", variant="primary", elem_id="extract-btn")
+        btn_infer   = gr.Button("Inference",        variant="primary", elem_id="inference-btn")
     with gr.Row():
-        clear_btn = gr.Button(
-            "Clear",
-            variant="secondary",
-            elem_classes="full-width",
-            elem_id="clear-btn"
-        )
+        clear_btn   = gr.Button("Clear", variant="secondary", elem_id="clear-btn")
 
-    # ───────────── Output Visualization ─────────────
+    # ───────────── Output Visualisation ─────────────
     output_html  = gr.HTML(elem_id="result-html")
 
     # ───────────── Event Wiring ─────────────
@@ -562,7 +533,7 @@ with gr.Blocks(theme=Soft(primary_hue="indigo", neutral_hue="slate"), css=css) a
         outputs=[output_html]
     )
     clear_btn.click(
-        fn=lambda: ("", "", None, "", None),
+        fn=clear_cb,
         inputs=[],
         outputs=[protein_seq, drug_seq, drug_idx, output_html, structure_file]
     )