Update: SMILES-to-SELFIES conversion, UI polish, and usage guide

.ipynb_checkpoints/app-checkpoint.py

@@ -0,0 +1,472 @@
+import os, sys, argparse, tempfile, shutil, base64, io
+from flask import Flask, request, render_template_string
+from werkzeug.utils import secure_filename
+from torch.utils.data import DataLoader
+import selfies
+from rdkit import Chem
+
+import torch
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from matplotlib import cm
+from typing import Optional
+
+from utils.drug_tokenizer import DrugTokenizer
+from transformers import EsmForMaskedLM, EsmTokenizer, AutoModel
+from utils.metric_learning_models_att_maps import Pre_encoded, FusionDTI
+from utils.foldseek_util import get_struc_seq
+
+# ───── Biopython fallback ───────────────────────────────────────
+from Bio.PDB import PDBParser, MMCIFParser
+from Bio.Data import IUPACData
+
+three2one = {k.upper(): v for k, v in IUPACData.protein_letters_3to1.items()}
+three2one.update({"SEC": "C", "PYL": "K"})
+def simple_seq_from_structure(path: str) -> str:
+    parser = MMCIFParser(QUIET=True) if path.endswith(".cif") else PDBParser(QUIET=True)
+    chain  = next(parser.get_structure("P", path).get_chains())
+    return "".join(three2one.get(res.get_resname().upper(), "X") for res in chain)
+
+# ───── global paths / args ──────────────────────────────────────
+FOLDSEEK_BIN = shutil.which("foldseek")
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+sys.path.append("..")
+
+def parse_config():
+    p = argparse.ArgumentParser()
+    p.add_argument("-f")
+    p.add_argument("--prot_encoder_path", default="westlake-repl/SaProt_650M_AF2")
+    p.add_argument("--drug_encoder_path", default="HUBioDataLab/SELFormer")
+    p.add_argument("--agg_mode", default="mean_all_tok", type=str, help="{cls|mean|mean_all_tok}")
+    p.add_argument("--group_size", type=int, default=1)
+    p.add_argument("--lr", type=float, default=1e-4)
+    p.add_argument("--fusion", default="CAN")
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--save_path_prefix", default="save_model_ckp/")
+    p.add_argument("--dataset", default="BindingDB"), help="Name of the dataset to use (e.g., 'BindingDB', 'Human', 'Biosnap')"
+    return p.parse_args()
+
+args = parse_config()
+DEVICE = args.device
+
+# ───── tokenisers & encoders ────────────────────────────────────
+prot_tokenizer = EsmTokenizer.from_pretrained(args.prot_encoder_path)
+prot_model     = EsmForMaskedLM.from_pretrained(args.prot_encoder_path)
+
+drug_tokenizer = DrugTokenizer()        # SELFIES
+drug_model     = AutoModel.from_pretrained(args.drug_encoder_path)
+
+encoding = Pre_encoded(prot_model, drug_model, args).to(DEVICE)
+
+# ─── collate fn ────────────────────────────────────────────────
+def collate_fn(batch):
+    query1, query2, scores = zip(*batch)
+    
+    query_encodings1 = prot_tokenizer.batch_encode_plus(
+        list(query1),
+        max_length=512,
+        padding="max_length",
+        truncation=True,
+        add_special_tokens=True,
+        return_tensors="pt",
+    )
+    query_encodings2 = drug_tokenizer.batch_encode_plus(
+        list(query2),
+        max_length=512,
+        padding="max_length",
+        truncation=True,
+        add_special_tokens=True,
+        return_tensors="pt",
+    )
+    scores = torch.tensor(list(scores))
+
+    attention_mask1 = query_encodings1["attention_mask"].bool()
+    attention_mask2 = query_encodings2["attention_mask"].bool()
+
+    return query_encodings1["input_ids"], attention_mask1, query_encodings2["input_ids"], attention_mask2, scores
+# def collate_fn_batch_encoding(batch):
+
+def smiles_to_selfies(smiles: str) -> Optional[str]:
+    try:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return None
+        selfies_str = selfies.encoder(smiles)
+        return selfies_str
+    except Exception:
+        return None
+
+
+# ───── single-case embedding ───────────────────────────────────
+def get_case_feature(model, loader):
+    model.eval()
+    with torch.no_grad():
+        for p_ids, p_mask, d_ids, d_mask, _ in loader:
+            p_ids, p_mask = p_ids.to(DEVICE), p_mask.to(DEVICE)
+            d_ids, d_mask = d_ids.to(DEVICE), d_mask.to(DEVICE)
+            p_emb, d_emb = model.encoding(p_ids, p_mask, d_ids, d_mask)
+            return [(p_emb.cpu(), d_emb.cpu(),
+                     p_ids.cpu(), d_ids.cpu(),
+                     p_mask.cpu(), d_mask.cpu(), None)]
+
+# ───── helper：过滤特殊 token ───────────────────────────────────
+def clean_tokens(ids, tokenizer):
+    toks = tokenizer.convert_ids_to_tokens(ids.tolist())
+    return [t for t in toks if t not in tokenizer.all_special_tokens]
+
+# ───── visualisation ───────────────────────────────────────────
+
+def visualize_attention(model, feats, drug_idx: Optional[int] = None) -> str:
+    """
+    Render a Protein → Drug cross-attention heat-map and, optionally, a
+    Top-20 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).
+    """
+    model.eval()
+    with torch.no_grad():
+        # ── unpack single-case tensors ───────────────────────────────────────────
+        p_emb, d_emb, p_ids, d_ids, p_mask, d_mask, _ = feats[0]
+        p_emb, d_emb = p_emb.to(DEVICE), d_emb.to(DEVICE)
+        p_mask, d_mask = p_mask.to(DEVICE), d_mask.to(DEVICE)
+
+        # ── 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)
+
+        # ── 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]
+
+        # ── 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)]
+
+        # ── 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)
+
+        if row_keep.sum() < 3:
+            row_keep[:] = True
+        if col_keep.sum() < 3:
+            col_keep[:] = True
+
+        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]
+
+        # ── 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]
+
+        # ── 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, 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)
+
+        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_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)
+
+        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-20 表（若需要） ─────────────────────
+        table_html = ""                   # 先设空串，方便后面统一拼接
+        if drug_idx is not None:
+            # 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(20, 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[col_pos]
+                orig_idx      = d_indices[col_pos]
+
+                table_html = (
+                    f"<h4 style='margin-bottom:6px'>"
+                    f"Drug token #{orig_idx} <code>{drug_tok_text}</code> "
+                    f"→ Top-20 Protein residues</h4>"
+                    "<table class='tg' style='margin-bottom:8px'>"
+                    f"<tr><th>Rank</th>{rank_hdr}</tr>"
+                    f"<tr><td>Residue</td>{res_row}</tr>"
+                    f"<tr><td>Position</td>{pos_row}</tr>"
+                    "</table>")
+
+        # ────────────────── 生成可放大 + 可下载的热图 ────────────────────
+        buf_png = io.BytesIO()
+        fig.savefig(buf_png, format="png", dpi=140)   # 预览（光栅）
+        buf_png.seek(0)
+
+        buf_pdf = io.BytesIO()
+        fig.savefig(buf_pdf, format="pdf")            # 高清下载（矢量）
+        buf_pdf.seek(0)
+        plt.close(fig)
+
+        png_b64 = base64.b64encode(buf_png.getvalue()).decode()
+        pdf_b64 = base64.b64encode(buf_pdf.getvalue()).decode()
+
+        html_heat = (
+            f"<a href='data:image/png;base64,{png_b64}' target='_blank' "
+            f"title='Click to enlarge'>"
+            f"<img src='data:image/png;base64,{png_b64}' "
+            f"style='max-width:100%;height:auto;cursor:zoom-in' /></a>"
+            f"<div style='margin-top:6px'>"
+            f"<a href='data:application/pdf;base64,{pdf_b64}' "
+            f"download='attention_heatmap.pdf'>Download PDF</a></div>"
+        )
+
+        # ───────────────────────── 返回最终 HTML ─────────────────────────
+        return table_html + html_heat
+
+
+# ───── Flask app ───────────────────────────────────────────────
+app = Flask(__name__)
+
+@app.route("/", methods=["GET", "POST"])
+def index():
+    protein_seq = drug_seq = structure_seq = ""; result_html = None
+    tmp_structure_path = ""; drug_idx = None
+
+    if request.method == "POST":
+        drug_idx_raw = request.form.get("drug_idx", "")
+        drug_idx = int(drug_idx_raw)-1 if drug_idx_raw.isdigit() else None
+
+        struct = request.files.get("structure_file")
+        if struct and struct.filename:
+            path = os.path.join(tempfile.gettempdir(), secure_filename(struct.filename))
+            struct.save(path); tmp_structure_path = path
+        else:
+            tmp_structure_path = request.form.get("tmp_structure_path", "")
+
+        if "clear" in request.form:
+            protein_seq = drug_seq = structure_seq = ""; tmp_structure_path = ""
+
+        elif "confirm_structure" in request.form and tmp_structure_path:
+            try:
+                parsed = get_struc_seq(FOLDSEEK_BIN, tmp_structure_path, None, plddt_mask=False)
+                chain  = list(parsed.keys())[0]; _, _, structure_seq = parsed[chain]
+            except Exception:
+                structure_seq = simple_seq_from_structure(tmp_structure_path)
+            protein_seq = structure_seq
+            drug_input = request.form.get("drug_sequence", "")
+            # Heuristically check if input is SMILES (not starting with [) and convert
+            if not drug_input.strip().startswith("["):
+                converted = smiles_to_selfies(drug_input.strip())
+                if converted:
+                    drug_seq = converted
+                else:
+                    drug_seq = ""
+                    result_html = "<p style='color:red'><strong>Failed to convert SMILES to SELFIES. Please check the input string.</strong></p>"
+            else:
+                drug_seq = drug_input
+
+        elif "Inference" in request.form:
+            protein_seq = request.form.get("protein_sequence", "")
+            drug_seq    = request.form.get("drug_sequence", "")
+            if protein_seq and drug_seq:
+                loader = DataLoader([(protein_seq, drug_seq, 1)], batch_size=1,
+                                    collate_fn=collate_fn)
+                feats  = get_case_feature(encoding, loader)
+                model  = FusionDTI(446, 768, args).to(DEVICE)
+                ckpt   = os.path.join(f"{args.save_path_prefix}{args.dataset}_{args.fusion}",
+                                      "best_model.ckpt")
+                if os.path.isfile(ckpt):
+                    model.load_state_dict(torch.load(ckpt, map_location=DEVICE))
+                result_html = visualize_attention(model, feats, drug_idx)
+
+    return render_template_string(
+    # ───────────── HTML (原 UI + 新输入框) ─────────────
+    """
+<!doctype html>
+<html lang="en"><head><meta charset="utf-8"><title>FusionDTI </title>
+<link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600&family=Poppins:wght@500;600&display=swap" rel="stylesheet">
+
+<style>
+:root{--bg:#f3f4f6;--card:#fff;--primary:#6366f1;--primary-dark:#4f46e5;--text:#111827;--border:#e5e7eb;}
+*{box-sizing:border-box;margin:0;padding:0}
+body{background:var(--bg);color:var(--text);font-family:Inter,system-ui,Arial,sans-serif;line-height:1.5;padding:32px 12px;}
+h1{font-family:Poppins,Inter,sans-serif;font-weight:600;font-size:1.7rem;text-align:center;margin-bottom:28px;letter-spacing:-.2px;}
+.card{max-width:1000px;margin:0 auto;background:var(--card);border:1px solid var(--border);
+      border-radius:12px;box-shadow:0 2px 6px rgba(0,0,0,.05);padding:32px 36px;}
+label{font-weight:500;margin-bottom:6px;display:block}
+textarea,input[type=file]{width:100%;font-size:.9rem;font-family:monospace;padding:10px 12px;
+      border:1px solid var(--border);border-radius:8px;background:#fff;resize:vertical;}
+textarea{min-height:90px}
+.btn{appearance:none;border:none;cursor:pointer;padding:12px 22px;border-radius:8px;font-weight:500;
+     font-family:Inter,sans-serif;transition:all .18s ease;color:#fff;}
+.btn-primary{background:var(--primary)}.btn-primary:hover{background:var(--primary-dark)}
+.btn-neutral{background:#9ca3af;}.btn-neutral:hover{background:#6b7280}
+.grid{display:grid;gap:22px}.grid-2{grid-template-columns:1fr 1fr}
+.vis-box{margin-top:28px;border:1px solid var(--border);border-radius:10px;overflow:auto;max-height:72vh;}
+pre{white-space:pre-wrap;word-break:break-all;font-family:monospace;margin-top:8px}
+
+/* ── tidy table for Top-20 list ─────────────────────────────── */
+table.tg{border-collapse:collapse;margin-top:4px;font-size:0.83rem}
+table.tg th,table.tg td{border:1px solid var(--border);padding:6px 8px;text-align:left}
+table.tg th{background:var(--bg);font-weight:600}
+</style>
+</head>
+<body>
+<h1> Token-level Visualiser for Drug-Target Interaction</h1>
+
+<!-- ───────────── Project Links (larger + spaced) ───────────── -->
+<div style="margin-top:24px; text-align:center;">
+  <a href="https://zhaohanm.github.io/FusionDTI.github.io/" target="_blank"
+     style="display:inline-block;margin:8px 18px;padding:10px 20px;
+            background:linear-gradient(to right,#10b981,#059669);color:white;
+            font-weight:600;border-radius:8px;font-size:0.9rem;
+            font-family:Inter,sans-serif;text-decoration:none;
+            box-shadow:0 2px 6px rgba(0,0,0,0.12);transition:all 0.2s ease-in-out;"
+     onmouseover="this.style.opacity='0.9'" onmouseout="this.style.opacity='1'">
+    🌐 Project Page
+  </a>
+
+  <a href="https://arxiv.org/abs/2406.01651" target="_blank"
+     style="display:inline-block;margin:8px 18px;padding:10px 20px;
+            background:linear-gradient(to right,#ef4444,#dc2626);color:white;
+            font-weight:600;border-radius:8px;font-size:0.9rem;
+            font-family:Inter,sans-serif;text-decoration:none;
+            box-shadow:0 2px 6px rgba(0,0,0,0.12);transition:all 0.2s ease-in-out;"
+     onmouseover="this.style.opacity='0.9'" onmouseout="this.style.opacity='1'">
+    📄 ArXiv: 2406.01651
+  </a>
+
+  <a href="https://github.com/ZhaohanM/FusionDTI" target="_blank"
+     style="display:inline-block;margin:8px 18px;padding:10px 20px;
+            background:linear-gradient(to right,#3b82f6,#2563eb);color:white;
+            font-weight:600;border-radius:8px;font-size:0.9rem;
+            font-family:Inter,sans-serif;text-decoration:none;
+            box-shadow:0 2px 6px rgba(0,0,0,0.12);transition:all 0.2s ease-in-out;"
+     onmouseover="this.style.opacity='0.9'" onmouseout="this.style.opacity='1'">
+    💻 GitHub Repo
+  </a>
+</div>
+
+<!-- ─────────────  Guidelines for Use  ───────────── -->
+<div class="card" style="margin-bottom:24px">
+  <h2 style="font-size:1.2rem;margin-bottom:14px">Guidelines for Use</h2>
+  <ul style="margin-left:18px;line-height:1.55;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
+        <a href="https://github.com/steineggerlab/foldseek" target="_blank">Foldseek</a>,
+        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>
+
+    <li>After inference, you can use the
+        “Download PDF” link to export a high-resolution vector version.</li>
+  </ul>
+</div>
+
+<div class="card">
+<form method="POST" enctype="multipart/form-data" class="grid">
+
+  <div><label>Protein Structure (.pdb / .cif)</label>
+       <input type="file" name="structure_file">
+       <input type="hidden" name="tmp_structure_path" value="{{ tmp_structure_path }}"></div>
+
+  <div><label>Protein Sequence</label>
+       <textarea name="protein_sequence" placeholder="Confirm / paste sequence…">{{ protein_seq }}</textarea></div>
+
+  <div><label>Drug Sequence (SELFIES/SMILES)</label>
+       <textarea name="drug_sequence" placeholder="[C][C][O]/cco …">{{ drug_seq }}</textarea></div>
+
+    <label>Drug atom/substructure index (1-based) – show Top-10 related protein residue</label>
+        <input type="number" name="drug_idx" min="1" style="width:120px">
+
+  <div class="grid grid-2">
+    <button class="btn btn-primary" type="Inference" name="confirm_structure">Confirm Structure</button>
+    <button class="btn btn-primary" type="Inference" name="Inference">Inference</button>
+  </div>
+  <button class="btn btn-neutral" style="width:100%" type="Inference" name="clear">Clear</button>
+</form>
+
+{% if structure_seq %}
+  <div style="margin-top:18px"><strong>Structure-aware sequence:</strong><pre>{{ structure_seq }}</pre></div>
+{% endif %}
+{% if result_html %}
+  <div class="vis-box" style="margin-top:26px">{{ result_html|safe }}</div>
+{% endif %}
+</div></body></html>
+    """,
+    protein_seq=protein_seq, drug_seq=drug_seq, structure_seq=structure_seq,
+    result_html=result_html, tmp_structure_path=tmp_structure_path)
+
+# ───── run ─────────────────────────────────────────────────────
+if __name__ == "__main__":
+    app.run(debug=True, host="0.0.0.0", port=7860)

.ipynb_checkpoints/requirements-checkpoint.txt

@@ -0,0 +1,11 @@
+Flask
+torch
+transformers
+IPython
+selfies
+rdkit
+biopython
+matplotlib
+scikit-learn
+numpy
+pandas

app.py

@@ -1,209 +1,66 @@
-import os
-import sys
-import argparse
-import torch
+import os, sys, argparse, tempfile, shutil, base64, io
+from flask import Flask, request, render_template_string
+from werkzeug.utils import secure_filename
 from torch.utils.data import DataLoader
-from transformers import EsmForMaskedLM, AutoModel, EsmTokenizer
+import selfies
+from rdkit import Chem
+
+import torch
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from matplotlib import cm
+from typing import Optional
+
 from utils.drug_tokenizer import DrugTokenizer
+from transformers import EsmForMaskedLM, EsmTokenizer, AutoModel
 from utils.metric_learning_models_att_maps import Pre_encoded, FusionDTI
-from bertviz import head_view
-import tempfile
-from flask import Flask, request, render_template_string
+from utils.foldseek_util import get_struc_seq
 
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-sys.path.append("../")
+# ───── Biopython fallback ───────────────────────────────────────
+from Bio.PDB import PDBParser, MMCIFParser
+from Bio.Data import IUPACData
 
-app = Flask(__name__)
+three2one = {k.upper(): v for k, v in IUPACData.protein_letters_3to1.items()}
+three2one.update({"SEC": "C", "PYL": "K"})
+def simple_seq_from_structure(path: str) -> str:
+    parser = MMCIFParser(QUIET=True) if path.endswith(".cif") else PDBParser(QUIET=True)
+    chain  = next(parser.get_structure("P", path).get_chains())
+    return "".join(three2one.get(res.get_resname().upper(), "X") for res in chain)
+
+# ───── global paths / args ──────────────────────────────────────
+FOLDSEEK_BIN = shutil.which("foldseek")
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+sys.path.append("..")
 
 def parse_config():
-    parser = argparse.ArgumentParser()
-    parser.add_argument('-f')
-    parser.add_argument("--prot_encoder_path", type=str, default="westlake-repl/SaProt_650M_AF2", help="path/name of protein encoder model located")
-    parser.add_argument("--drug_encoder_path", type=str, default="HUBioDataLab/SELFormer", help="path/name of SMILE pre-trained language model")
-    parser.add_argument("--agg_mode", default="mean_all_tok", type=str, help="{cls|mean|mean_all_tok}")
-    parser.add_argument("--fusion", default="CAN", type=str, help="{CAN|BAN}")
-    parser.add_argument("--batch_size", type=int, default=64)
-    parser.add_argument("--group_size", type=int, default=1)
-    parser.add_argument("--lr", type=float, default=1e-4)
-    parser.add_argument("--dropout", type=float, default=0.1)
-    parser.add_argument("--test", type=int, default=0)
-    parser.add_argument("--use_pooled", action="store_true", default=True)
-    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
-    parser.add_argument("--save_path_prefix", type=str, default="save_model_ckp/", help="save the result in which directory")
-    parser.add_argument("--save_name", default="fine_tune", type=str, help="the name of the saved file")
-    parser.add_argument("--dataset", type=str, default="Human", help="Name of the dataset to use (e.g., 'BindingDB', 'Human', 'Biosnap')")
-    return parser.parse_args()
+    p = argparse.ArgumentParser()
+    p.add_argument("-f")
+    p.add_argument("--prot_encoder_path", default="westlake-repl/SaProt_650M_AF2")
+    p.add_argument("--drug_encoder_path", default="HUBioDataLab/SELFormer")
+    p.add_argument("--agg_mode", default="mean_all_tok", type=str, help="{cls|mean|mean_all_tok}")
+    p.add_argument("--group_size", type=int, default=1)
+    p.add_argument("--lr", type=float, default=1e-4)
+    p.add_argument("--fusion", default="CAN")
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--save_path_prefix", default="save_model_ckp/")
+    p.add_argument("--dataset", default="BindingDB"), help="Name of the dataset to use (e.g., 'BindingDB', 'Human', 'Biosnap')"
+    return p.parse_args()
 
 args = parse_config()
-device = args.device
+DEVICE = args.device
 
+# ───── tokenisers & encoders ────────────────────────────────────
 prot_tokenizer = EsmTokenizer.from_pretrained(args.prot_encoder_path)
-drug_tokenizer = DrugTokenizer()
-
-prot_model = EsmForMaskedLM.from_pretrained(args.prot_encoder_path)
-drug_model = AutoModel.from_pretrained(args.drug_encoder_path)
+prot_model     = EsmForMaskedLM.from_pretrained(args.prot_encoder_path)
 
-encoding = Pre_encoded(prot_model, drug_model, args).to(device)
+drug_tokenizer = DrugTokenizer()        # SELFIES
+drug_model     = AutoModel.from_pretrained(args.drug_encoder_path)
 
-def get_case_feature(model, dataloader, device):
-    with torch.no_grad():
-        for step, batch in enumerate(dataloader):
-            prot_input_ids, prot_attention_mask, drug_input_ids, drug_attention_mask, label = batch
-            prot_input_ids, prot_attention_mask, drug_input_ids, drug_attention_mask = \
-                prot_input_ids.to(device), prot_attention_mask.to(device), drug_input_ids.to(device), drug_attention_mask.to(device)
-
-            prot_embed, drug_embed = model.encoding(prot_input_ids, prot_attention_mask, drug_input_ids, drug_attention_mask)
-            prot_embed, drug_embed = prot_embed.cpu(), drug_embed.cpu()
-            prot_input_ids, drug_input_ids = prot_input_ids.cpu(), drug_input_ids.cpu()
-            prot_attention_mask, drug_attention_mask = prot_attention_mask.cpu(), drug_attention_mask.cpu()
-            label = label.cpu()
-
-            return [(prot_embed, drug_embed, prot_input_ids, drug_input_ids, prot_attention_mask, drug_attention_mask, label)]
+encoding = Pre_encoded(prot_model, drug_model, args).to(DEVICE)
 
-def visualize_attention(model, case_features, device, prot_tokenizer, drug_tokenizer):
-    model.eval()
-    with torch.no_grad():
-        for batch in case_features:
-            prot, drug, prot_ids, drug_ids, prot_mask, drug_mask, label = batch
-            prot, drug = prot.to(device), drug.to(device)
-            prot_mask, drug_mask = prot_mask.to(device), drug_mask.to(device)
-
-            output, attention_weights = model(prot, drug, prot_mask, drug_mask)
-            prot_tokens = [prot_tokenizer.decode([pid.item()], skip_special_tokens=True) for pid in prot_ids.squeeze()]
-            drug_tokens = [drug_tokenizer.decode([did.item()], skip_special_tokens=True) for did in drug_ids.squeeze()]
-            tokens = prot_tokens + drug_tokens
-
-            attention_weights = attention_weights.unsqueeze(1)
-
-            # Generate HTML content using head_view with html_action='return'
-            html_head_view = head_view(attention_weights, tokens, sentence_b_start=512, html_action='return')
-
-            # Parse the HTML and modify it to replace sentence labels
-            html_content = html_head_view.data
-            html_content = html_content.replace("Sentence A -> Sentence A", "Protein -> Protein")
-            html_content = html_content.replace("Sentence B -> Sentence B", "Drug -> Drug")
-            html_content = html_content.replace("Sentence A -> Sentence B", "Protein -> Drug")
-            html_content = html_content.replace("Sentence B -> Sentence A", "Drug -> Protein")
-
-            # Save the modified HTML content to a temporary file
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".html") as f:
-                f.write(html_content.encode('utf-8'))
-                temp_file_path = f.name
-            
-            return temp_file_path
-
-@app.route('/', methods=['GET', 'POST'])
-def index():
-    protein_sequence = ""
-    drug_sequence = ""
-    result = None
-
-    if request.method == 'POST':
-        if 'clear' in request.form:
-            protein_sequence = ""
-            drug_sequence = ""
-        else:
-            protein_sequence = request.form['protein_sequence']
-            drug_sequence = request.form['drug_sequence']
-            
-            dataset = [(protein_sequence, drug_sequence, 1)]
-            dataloader = DataLoader(dataset, batch_size=1, collate_fn=collate_fn_batch_encoding)
-            
-            case_features = get_case_feature(encoding, dataloader, device)
-            model = FusionDTI(446, 768, args).to(device)
-            
-            best_model_dir = f"{args.save_path_prefix}{args.dataset}_{args.fusion}"  
-            checkpoint_path = os.path.join(best_model_dir, 'best_model.ckpt')
-            
-            if os.path.exists(checkpoint_path):
-                model.load_state_dict(torch.load(checkpoint_path, map_location=device))
-            
-            html_file_path = visualize_attention(model, case_features, device, prot_tokenizer, drug_tokenizer)
-            
-            with open(html_file_path, 'r') as f:
-                result = f.read()
-    
-    return render_template_string('''
-        <html>
-            <head>
-                <title>Drug Target Interaction Visualization</title>
-                <style>
-                    body { font-family: 'Times New Roman', Times, serif; margin: 40px; }
-                    h2 { color: #333; }
-                    .container { display: flex; }
-                    .left { flex: 1; padding-right: 20px; }
-                    .right { flex: 1; }
-                    textarea {
-                        width: 100%;
-                        padding: 12px 20px;
-                        margin: 8px 0;
-                        display: inline-block;
-                        border: 1px solid #ccc;
-                        border-radius: 4px;
-                        box-sizing: border-box;
-                        font-size: 16px;
-                        font-family: 'Times New Roman', Times, serif;
-                    }
-                    .button-container {
-                        display: flex;
-                        justify-content: space-between;
-                    }
-                    input[type="submit"], .button {
-                        width: 48%;
-                        color: white;
-                        padding: 14px 20px;
-                        margin: 8px 0;
-                        border: none;
-                        border-radius: 4px;
-                        cursor: pointer;
-                        font-size: 16px;
-                        font-family: 'Times New Roman', Times, serif;
-                    }
-                    .submit {
-                        background-color: #FFA500;
-                    }
-                    .submit:hover {
-                        background-color: #FF8C00;
-                    }
-                    .clear {
-                        background-color: #D3D3D3;
-                    }
-                    .clear:hover {
-                        background-color: #A9A9A9;
-                    }
-                    .result {
-                        font-size: 18px;
-                    }
-                </style>
-            </head>
-            <body>
-                <h2 style="text-align: center;">Drug Target Interaction Visualization</h2>
-                <div class="container">
-                    <div class="left">
-                        <form method="post">
-                            <label for="protein_sequence">Protein Sequence:</label>
-                            <textarea id="protein_sequence" name="protein_sequence" rows="4" placeholder="Enter protein sequence here..." required>{{ protein_sequence }}</textarea><br>
-                            <label for="drug_sequence">Drug Sequence:</label>
-                            <textarea id="drug_sequence" name="drug_sequence" rows="4" placeholder="Enter drug sequence here..." required>{{ drug_sequence }}</textarea><br>
-                            <div class="button-container">
-                                <input type="submit" name="submit" class="button submit" value="Submit">
-                                <input type="submit" name="clear" class="button clear" value="Clear">
-                            </div>
-                        </form>
-                    </div>
-                    <div class="right" style="display: flex; justify-content: center; align-items: center;">
-                        {% if result %}
-                            <div class="result">
-                                {{ result|safe }}
-                            </div>
-                        {% endif %}
-                    </div>
-                </div>
-            </body>
-        </html>
-    ''', protein_sequence=protein_sequence, drug_sequence=drug_sequence, result=result)
-
-def collate_fn_batch_encoding(batch):
+# ─── collate fn ────────────────────────────────────────────────
+def collate_fn(batch):
     query1, query2, scores = zip(*batch)
     
     query_encodings1 = prot_tokenizer.batch_encode_plus(
@@ -228,6 +85,388 @@ def collate_fn_batch_encoding(batch):
     attention_mask2 = query_encodings2["attention_mask"].bool()
 
     return query_encodings1["input_ids"], attention_mask1, query_encodings2["input_ids"], attention_mask2, scores
+# def collate_fn_batch_encoding(batch):
+
+def smiles_to_selfies(smiles: str) -> Optional[str]:
+    try:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return None
+        selfies_str = selfies.encoder(smiles)
+        return selfies_str
+    except Exception:
+        return None
+
+
+# ───── single-case embedding ───────────────────────────────────
+def get_case_feature(model, loader):
+    model.eval()
+    with torch.no_grad():
+        for p_ids, p_mask, d_ids, d_mask, _ in loader:
+            p_ids, p_mask = p_ids.to(DEVICE), p_mask.to(DEVICE)
+            d_ids, d_mask = d_ids.to(DEVICE), d_mask.to(DEVICE)
+            p_emb, d_emb = model.encoding(p_ids, p_mask, d_ids, d_mask)
+            return [(p_emb.cpu(), d_emb.cpu(),
+                     p_ids.cpu(), d_ids.cpu(),
+                     p_mask.cpu(), d_mask.cpu(), None)]
+
+# ───── helper：过滤特殊 token ───────────────────────────────────
+def clean_tokens(ids, tokenizer):
+    toks = tokenizer.convert_ids_to_tokens(ids.tolist())
+    return [t for t in toks if t not in tokenizer.all_special_tokens]
+
+# ───── visualisation ───────────────────────────────────────────
+
+def visualize_attention(model, feats, drug_idx: Optional[int] = None) -> str:
+    """
+    Render a Protein → Drug cross-attention heat-map and, optionally, a
+    Top-20 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).
+    """
+    model.eval()
+    with torch.no_grad():
+        # ── unpack single-case tensors ───────────────────────────────────────────
+        p_emb, d_emb, p_ids, d_ids, p_mask, d_mask, _ = feats[0]
+        p_emb, d_emb = p_emb.to(DEVICE), d_emb.to(DEVICE)
+        p_mask, d_mask = p_mask.to(DEVICE), d_mask.to(DEVICE)
+
+        # ── 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)
+
+        # ── 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]
+
+        # ── 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)]
+
+        # ── 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)
+
+        if row_keep.sum() < 3:
+            row_keep[:] = True
+        if col_keep.sum() < 3:
+            col_keep[:] = True
+
+        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]
+
+        # ── 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]
+
+        # ── 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, 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)
+
+        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_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)
+
+        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-20 表（若需要） ─────────────────────
+        table_html = ""                   # 先设空串，方便后面统一拼接
+        if drug_idx is not None:
+            # 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(20, 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[col_pos]
+                orig_idx      = d_indices[col_pos]
+
+                table_html = (
+                    f"<h4 style='margin-bottom:6px'>"
+                    f"Drug token #{orig_idx} <code>{drug_tok_text}</code> "
+                    f"→ Top-20 Protein residues</h4>"
+                    "<table class='tg' style='margin-bottom:8px'>"
+                    f"<tr><th>Rank</th>{rank_hdr}</tr>"
+                    f"<tr><td>Residue</td>{res_row}</tr>"
+                    f"<tr><td>Position</td>{pos_row}</tr>"
+                    "</table>")
+
+        # ────────────────── 生成可放大 + 可下载的热图 ────────────────────
+        buf_png = io.BytesIO()
+        fig.savefig(buf_png, format="png", dpi=140)   # 预览（光栅）
+        buf_png.seek(0)
+
+        buf_pdf = io.BytesIO()
+        fig.savefig(buf_pdf, format="pdf")            # 高清下载（矢量）
+        buf_pdf.seek(0)
+        plt.close(fig)
+
+        png_b64 = base64.b64encode(buf_png.getvalue()).decode()
+        pdf_b64 = base64.b64encode(buf_pdf.getvalue()).decode()
+
+        html_heat = (
+            f"<a href='data:image/png;base64,{png_b64}' target='_blank' "
+            f"title='Click to enlarge'>"
+            f"<img src='data:image/png;base64,{png_b64}' "
+            f"style='max-width:100%;height:auto;cursor:zoom-in' /></a>"
+            f"<div style='margin-top:6px'>"
+            f"<a href='data:application/pdf;base64,{pdf_b64}' "
+            f"download='attention_heatmap.pdf'>Download PDF</a></div>"
+        )
+
+        # ───────────────────────── 返回最终 HTML ─────────────────────────
+        return table_html + html_heat
+
+
+# ───── Flask app ───────────────────────────────────────────────
+app = Flask(__name__)
+
+@app.route("/", methods=["GET", "POST"])
+def index():
+    protein_seq = drug_seq = structure_seq = ""; result_html = None
+    tmp_structure_path = ""; drug_idx = None
+
+    if request.method == "POST":
+        drug_idx_raw = request.form.get("drug_idx", "")
+        drug_idx = int(drug_idx_raw)-1 if drug_idx_raw.isdigit() else None
+
+        struct = request.files.get("structure_file")
+        if struct and struct.filename:
+            path = os.path.join(tempfile.gettempdir(), secure_filename(struct.filename))
+            struct.save(path); tmp_structure_path = path
+        else:
+            tmp_structure_path = request.form.get("tmp_structure_path", "")
+
+        if "clear" in request.form:
+            protein_seq = drug_seq = structure_seq = ""; tmp_structure_path = ""
+
+        elif "confirm_structure" in request.form and tmp_structure_path:
+            try:
+                parsed = get_struc_seq(FOLDSEEK_BIN, tmp_structure_path, None, plddt_mask=False)
+                chain  = list(parsed.keys())[0]; _, _, structure_seq = parsed[chain]
+            except Exception:
+                structure_seq = simple_seq_from_structure(tmp_structure_path)
+            protein_seq = structure_seq
+            drug_input = request.form.get("drug_sequence", "")
+            # Heuristically check if input is SMILES (not starting with [) and convert
+            if not drug_input.strip().startswith("["):
+                converted = smiles_to_selfies(drug_input.strip())
+                if converted:
+                    drug_seq = converted
+                else:
+                    drug_seq = ""
+                    result_html = "<p style='color:red'><strong>Failed to convert SMILES to SELFIES. Please check the input string.</strong></p>"
+            else:
+                drug_seq = drug_input
+
+        elif "Inference" in request.form:
+            protein_seq = request.form.get("protein_sequence", "")
+            drug_seq    = request.form.get("drug_sequence", "")
+            if protein_seq and drug_seq:
+                loader = DataLoader([(protein_seq, drug_seq, 1)], batch_size=1,
+                                    collate_fn=collate_fn)
+                feats  = get_case_feature(encoding, loader)
+                model  = FusionDTI(446, 768, args).to(DEVICE)
+                ckpt   = os.path.join(f"{args.save_path_prefix}{args.dataset}_{args.fusion}",
+                                      "best_model.ckpt")
+                if os.path.isfile(ckpt):
+                    model.load_state_dict(torch.load(ckpt, map_location=DEVICE))
+                result_html = visualize_attention(model, feats, drug_idx)
+
+    return render_template_string(
+    # ───────────── HTML (原 UI + 新输入框) ─────────────
+    """
+<!doctype html>
+<html lang="en"><head><meta charset="utf-8"><title>FusionDTI </title>
+<link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600&family=Poppins:wght@500;600&display=swap" rel="stylesheet">
+
+<style>
+:root{--bg:#f3f4f6;--card:#fff;--primary:#6366f1;--primary-dark:#4f46e5;--text:#111827;--border:#e5e7eb;}
+*{box-sizing:border-box;margin:0;padding:0}
+body{background:var(--bg);color:var(--text);font-family:Inter,system-ui,Arial,sans-serif;line-height:1.5;padding:32px 12px;}
+h1{font-family:Poppins,Inter,sans-serif;font-weight:600;font-size:1.7rem;text-align:center;margin-bottom:28px;letter-spacing:-.2px;}
+.card{max-width:1000px;margin:0 auto;background:var(--card);border:1px solid var(--border);
+      border-radius:12px;box-shadow:0 2px 6px rgba(0,0,0,.05);padding:32px 36px;}
+label{font-weight:500;margin-bottom:6px;display:block}
+textarea,input[type=file]{width:100%;font-size:.9rem;font-family:monospace;padding:10px 12px;
+      border:1px solid var(--border);border-radius:8px;background:#fff;resize:vertical;}
+textarea{min-height:90px}
+.btn{appearance:none;border:none;cursor:pointer;padding:12px 22px;border-radius:8px;font-weight:500;
+     font-family:Inter,sans-serif;transition:all .18s ease;color:#fff;}
+.btn-primary{background:var(--primary)}.btn-primary:hover{background:var(--primary-dark)}
+.btn-neutral{background:#9ca3af;}.btn-neutral:hover{background:#6b7280}
+.grid{display:grid;gap:22px}.grid-2{grid-template-columns:1fr 1fr}
+.vis-box{margin-top:28px;border:1px solid var(--border);border-radius:10px;overflow:auto;max-height:72vh;}
+pre{white-space:pre-wrap;word-break:break-all;font-family:monospace;margin-top:8px}
+
+/* ── tidy table for Top-20 list ─────────────────────────────── */
+table.tg{border-collapse:collapse;margin-top:4px;font-size:0.83rem}
+table.tg th,table.tg td{border:1px solid var(--border);padding:6px 8px;text-align:left}
+table.tg th{background:var(--bg);font-weight:600}
+</style>
+</head>
+<body>
+<h1> Token-level Visualiser for Drug-Target Interaction</h1>
+
+<!-- ───────────── Project Links (larger + spaced) ───────────── -->
+<div style="margin-top:24px; text-align:center;">
+  <a href="https://zhaohanm.github.io/FusionDTI.github.io/" target="_blank"
+     style="display:inline-block;margin:8px 18px;padding:10px 20px;
+            background:linear-gradient(to right,#10b981,#059669);color:white;
+            font-weight:600;border-radius:8px;font-size:0.9rem;
+            font-family:Inter,sans-serif;text-decoration:none;
+            box-shadow:0 2px 6px rgba(0,0,0,0.12);transition:all 0.2s ease-in-out;"
+     onmouseover="this.style.opacity='0.9'" onmouseout="this.style.opacity='1'">
+    🌐 Project Page
+  </a>
+
+  <a href="https://arxiv.org/abs/2406.01651" target="_blank"
+     style="display:inline-block;margin:8px 18px;padding:10px 20px;
+            background:linear-gradient(to right,#ef4444,#dc2626);color:white;
+            font-weight:600;border-radius:8px;font-size:0.9rem;
+            font-family:Inter,sans-serif;text-decoration:none;
+            box-shadow:0 2px 6px rgba(0,0,0,0.12);transition:all 0.2s ease-in-out;"
+     onmouseover="this.style.opacity='0.9'" onmouseout="this.style.opacity='1'">
+    📄 ArXiv: 2406.01651
+  </a>
+
+  <a href="https://github.com/ZhaohanM/FusionDTI" target="_blank"
+     style="display:inline-block;margin:8px 18px;padding:10px 20px;
+            background:linear-gradient(to right,#3b82f6,#2563eb);color:white;
+            font-weight:600;border-radius:8px;font-size:0.9rem;
+            font-family:Inter,sans-serif;text-decoration:none;
+            box-shadow:0 2px 6px rgba(0,0,0,0.12);transition:all 0.2s ease-in-out;"
+     onmouseover="this.style.opacity='0.9'" onmouseout="this.style.opacity='1'">
+    💻 GitHub Repo
+  </a>
+</div>
+
+<!-- ─────────────  Guidelines for Use  ───────────── -->
+<div class="card" style="margin-bottom:24px">
+  <h2 style="font-size:1.2rem;margin-bottom:14px">Guidelines for Use</h2>
+  <ul style="margin-left:18px;line-height:1.55;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
+        <a href="https://github.com/steineggerlab/foldseek" target="_blank">Foldseek</a>,
+        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>
+
+    <li>After inference, you can use the
+        “Download PDF” link to export a high-resolution vector version.</li>
+  </ul>
+</div>
+
+<div class="card">
+<form method="POST" enctype="multipart/form-data" class="grid">
+
+  <div><label>Protein Structure (.pdb / .cif)</label>
+       <input type="file" name="structure_file">
+       <input type="hidden" name="tmp_structure_path" value="{{ tmp_structure_path }}"></div>
+
+  <div><label>Protein Sequence</label>
+       <textarea name="protein_sequence" placeholder="Confirm / paste sequence…">{{ protein_seq }}</textarea></div>
+
+  <div><label>Drug Sequence (SELFIES/SMILES)</label>
+       <textarea name="drug_sequence" placeholder="[C][C][O]/cco …">{{ drug_seq }}</textarea></div>
+
+    <label>Drug atom/substructure index (1-based) – show Top-10 related protein residue</label>
+        <input type="number" name="drug_idx" min="1" style="width:120px">
+
+  <div class="grid grid-2">
+    <button class="btn btn-primary" type="Inference" name="confirm_structure">Confirm Structure</button>
+    <button class="btn btn-primary" type="Inference" name="Inference">Inference</button>
+  </div>
+  <button class="btn btn-neutral" style="width:100%" type="Inference" name="clear">Clear</button>
+</form>
+
+{% if structure_seq %}
+  <div style="margin-top:18px"><strong>Structure-aware sequence:</strong><pre>{{ structure_seq }}</pre></div>
+{% endif %}
+{% if result_html %}
+  <div class="vis-box" style="margin-top:26px">{{ result_html|safe }}</div>
+{% endif %}
+</div></body></html>
+    """,
+    protein_seq=protein_seq, drug_seq=drug_seq, structure_seq=structure_seq,
+    result_html=result_html, tmp_structure_path=tmp_structure_path)
 
-if __name__ == '__main__':
+# ───── run ─────────────────────────────────────────────────────
+if __name__ == "__main__":
     app.run(debug=True, host="0.0.0.0", port=7860)

requirements.txt

@@ -1,5 +1,11 @@
 Flask
 torch
 transformers
-bertviz
-IPython
+IPython
+selfies
+rdkit
+biopython
+matplotlib
+scikit-learn
+numpy
+pandas

utils/.ipynb_checkpoints/drug_tokenizer-checkpoint.py

@@ -0,0 +1,73 @@
+import json
+import re
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class DrugTokenizer:
+    def __init__(self, vocab_path="data/Tokenizer/vocab.json", special_tokens_path="data/Tokenizer/special_tokens_map.json"):
+        self.vocab, self.special_tokens = self.load_vocab_and_special_tokens(vocab_path, special_tokens_path)
+        self.cls_token_id = self.vocab[self.special_tokens['cls_token']]
+        self.sep_token_id = self.vocab[self.special_tokens['sep_token']]
+        self.unk_token_id = self.vocab[self.special_tokens['unk_token']]
+        self.pad_token_id = self.vocab[self.special_tokens['pad_token']]
+        self.id_to_token = {v: k for k, v in self.vocab.items()}
+
+        self.all_special_tokens = list(self.special_tokens.values())
+
+    def load_vocab_and_special_tokens(self, vocab_path, special_tokens_path):
+        with open(vocab_path, 'r', encoding='utf-8') as vocab_file:
+            vocab = json.load(vocab_file)
+        with open(special_tokens_path, 'r', encoding='utf-8') as special_tokens_file:
+            special_tokens_raw = json.load(special_tokens_file)
+
+        special_tokens = {key: value['content'] for key, value in special_tokens_raw.items()}
+        return vocab, special_tokens
+
+    def encode(self, sequence):
+        tokens = re.findall(r'\[([^\[\]]+)\]', sequence)
+        input_ids = [self.cls_token_id] + [self.vocab.get(token, self.unk_token_id) for token in tokens] + [self.sep_token_id]
+        attention_mask = [1] * len(input_ids)
+        return {
+            'input_ids': input_ids,
+            'attention_mask': attention_mask
+        }
+
+    def batch_encode_plus(self, sequences, max_length, padding, truncation, add_special_tokens, return_tensors):
+        input_ids_list = []
+        attention_mask_list = []
+
+        for sequence in sequences:
+            encoded = self.encode(sequence)
+            input_ids = encoded['input_ids']
+            attention_mask = encoded['attention_mask']
+
+            if len(input_ids) > max_length:
+                input_ids = input_ids[:max_length]
+                attention_mask = attention_mask[:max_length]
+            elif len(input_ids) < max_length:
+                pad_length = max_length - len(input_ids)
+                input_ids = input_ids + [self.vocab[self.special_tokens['pad_token']]] * pad_length
+                attention_mask = attention_mask + [0] * pad_length
+
+            input_ids_list.append(input_ids)
+            attention_mask_list.append(attention_mask)
+
+        return {
+            'input_ids': torch.tensor(input_ids_list, dtype=torch.long),
+            'attention_mask': torch.tensor(attention_mask_list, dtype=torch.long)
+        }
+
+    def decode(self, input_ids, skip_special_tokens=False):
+        tokens = []
+        for id in input_ids:
+            if skip_special_tokens and id in [self.cls_token_id, self.sep_token_id, self.pad_token_id]:
+                continue
+            tokens.append(self.id_to_token.get(id, self.special_tokens['unk_token']))
+        sequence = ''.join([f'[{token}]' for token in tokens])
+        return sequence
+
+        # --- 新增 ---
+    def convert_ids_to_tokens(self, ids):
+        """list[int] → list[str]，跳过未知 id"""
+        return [self.id_to_token.get(i, self.special_tokens['unk_token']) for i in ids]

utils/.ipynb_checkpoints/metric_learning_models_att_maps-checkpoint.py

@@ -0,0 +1,325 @@
+import logging
+import os
+import sys
+
+sys.path.append("../")
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.cuda.amp import autocast
+from torch.nn import Module
+from tqdm import tqdm
+from torch.nn.utils.weight_norm import weight_norm
+from torch.utils.data import Dataset
+
+LOGGER = logging.getLogger(__name__)
+
+class FusionDTI(nn.Module):
+    def __init__(self, prot_out_dim, disease_out_dim, args):
+        super(FusionDTI, self).__init__()
+        self.fusion = args.fusion
+        self.drug_reg = nn.Linear(disease_out_dim, 512)
+        self.prot_reg = nn.Linear(prot_out_dim, 512)
+
+        if self.fusion == "CAN":
+            self.can_layer = CAN_Layer(hidden_dim=512, num_heads=8, args=args)
+            self.mlp_classifier = MlPdecoder_CAN(input_dim=1024)
+        elif self.fusion == "BAN":
+            self.ban_layer = weight_norm(BANLayer(512, 512, 256, 2), name='h_mat', dim=None)
+            self.mlp_classifier = MlPdecoder_CAN(input_dim=256)
+        elif self.fusion == "Nan":
+            self.mlp_classifier_nan = MlPdecoder_CAN(input_dim=1214)
+
+    def forward(self, prot_embed, drug_embed, prot_mask, drug_mask):
+        # print("drug_embed", drug_embed.shape)
+        if self.fusion == "Nan":
+            prot_embed = prot_embed.mean(1)  # query : [batch_size, hidden]
+            drug_embed = drug_embed.mean(1)  # query : [batch_size, hidden]
+            joint_embed = torch.cat([prot_embed, drug_embed], dim=1)
+            score = self.mlp_classifier_nan(joint_embed)
+        else:
+            prot_embed = self.prot_reg(prot_embed)
+            drug_embed = self.drug_reg(drug_embed)
+
+            if self.fusion == "CAN":
+                joint_embed, att = self.can_layer(prot_embed, drug_embed, prot_mask, drug_mask)
+            elif self.fusion == "BAN":
+                joint_embed, att = self.ban_layer(prot_embed, drug_embed)
+
+            score = self.mlp_classifier(joint_embed)
+
+        return score, att
+    
+class Pre_encoded(nn.Module):
+    def __init__(
+            self, prot_encoder, drug_encoder, args
+    ):
+        """Constructor for the model.
+
+        Args:
+            prot_encoder (_type_): Protein sturcture-aware sequence encoder.
+            drug_encoder (_type_): Drug SFLFIES encoder.
+            args (_type_): _description_
+        """
+        super(Pre_encoded, self).__init__()
+        self.prot_encoder = prot_encoder
+        self.drug_encoder = drug_encoder
+        
+    def encoding(self, prot_input_ids, prot_attention_mask, drug_input_ids, drug_attention_mask):
+        # Process inputs through encoders
+        prot_embed = self.prot_encoder(
+            input_ids=prot_input_ids, attention_mask=prot_attention_mask, return_dict=True
+        ).logits
+        # prot_embed = self.prot_reg(prot_embed)
+
+        drug_embed = self.drug_encoder(
+            input_ids=drug_input_ids, attention_mask=drug_attention_mask, return_dict=True
+        ).last_hidden_state  # .last_hidden_state
+        
+        # print("drug_embed", drug_embed.shape)
+        
+        return prot_embed, drug_embed
+    
+    
+class CAN_Layer(nn.Module):
+    def __init__(self, hidden_dim, num_heads, args):
+        super(CAN_Layer, self).__init__()
+        self.agg_mode = args.agg_mode
+        self.group_size = args.group_size  #  Control Fusion Scale
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.head_size = hidden_dim // num_heads
+
+        self.query_p = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.key_p = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.value_p = nn.Linear(hidden_dim, hidden_dim, bias=False)
+
+        self.query_d = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.key_d = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.value_d = nn.Linear(hidden_dim, hidden_dim, bias=False)
+
+    def alpha_logits(self, logits, mask_row, mask_col, inf=1e6):
+        N, L1, L2, H = logits.shape
+        mask_row = mask_row.view(N, L1, 1).repeat(1, 1, H)
+        mask_col = mask_col.view(N, L2, 1).repeat(1, 1, H)
+        mask_pair = torch.einsum('blh, bkh->blkh', mask_row, mask_col)
+
+        logits = torch.where(mask_pair, logits, logits - inf)
+        alpha = torch.softmax(logits, dim=2)
+        mask_row = mask_row.view(N, L1, 1, H).repeat(1, 1, L2, 1)
+        alpha = torch.where(mask_row, alpha, torch.zeros_like(alpha))
+        return alpha
+
+    def apply_heads(self, x, n_heads, n_ch):
+        s = list(x.size())[:-1] + [n_heads, n_ch]
+        return x.view(*s)
+
+    def group_embeddings(self, x, mask, group_size):
+        N, L, D = x.shape
+        groups = L // group_size
+        x_grouped = x.view(N, groups, group_size, D).mean(dim=2)
+        mask_grouped = mask.view(N, groups, group_size).any(dim=2)
+        return x_grouped, mask_grouped
+
+    def forward(self, protein, drug, mask_prot, mask_drug):
+        # Group embeddings before applying multi-head attention
+        protein_grouped, mask_prot_grouped = self.group_embeddings(protein, mask_prot, self.group_size)
+        drug_grouped, mask_drug_grouped = self.group_embeddings(drug, mask_drug, self.group_size)
+        
+        # print("protein_grouped:", protein_grouped.shape)
+        # print("mask_prot_grouped:", mask_prot_grouped.shape)
+
+        # Compute queries, keys, values for both protein and drug after grouping
+        query_prot = self.apply_heads(self.query_p(protein_grouped), self.num_heads, self.head_size)
+        key_prot = self.apply_heads(self.key_p(protein_grouped), self.num_heads, self.head_size)
+        value_prot = self.apply_heads(self.value_p(protein_grouped), self.num_heads, self.head_size)
+
+        query_drug = self.apply_heads(self.query_d(drug_grouped), self.num_heads, self.head_size)
+        key_drug = self.apply_heads(self.key_d(drug_grouped), self.num_heads, self.head_size)
+        value_drug = self.apply_heads(self.value_d(drug_grouped), self.num_heads, self.head_size)
+
+        # Compute attention scores
+        logits_pp = torch.einsum('blhd, bkhd->blkh', query_prot, key_prot)
+        logits_pd = torch.einsum('blhd, bkhd->blkh', query_prot, key_drug)
+        logits_dp = torch.einsum('blhd, bkhd->blkh', query_drug, key_prot)
+        logits_dd = torch.einsum('blhd, bkhd->blkh', query_drug, key_drug)
+        # print("logits_pp:", logits_pp.shape)
+
+        alpha_pp = self.alpha_logits(logits_pp, mask_prot_grouped, mask_prot_grouped)
+        alpha_pd = self.alpha_logits(logits_pd, mask_prot_grouped, mask_drug_grouped)
+        alpha_dp = self.alpha_logits(logits_dp, mask_drug_grouped, mask_prot_grouped)
+        alpha_dd = self.alpha_logits(logits_dd, mask_drug_grouped, mask_drug_grouped)
+
+        prot_embedding = (torch.einsum('blkh, bkhd->blhd', alpha_pp, value_prot).flatten(-2) +
+                   torch.einsum('blkh, bkhd->blhd', alpha_pd, value_drug).flatten(-2)) / 2
+        drug_embedding = (torch.einsum('blkh, bkhd->blhd', alpha_dp, value_prot).flatten(-2) +
+                   torch.einsum('blkh, bkhd->blhd', alpha_dd, value_drug).flatten(-2)) / 2
+        
+        # print("prot_embedding:", prot_embedding.shape)
+        
+        # Continue as usual with the aggregation mode
+        if self.agg_mode == "cls":
+            prot_embed = prot_embedding[:, 0]  # query : [batch_size, hidden]
+            drug_embed = drug_embedding[:, 0]  # query : [batch_size, hidden]
+        elif self.agg_mode == "mean_all_tok":
+            prot_embed = prot_embedding.mean(1)  # query : [batch_size, hidden]
+            drug_embed = drug_embedding.mean(1)  # query : [batch_size, hidden]
+        elif self.agg_mode == "mean":
+            prot_embed = (prot_embedding * mask_prot_grouped.unsqueeze(-1)).sum(1) / mask_prot_grouped.sum(-1).unsqueeze(-1)
+            drug_embed = (drug_embedding * mask_drug_grouped.unsqueeze(-1)).sum(1) / mask_drug_grouped.sum(-1).unsqueeze(-1)
+        else:
+            raise NotImplementedError()
+            
+        # print("prot_embed:", prot_embed.shape)
+
+        query_embed = torch.cat([prot_embed, drug_embed], dim=1)
+        
+        att_pd = alpha_pd.mean(dim=-1)
+
+        # print("query_embed:", query_embed.shape)
+        return query_embed, att_pd
+
+class MlPdecoder_CAN(nn.Module):
+    def __init__(self, input_dim):
+        super(MlPdecoder_CAN, self).__init__()
+        self.fc1 = nn.Linear(input_dim, input_dim)
+        self.bn1 = nn.BatchNorm1d(input_dim)
+        self.fc2 = nn.Linear(input_dim, input_dim // 2)
+        self.bn2 = nn.BatchNorm1d(input_dim // 2)
+        self.fc3 = nn.Linear(input_dim // 2, input_dim // 4)
+        self.bn3 = nn.BatchNorm1d(input_dim // 4)
+        self.output = nn.Linear(input_dim // 4, 1)
+
+    def forward(self, x):
+        x = self.bn1(torch.relu(self.fc1(x)))
+        x = self.bn2(torch.relu(self.fc2(x)))
+        x = self.bn3(torch.relu(self.fc3(x)))
+        x = torch.sigmoid(self.output(x))
+        return x
+    
+class MLPdecoder_BAN(nn.Module):
+    def __init__(self, in_dim, hidden_dim, out_dim, binary=1):
+        super(MLPdecoder_BAN, self).__init__()
+        self.fc1 = nn.Linear(in_dim, hidden_dim)
+        self.bn1 = nn.BatchNorm1d(hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.bn2 = nn.BatchNorm1d(hidden_dim)
+        self.fc3 = nn.Linear(hidden_dim, out_dim)
+        self.bn3 = nn.BatchNorm1d(out_dim)
+        self.fc4 = nn.Linear(out_dim, binary)
+
+    def forward(self, x):
+        x = self.bn1(F.relu(self.fc1(x)))
+        x = self.bn2(F.relu(self.fc2(x)))
+        x = self.bn3(F.relu(self.fc3(x)))
+        # x = self.fc4(x)
+        x = torch.sigmoid(self.fc4(x)) 
+        return x
+
+class BANLayer(nn.Module):
+    """ Bilinear attention network
+    Modified from https://github.com/peizhenbai/DrugBAN/blob/main/ban.py
+    """
+    def __init__(self, v_dim, q_dim, h_dim, h_out, act='ReLU', dropout=0.2, k=3):
+        super(BANLayer, self).__init__()
+
+        self.c = 32
+        self.k = k
+        self.v_dim = v_dim
+        self.q_dim = q_dim
+        self.h_dim = h_dim
+        self.h_out = h_out
+
+        self.v_net = FCNet([v_dim, h_dim * self.k], act=act, dropout=dropout)
+        self.q_net = FCNet([q_dim, h_dim * self.k], act=act, dropout=dropout)
+        # self.dropout = nn.Dropout(dropout[1])
+        if 1 < k:
+            self.p_net = nn.AvgPool1d(self.k, stride=self.k)
+
+        if h_out <= self.c:
+            self.h_mat = nn.Parameter(torch.Tensor(1, h_out, 1, h_dim * self.k).normal_())
+            self.h_bias = nn.Parameter(torch.Tensor(1, h_out, 1, 1).normal_())
+        else:
+            self.h_net = weight_norm(nn.Linear(h_dim * self.k, h_out), dim=None)
+
+        self.bn = nn.BatchNorm1d(h_dim)
+
+    def attention_pooling(self, v, q, att_map):
+        fusion_logits = torch.einsum('bvk,bvq,bqk->bk', (v, att_map, q))
+        if 1 < self.k:
+            fusion_logits = fusion_logits.unsqueeze(1)  # b x 1 x d
+            fusion_logits = self.p_net(fusion_logits).squeeze(1) * self.k  # sum-pooling
+        return fusion_logits
+
+    def forward(self, v, q, softmax=False):
+        v_num = v.size(1)
+        q_num = q.size(1)
+        # print("v_num", v_num)
+        # print("v_num ", v_num)
+        if self.h_out <= self.c:
+            v_ = self.v_net(v)
+            q_ = self.q_net(q)
+            # print("v_", v_.shape)
+            # print("q_ ", q_.shape)
+            att_maps = torch.einsum('xhyk,bvk,bqk->bhvq', (self.h_mat, v_, q_)) + self.h_bias
+            # print("Attention map_1",att_maps.shape)
+        else:
+            v_ = self.v_net(v).transpose(1, 2).unsqueeze(3)
+            q_ = self.q_net(q).transpose(1, 2).unsqueeze(2)
+            d_ = torch.matmul(v_, q_)  # b x h_dim x v x q
+            att_maps = self.h_net(d_.transpose(1, 2).transpose(2, 3))  # b x v x q x h_out
+            att_maps = att_maps.transpose(2, 3).transpose(1, 2)  # b x h_out x v x q
+            # print("Attention map_2",att_maps.shape)
+        if softmax:
+            p = nn.functional.softmax(att_maps.view(-1, self.h_out, v_num * q_num), 2)
+            att_maps = p.view(-1, self.h_out, v_num, q_num)
+            # print("Attention map_softmax", att_maps.shape)
+        logits = self.attention_pooling(v_, q_, att_maps[:, 0, :, :])
+        for i in range(1, self.h_out):
+            logits_i = self.attention_pooling(v_, q_, att_maps[:, i, :, :])
+            logits += logits_i
+        logits = self.bn(logits)
+        return logits, att_maps
+
+    
+class FCNet(nn.Module):
+    """Simple class for non-linear fully connect network
+    Modified from https://github.com/jnhwkim/ban-vqa/blob/master/fc.py
+    """
+
+    def __init__(self, dims, act='ReLU', dropout=0):
+        super(FCNet, self).__init__()
+
+        layers = []
+        for i in range(len(dims) - 2):
+            in_dim = dims[i]
+            out_dim = dims[i + 1]
+            if 0 < dropout:
+                layers.append(nn.Dropout(dropout))
+            layers.append(weight_norm(nn.Linear(in_dim, out_dim), dim=None))
+            if '' != act:
+                layers.append(getattr(nn, act)())
+        if 0 < dropout:
+            layers.append(nn.Dropout(dropout))
+        layers.append(weight_norm(nn.Linear(dims[-2], dims[-1]), dim=None))
+        if '' != act:
+            layers.append(getattr(nn, act)())
+
+        self.main = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.main(x)
+    
+
+class BatchFileDataset_Case(Dataset):
+    def __init__(self, file_list):
+        self.file_list = file_list
+
+    def __len__(self):
+        return len(self.file_list)
+
+    def __getitem__(self, idx):
+        batch_file = self.file_list[idx]
+        data = torch.load(batch_file)
+        return data['prot'], data['drug'], data['prot_ids'], data['drug_ids'], data['prot_mask'], data['drug_mask'], data['y']

utils/drug_tokenizer.py

@@ -5,7 +5,7 @@ import torch.nn as nn
 from torch.nn import functional as F
 
 class DrugTokenizer:
-    def __init__(self, vocab_path="tokenizer/vocab.json", special_tokens_path="tokenizer/special_tokens_map.json"):
+    def __init__(self, vocab_path="data/Tokenizer/vocab.json", special_tokens_path="data/Tokenizer/special_tokens_map.json"):
         self.vocab, self.special_tokens = self.load_vocab_and_special_tokens(vocab_path, special_tokens_path)
         self.cls_token_id = self.vocab[self.special_tokens['cls_token']]
         self.sep_token_id = self.vocab[self.special_tokens['sep_token']]
@@ -13,6 +13,8 @@ class DrugTokenizer:
         self.pad_token_id = self.vocab[self.special_tokens['pad_token']]
         self.id_to_token = {v: k for k, v in self.vocab.items()}
 
+        self.all_special_tokens = list(self.special_tokens.values())
+
     def load_vocab_and_special_tokens(self, vocab_path, special_tokens_path):
         with open(vocab_path, 'r', encoding='utf-8') as vocab_file:
             vocab = json.load(vocab_file)
@@ -64,3 +66,8 @@ class DrugTokenizer:
             tokens.append(self.id_to_token.get(id, self.special_tokens['unk_token']))
         sequence = ''.join([f'[{token}]' for token in tokens])
         return sequence
+
+        # --- 新增 ---
+    def convert_ids_to_tokens(self, ids):
+        """list[int] → list[str]，跳过未知 id"""
+        return [self.id_to_token.get(i, self.special_tokens['unk_token']) for i in ids]

utils/foldseek_util.py

@@ -0,0 +1,167 @@
+import os
+import time
+import json
+import numpy as np
+import re
+import sys
+
+from Bio.PDB import PDBParser, MMCIFParser
+
+
+sys.path.append(".")
+
+
+# Get structural seqs from pdb file
+def get_struc_seq(foldseek,
+                  path,
+                  chains: list = None,
+                  process_id: int = 0,
+                  plddt_mask: bool = "auto",
+                  plddt_threshold: float = 70.,
+                  foldseek_verbose: bool = False) -> dict:
+    """
+
+    Args:
+        foldseek: Binary executable file of foldseek
+
+        path: Path to pdb file
+
+        chains: Chains to be extracted from pdb file. If None, all chains will be extracted.
+
+        process_id: Process ID for temporary files. This is used for parallel processing.
+
+        plddt_mask: If True, mask regions with plddt < plddt_threshold. plddt scores are from the pdb file.
+
+        plddt_threshold: Threshold for plddt. If plddt is lower than this value, the structure will be masked.
+
+        foldseek_verbose: If True, foldseek will print verbose messages.
+
+    Returns:
+        seq_dict: A dict of structural seqs. The keys are chain IDs. The values are tuples of
+        (seq, struc_seq, combined_seq).
+    """
+    assert os.path.exists(foldseek), f"Foldseek not found: {foldseek}"
+    assert os.path.exists(path), f"PDB file not found: {path}"
+    
+    tmp_save_path = f"get_struc_seq_{process_id}_{time.time()}.tsv"
+    if foldseek_verbose:
+        cmd = f"{foldseek} structureto3didescriptor --threads 1 --chain-name-mode 1 {path} {tmp_save_path}"
+    else:
+        cmd = f"{foldseek} structureto3didescriptor -v 0 --threads 1 --chain-name-mode 1 {path} {tmp_save_path}"
+    os.system(cmd)
+    
+    # Check whether the structure is predicted by AlphaFold2
+    if plddt_mask == "auto":
+        with open(path, "r") as r:
+            plddt_mask = True if "alphafold" in r.read().lower() else False
+    
+    seq_dict = {}
+    name = os.path.basename(path)
+    with open(tmp_save_path, "r") as r:
+        for i, line in enumerate(r):
+            desc, seq, struc_seq = line.split("\t")[:3]
+            
+            # Mask low plddt
+            if plddt_mask:
+                try:
+                    plddts = extract_plddt(path)
+                    assert len(plddts) == len(struc_seq), f"Length mismatch: {len(plddts)} != {len(struc_seq)}"
+                    
+                    # Mask regions with plddt < threshold
+                    indices = np.where(plddts < plddt_threshold)[0]
+                    np_seq = np.array(list(struc_seq))
+                    np_seq[indices] = "#"
+                    struc_seq = "".join(np_seq)
+                
+                except Exception as e:
+                    print(f"Error: {e}")
+                    print(f"Failed to mask plddt for {name}")
+            
+            name_chain = desc.split(" ")[0]
+            chain = name_chain.replace(name, "").split("_")[-1]
+            
+            if chains is None or chain in chains:
+                if chain not in seq_dict:
+                    combined_seq = "".join([a + b.lower() for a, b in zip(seq, struc_seq)])
+                    seq_dict[chain] = (seq, struc_seq, combined_seq)
+    
+    os.remove(tmp_save_path)
+    os.remove(tmp_save_path + ".dbtype")
+    return seq_dict
+
+
+def extract_plddt(pdb_path: str) -> np.ndarray:
+    """
+    Extract plddt scores from pdb file.
+    Args:
+        pdb_path: Path to pdb file.
+
+    Returns:
+        plddts: plddt scores.
+    """
+
+    # Initialize parser
+    if pdb_path.endswith(".cif"):
+        parser = MMCIFParser()
+    elif pdb_path.endswith(".pdb"):
+        parser = PDBParser()
+    else:
+        raise ValueError("Invalid file format for plddt extraction. Must be '.cif' or '.pdb'.")
+    
+    structure = parser.get_structure('protein', pdb_path)
+    model = structure[0]
+    chain = model["A"]
+
+    # Extract plddt scores
+    plddts = []
+    for residue in chain:
+        residue_plddts = []
+        for atom in residue:
+            plddt = atom.get_bfactor()
+            residue_plddts.append(plddt)
+        
+        plddts.append(np.mean(residue_plddts))
+
+    plddts = np.array(plddts)
+    return plddts
+
+
+def transform_pdb_dir(foldseek: str, pdb_dir: str, seq_type: str, save_path: str):
+    """
+    Transform a directory of pdb files into a fasta file.
+    Args:
+        foldseek: Binary executable file of foldseek.
+        
+        pdb_dir: Directory of pdb files.
+        
+        seq_type: Type of sequence to be extracted. Must be "aa" or "foldseek"
+        
+        save_path: Path to save the fasta file.
+    """
+    assert os.path.exists(foldseek), f"Foldseek not found: {foldseek}"
+    assert seq_type in ["aa", "foldseek"], f"seq_type must be 'aa' or 'foldseek'!"
+    
+    tmp_save_path = f"get_struc_seq_{time.time()}.tsv"
+    cmd = f"{foldseek} structureto3didescriptor --chain-name-mode 1 {pdb_dir} {tmp_save_path}"
+    os.system(cmd)
+    
+    with open(tmp_save_path, "r") as r, open(save_path, "w") as w:
+        for line in r:
+            protein_id, aa_seq, foldseek_seq = line.strip().split("\t")[:3]
+            
+            if seq_type == "aa":
+                w.write(f">{protein_id}\n{aa_seq}\n")
+            else:
+                w.write(f">{protein_id}\n{foldseek_seq.lower()}\n")
+    
+    os.remove(tmp_save_path)
+    os.remove(tmp_save_path + ".dbtype")
+    
+
+if __name__ == '__main__':
+    foldseek = "/sujin/bin/foldseek"
+    # test_path = "/sujin/Datasets/PDB/all/6xtd.cif"
+    test_path = "/sujin/Datasets/FLIP/meltome/af2_structures/A0A061ACX4.pdb"
+    plddt_path = "/sujin/Datasets/FLIP/meltome/af2_plddts/A0A061ACX4.json"
+    res = get_struc_seq(foldseek, test_path, plddt_path=plddt_path, plddt_threshold=70.)
+    print(res["A"][1].lower())

utils/metric_learning_models_att_maps.py

@@ -175,15 +175,10 @@ class CAN_Layer(nn.Module):
 
         query_embed = torch.cat([prot_embed, drug_embed], dim=1)
         
-
-        att = torch.zeros(1, 1, 1024, 1024)
-        att[:, :, :512, :512] = alpha_pp.mean(dim=-1)  # Protein to Protein
-        att[:, :, :512, 512:] = alpha_pd.mean(dim=-1)  # Protein to Drug
-        att[:, :, 512:, :512] = alpha_dp.mean(dim=-1)  # Drug to Protein
-        att[:, :, 512:, 512:] = alpha_dd.mean(dim=-1)  # Drug to Drug
+        att_pd = alpha_pd.mean(dim=-1)
 
         # print("query_embed:", query_embed.shape)
-        return query_embed, att
+        return query_embed, att_pd
 
 class MlPdecoder_CAN(nn.Module):
     def __init__(self, input_dim):