Add ExplainBind Demo

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+bin/foldseek filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+.DS_Store
+.ipynb_checkpoints/
+__pycache__/
+*.pyc
+bin/
+*.log

README.md

@@ -8,7 +8,8 @@ sdk_version: 6.6.0
 app_file: app.py
 pinned: false
 license: mit
-short_description: An interaction-aware demo UI
+short_description: **ExplainBind** is an interaction-aware framework for **protein–ligand binding (PLB)** prediction.  
+It supervises token-level cross-attention using **non-covalent interaction maps** (e.g. hydrogen bonds, salt bridges, hydrophobic contacts, van der Waals, π–π, and cation–π interactions) derived from curated **PDB** protein–ligand complexes in **InteractBind**. By aligning model attention with these physically grounded signals, ExplainBind transforms PLB prediction from a black-box reasoning into an **chemistry-grounded** process suitable for large-scale screening.
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,1246 @@
+# ─── Patch gradio_client so boolean schemas don’t crash json_schema_to_python_type ───
+import gradio_client.utils as _gc_utils
+_orig_get_type   = _gc_utils.get_type
+_orig_json2py    = _gc_utils._json_schema_to_python_type
+def _patched_get_type(schema):
+    if isinstance(schema, bool):
+        schema = {}
+    return _orig_get_type(schema)
+def _patched_json_schema_to_python_type(schema, defs=None):
+    if isinstance(schema, bool):
+        schema = {}
+    return _orig_json2py(schema, defs)
+_gc_utils.get_type                    = _patched_get_type
+_gc_utils._json_schema_to_python_type = _patched_json_schema_to_python_type
+
+# ─── Imports ───────────────────────────────────────────────────────────────────
+import os
+import io
+import base64
+import argparse
+from typing import Optional, List, Tuple
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+
+import selfies
+from rdkit import Chem
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from matplotlib import cm
+
+from transformers import EsmForMaskedLM, EsmTokenizer, AutoModel, AutoTokenizer
+from Bio.PDB import PDBParser, MMCIFParser
+from Bio.Data import IUPACData
+
+import gradio as gr
+
+# Project utils (ensure these exist in your repository)
+from utils.metric_learning_models_att_maps import Pre_encoded, ExplainBind
+from utils.foldseek_util import get_struc_seq
+
+# ───────────────────── Paths & Logos ─────────────────────
+ROOT = os.path.dirname(os.path.abspath(__file__))
+ASSET_DIR = os.path.join(ROOT, "assets")
+
+LOSCAZLO_LOGO = os.path.join(ASSET_DIR, "loscalzo.png")
+
+def _load_logo_b64(path):
+    if not os.path.exists(path):
+        return ""
+    with open(path, "rb") as f:
+        return base64.b64encode(f.read()).decode("utf-8")
+
+LOSCAZLO_B64 = _load_logo_b64(LOSCAZLO_LOGO)
+
+
+# ───────────────────── Configurable constants ─────────────────────
+# UI-visible names (Halogen bonding removed)
+INTERACTION_NAMES = [
+    "Hydrogen bonding",
+    "Salt Bridging",
+    "π–π Stacking",
+    "Cation–π",
+    "Hydrophobic",
+    "Van der Waals",
+    "Overall Interaction",
+]
+
+# Map visible indices (0..5 = specific, 6 = combined) to underlying channel indices
+# Underlying channels originally had Halogen at index=5 (0-based). We skip 5 entirely.
+VISIBLE2UNDERLYING = [1, 2, 3, 4, 6, 0]  # HB, Salt, Pi, Cation-Pi, Hydro, VdW
+N_VISIBLE_SPEC = len(VISIBLE2UNDERLYING)  # 6
+
+# ───── Helper utilities ───────────────────────────────────────────
+three2one = {k.upper(): v for k, v in IUPACData.protein_letters_3to1.items()}
+three2one.update({"MSE": "M", "SEC": "C", "PYL": "K"})
+STANDARD_AA_SET = set("ACDEFGHIKLMNPQRSTVWY")  # Uppercase FASTA amino acids
+
+
+def simple_seq_from_structure(path: str) -> str:
+    """Extract the longest chain from a PDB/mmCIF file and return a simple 1-letter sequence."""
+    parser = MMCIFParser(QUIET=True) if path.endswith(".cif") else PDBParser(QUIET=True)
+    structure = parser.get_structure("P", path)
+    chains = list(structure.get_chains())
+    if not chains:
+        return ""
+    chain = max(chains, key=lambda c: len(list(c.get_residues())))
+    return "".join(three2one.get(res.get_resname().upper(), "X") for res in chain)
+
+
+def smiles_to_selfies(smiles_text: str) -> Optional[str]:
+    """Validate and convert SMILES to SELFIES; return None if invalid."""
+    try:
+        mol = Chem.MolFromSmiles(smiles_text)
+        if mol is None:
+            return None
+        return selfies.encoder(smiles_text)
+    except Exception:
+        return None
+
+
+def detect_protein_type(seq: str) -> str:
+    """
+    Heuristic for protein input:
+    - All uppercase and only the standard 20 amino acids → 'fasta'
+    - Otherwise (contains lowercase or non-standard characters) → 'sa'
+    """
+    s = (seq or "").strip()
+    if not s:
+        return "fasta"
+    up = s.upper()
+    only_aa = all(ch in STANDARD_AA_SET for ch in up)
+    all_upper = (s == up)
+    return "fasta" if (only_aa and all_upper) else "sa"
+
+
+def detect_ligand_type(text: str) -> str:
+    """
+    Heuristic for ligand input:
+    - Starts with '[' and contains ']' → 'selfies'
+    - Otherwise → 'smiles'
+    """
+    t = (text or "").strip()
+    if not t:
+        return "smiles"
+    return "selfies" if (t.startswith("[") and ("]" in t)) else "smiles"
+
+
+def parse_config():
+    """Parse command-line options."""
+    p = argparse.ArgumentParser()
+    p.add_argument("--agg_mode", type=str, default="mean_all_tok")
+    p.add_argument("--group_size", type=int, default=1)
+    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/")  # Root folder containing checkpoints
+    p.add_argument("--dataset", default="Human")
+    return p.parse_args()
+
+
+args = parse_config()
+DEVICE = args.device
+
+# ───── Dynamic model registry ─────────────────────────────────────
+PROT_MODELS = {
+    "sa":    "westlake-repl/SaProt_650M_AF2",
+    "fasta": "facebook/esm2_t33_650M_UR50D",
+}
+DRUG_MODELS = {
+    "selfies": "HUBioDataLab/SELFormer",
+    # "smiles":  "ibm/MoLFormer-XL-both-10pct",
+}
+
+
+def load_encoders(ptype: str, ltype: str, args):
+    """
+    Dynamically load encoders and tokenisers based on input types.
+    Returns: (prot_tokenizer, prot_model, drug_tokenizer, drug_model, encoding_module)
+    """
+    # Protein encoder
+    if ptype == "fasta":
+        prot_path = PROT_MODELS["fasta"]
+        prot_tokenizer = EsmTokenizer.from_pretrained(prot_path, do_lower_case=False)
+        prot_model     = EsmForMaskedLM.from_pretrained(prot_path)
+    else:  # 'sa'
+        prot_path = PROT_MODELS["sa"]
+        prot_tokenizer = EsmTokenizer.from_pretrained(prot_path)
+        prot_model     = EsmForMaskedLM.from_pretrained(prot_path)
+
+    drug_path = DRUG_MODELS["selfies"]
+    drug_tokenizer = AutoTokenizer.from_pretrained(drug_path)
+    drug_model     = AutoModel.from_pretrained(drug_path)
+    # Ligand encoder
+    # if ltype == "smiles":
+    #     drug_path = DRUG_MODELS["smiles"]
+    #     drug_tokenizer = AutoTokenizer.from_pretrained(drug_path, trust_remote_code=True)
+    #     drug_model     = AutoModel.from_pretrained(drug_path, deterministic_eval=True, trust_remote_code=True)
+    # else:  # 'selfies'
+    #     drug_path = DRUG_MODELS["selfies"]
+    #     drug_tokenizer = AutoTokenizer.from_pretrained(drug_path)
+    #     drug_model     = AutoModel.from_pretrained(drug_path)
+
+    # Wrap encoders with Pre_encoded module
+    encoding = Pre_encoded(prot_model, drug_model, args).to(DEVICE)
+    return prot_tokenizer, prot_model, drug_tokenizer, drug_model, encoding
+
+
+def make_collate_fn(prot_tokenizer, drug_tokenizer):
+    """Create a batch collation function using the given tokenisers."""
+    def _collate_fn(batch):
+        query1, query2, scores = zip(*batch)
+        query_encodings1 = prot_tokenizer(
+            list(query1), max_length=512, padding="max_length", truncation=True,
+            add_special_tokens=True, return_tensors="pt",
+        )
+        query_encodings2 = drug_tokenizer(
+            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)
+    return _collate_fn
+
+
+def get_case_feature(model, loader):
+    """Generate features for one protein–ligand pair using the provided model."""
+    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)]
+
+# ─────────────── SELFIES grouping by ORIGINAL string ─────────────
+def _group_rows_by_selfies_string(n_rows: int, selfies_str: str):
+    """
+    Partition the attention matrix's n_rows along ligand axis into groups per SELFIES token '[ ... ]'.
+    Each group is a contiguous row span; we assign rows ≈ equally using linspace.
+    Returns:
+        groups: List[(start_row, end_row)] inclusive
+        labels: List['[X]','[=O]', ...]
+    """
+    if n_rows <= 0:
+        return [], []
+
+    try:
+        toks = list(selfies.split_selfies((selfies_str or "").strip()))
+    except Exception:
+        toks = []
+
+    if not toks:
+        # Fallback: treat whole ligand as one token
+        return [(0, n_rows - 1)], [selfies_str or "[?]"]
+
+    g = len(toks)
+    edges = np.linspace(0, n_rows, g + 1, dtype=int)
+    groups = []
+    for i in range(g):
+        s, e = edges[i], edges[i + 1] - 1
+        if e < s:
+            e = s
+        groups.append((s, e))
+    return groups, toks
+
+
+
+def _connected_components_2d(mask: torch.Tensor) -> List[List[Tuple[int, int]]]:
+    """4-connected components over a 2D boolean mask (rows=ligand tokens, cols=protein residues)."""
+    h, w = mask.shape
+    visited = torch.zeros_like(mask, dtype=torch.bool)
+    comps: List[List[Tuple[int,int]]] = []
+    for i in range(h):
+        for j in range(w):
+            if mask[i, j] and not visited[i, j]:
+                stack = [(i, j)]
+                visited[i, j] = True
+                comp = []
+                while stack:
+                    r, c = stack.pop()
+                    comp.append((r, c))
+                    for dr, dc in ((1,0), (-1,0), (0,1), (0,-1)):
+                        rr, cc = r + dr, c + dc
+                        if 0 <= rr < h and 0 <= cc < w and mask[rr, cc] and not visited[rr, cc]:
+                            visited[rr, cc] = True
+                            stack.append((rr, cc))
+                comps.append(comp)
+    return comps
+
+def _format_component_table(
+    components,
+    p_tokens,
+    d_tokens,
+    *,
+    mode: str = "pair",   # "pair" | "residue"
+):
+    """
+    Render HTML table for highlighted interaction components.
+
+    Parameters
+    ----------
+    components : List[List[Tuple[int,int]]]
+        Each component is a list of (ligand_index, protein_index) pairs.
+    p_tokens : List[str]
+        Protein token strings.
+    d_tokens : List[str]
+        Ligand token strings.
+    mode : str
+        "pair"    -> show Protein range + Ligand range
+        "residue" -> show Protein residue(s) only
+    """
+
+    # ----------------------------
+    # Residue-only mode
+    # ----------------------------
+    if mode == "residue":
+        if not components:
+            return (
+                "<h4 style='margin:12px 0 6px'>Highlighted protein residues</h4>"
+                "<p>No residues selected.</p>"
+            )
+
+        rows = []
+        for comp in components:
+            # comp = [(lig_idx, prot_idx), ...]
+            prot_indices = [j for (_, j) in comp]
+            p_start, p_end = min(prot_indices), max(prot_indices)
+
+            p_s_idx, p_e_idx = p_start + 1, p_end + 1
+            p_s_tok = p_tokens[p_start] if p_start < len(p_tokens) else "?"
+            p_e_tok = p_tokens[p_end]   if p_end   < len(p_tokens) else "?"
+
+            if p_start == p_end:
+                label = f"{p_s_idx}:{p_s_tok}"
+            else:
+                label = f"{p_s_idx}:{p_s_tok} – {p_e_idx}:{p_e_tok}"
+
+            rows.append(
+                f"<tr>"
+                f"<td style='border:1px solid #ddd;padding:6px'>"
+                f"<strong>{label}</strong>"
+                f"</td>"
+                f"</tr>"
+            )
+
+        return (
+            "<h4 style='margin:12px 0 6px'>Highlighted protein residues</h4>"
+            "<table style='border-collapse:collapse;margin:6px 0 16px;width:60%'>"
+            "<thead><tr style='background:#f5f5f5'>"
+            "<th style='border:1px solid #ddd;padding:6px'>Protein residue(s)</th>"
+            "</tr></thead>"
+            f"<tbody>{''.join(rows)}</tbody></table>"
+        )
+
+    # ----------------------------
+    # Pair mode (default behaviour)
+    # ----------------------------
+    if not components:
+        return (
+            "<h4 style='margin:12px 0 6px'>Highlighted interaction segments</h4>"
+            "<p>No interaction pairs selected.</p>"
+        )
+
+    rows = []
+    for comp in components:
+        lig_indices = [i for (i, _) in comp]
+        prot_indices = [j for (_, j) in comp]
+
+        d_start, d_end = min(lig_indices), max(lig_indices)
+        p_start, p_end = min(prot_indices), max(prot_indices)
+
+        d_s_idx, d_e_idx = d_start + 1, d_end + 1
+        p_s_idx, p_e_idx = p_start + 1, p_end + 1
+
+        d_s_tok = d_tokens[d_start] if d_start < len(d_tokens) else "?"
+        d_e_tok = d_tokens[d_end]   if d_end   < len(d_tokens) else "?"
+        p_s_tok = p_tokens[p_start] if p_start < len(p_tokens) else "?"
+        p_e_tok = p_tokens[p_end]   if p_end   < len(p_tokens) else "?"
+
+        rows.append(
+            f"<tr>"
+            f"<td style='border:1px solid #ddd;padding:6px'>Protein: "
+            f"<strong>{p_s_idx}:{p_s_tok}</strong>"
+            f"{' – <strong>'+str(p_e_idx)+':'+p_e_tok+'</strong>' if p_end > p_start else ''}"
+            f"</td>"
+            f"<td style='border:1px solid #ddd;padding:6px'>Ligand: "
+            f"<strong>{d_s_idx}:{d_s_tok}</strong>"
+            f"{' – <strong>'+str(d_e_idx)+':'+d_e_tok+'</strong>' if d_end > d_start else ''}"
+            f"</td>"
+            f"</tr>"
+        )
+
+    return (
+        "<h4 style='margin:12px 0 6px'></h4>"
+        "<table style='border-collapse:collapse;margin:6px 0 16px;width:100%'>"
+        "<thead><tr style='background:#f5f5f5'>"
+        "<th style='border:1px solid #ddd;padding:6px'>Protein range</th>"
+        "<th style='border:1px solid #ddd;padding:6px'>Ligand range</th>"
+        "</tr></thead>"
+        f"<tbody>{''.join(rows)}</tbody></table>"
+    )
+
+
+def visualize_attention_and_ranges(
+    model,
+    feats,
+    head_idx: int,
+    *,
+    mode: str = "pair",                 # "pair" | "residue"
+    topk_pairs: int = 1,                 # Top-K interaction pairs (default=1)
+    topk_residues: int = 1,              # Top-K residues (1–20, default=1)
+    prot_tokenizer=None,
+    drug_tokenizer=None,
+    ligand_type: str = "selfies",
+    raw_selfies: Optional[str] = None,
+) -> Tuple[str, str]:
+    """
+    Visualise interaction attention with two complementary Top-K modes.
+
+    Modes
+    -----
+    mode="pair":
+        - Select Top-K highest-scoring (ligand token, protein residue) pairs
+        - Project selected pairs onto protein axis (evaluation-aligned)
+        - Default K = 1 (user-controlled)
+
+    mode="residue":
+        - Aggregate attention over ligand dimension
+        - Rank residues by aggregated score
+        - Select Top-K residues (1–20)
+        - Default K = 1 (binding pocket discovery)
+
+    Notes
+    -----
+    - Per-head GLOBAL SUM normalisation (matches test()).
+    - Specific heads mapped exactly to GT channels.
+    - Combined head = sum of 6 specific heads (NOT overall=7).
+    """
+
+    assert mode in {"pair", "residue"}
+    assert topk_pairs >= 1
+    assert 1 <= topk_residues <= 20
+
+    model.eval()
+    with torch.no_grad():
+        # --------------------------------------------------
+        # Unpack features
+        # --------------------------------------------------
+        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
+        # --------------------------------------------------
+        _, att_pd = model(p_emb, d_emb, p_mask, d_mask)
+        att = att_pd.squeeze(0)
+        # expected: [Ld, Lp, 8] or [8, Ld, Lp]
+
+        # --------------------------------------------------
+        # Channel mapping (must match test())
+        # --------------------------------------------------
+        VISIBLE2UNDERLYING = [1, 2, 3, 4, 6, 0]  # HB, Salt, Pi, Cat-Pi, Hydro, VdW
+        N_VISIBLE_SPEC = 6
+
+        def select_channel_map(att_):
+            if att_.dim() == 3 and att_.shape[-1] >= 8:
+                if head_idx < N_VISIBLE_SPEC:
+                    return att_[:, :, VISIBLE2UNDERLYING[head_idx]].cpu()
+                return att_[:, :, VISIBLE2UNDERLYING].sum(dim=2).cpu()
+            if att_.dim() == 3 and att_.shape[0] >= 8:
+                if head_idx < N_VISIBLE_SPEC:
+                    return att_[VISIBLE2UNDERLYING[head_idx]].cpu()
+                return att_[VISIBLE2UNDERLYING].sum(dim=0).cpu()
+            return att_.squeeze().cpu()
+
+        att2d_raw = select_channel_map(att)  # [Ld, Lp]
+
+        # --------------------------------------------------
+        # Per-head GLOBAL SUM normalisation (critical)
+        # --------------------------------------------------
+        att2d_raw = att2d_raw / (att2d_raw.sum() + 1e-8)
+
+        # --------------------------------------------------
+        # Token decoding & trimming
+        # --------------------------------------------------
+        def clean_tokens(ids, tokenizer):
+            toks = tokenizer.convert_ids_to_tokens(ids.tolist())
+            if hasattr(tokenizer, "all_special_tokens"):
+                toks = [t for t in toks if t not in tokenizer.all_special_tokens]
+            return toks
+
+        p_tokens_full = clean_tokens(p_ids[0], prot_tokenizer)
+        d_tokens_full = clean_tokens(d_ids[0], drug_tokenizer)
+
+        n_d = min(len(d_tokens_full), att2d_raw.size(0))
+        n_p = min(len(p_tokens_full), att2d_raw.size(1))
+
+        att2d = att2d_raw[:n_d, :n_p]
+        p_tokens = p_tokens_full[:n_p]
+        d_tokens = d_tokens_full[:n_d]
+
+        p_indices = list(range(1, n_p + 1))
+        d_indices = list(range(1, n_d + 1))
+
+        # --------------------------------------------------
+        # SELFIES row merging (for interpretability)
+        # --------------------------------------------------
+        if ligand_type == "selfies" and raw_selfies:
+            groups, labels = _group_rows_by_selfies_string(att2d.size(0), raw_selfies)
+            if groups:
+                merged = []
+                for s, e in groups:
+                    merged.append(att2d[s:e + 1].mean(dim=0, keepdim=True))
+                att2d = torch.cat(merged, dim=0)
+                d_tokens = labels
+                d_indices = list(range(1, len(labels) + 1))
+
+        # --------------------------------------------------
+        # Top-K selection (two modes)
+        # --------------------------------------------------
+        if mode == "pair":
+            # --- Top-K interaction pairs ---
+            flat = att2d.reshape(-1)
+            k_eff = min(topk_pairs, flat.numel())
+            idx = torch.topk(flat, k=k_eff).indices
+
+            mask_top = torch.zeros_like(flat, dtype=torch.bool)
+            mask_top[idx] = True
+            mask_top = mask_top.view_as(att2d)
+
+        else:
+            # --- Top-K residues ---
+            residue_score = att2d.sum(dim=0)          # [protein]
+            k_eff = min(topk_residues, residue_score.numel())
+            topk_res_idx = torch.topk(residue_score, k=k_eff).indices
+
+            mask_top = torch.zeros_like(att2d, dtype=torch.bool)
+            mask_top[:, topk_res_idx] = True           # keep all ligand rows
+
+        # --------------------------------------------------
+        # Connected components (visual coherence)
+        # --------------------------------------------------
+        components = _connected_components_2d(mask_top)
+        ranges_html = _format_component_table(
+            components,
+            p_tokens,
+            d_tokens,
+            mode=mode,
+        )
+
+
+        # --------------------------------------------------
+        # Crop to union of selected rows / columns
+        # --------------------------------------------------
+        rows_keep = mask_top.any(dim=1)
+        cols_keep = mask_top.any(dim=0)
+
+        if not rows_keep.any():
+            rows_keep[:] = True
+        if not cols_keep.any():
+            cols_keep[:] = True
+
+        vis = att2d[rows_keep][:, cols_keep]
+
+        d_tokens_vis = [t for k, t in zip(rows_keep.tolist(), d_tokens) if k]
+        p_tokens_vis = [t for k, t in zip(cols_keep.tolist(), p_tokens) if k]
+        d_indices_vis = [i for k, i in zip(rows_keep.tolist(), d_indices) if k]
+        p_indices_vis = [i for k, i in zip(cols_keep.tolist(), p_indices) if k]
+
+        # Cap columns for readability
+        if vis.size(1) > 150:
+            topc = torch.topk(vis.sum(0), k=150).indices
+            vis = vis[:, topc]
+            p_tokens_vis = [p_tokens_vis[i] for i in topc]
+            p_indices_vis = [p_indices_vis[i] for i in topc]
+
+        # --------------------------------------------------
+        # Plot
+        # --------------------------------------------------
+        x_labels = [f"{i}:{t}" for i, t in zip(p_indices_vis, p_tokens_vis)]
+        y_labels = [f"{i}:{t}" for i, t in zip(d_indices_vis, d_tokens_vis)]
+
+        fig_w = min(22, max(6, 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(vis.numpy(), aspect="auto", cmap=cm.viridis)
+
+        title = INTERACTION_NAMES[head_idx]
+        suffix = "Top-K pairs" if mode == "pair" else "Top-K residues"
+        ax.set_title(f"Ligand × Protein — {title} ({suffix})", fontsize=10, pad=8)
+        ax.set_xlabel("Protein residues")
+        ax.set_ylabel("Ligand tokens")
+
+        ax.set_xticks(range(len(x_labels)))
+        ax.set_xticklabels(x_labels, rotation=90, fontsize=8)
+        ax.set_yticks(range(len(y_labels)))
+        ax.set_yticklabels(y_labels, fontsize=7)
+
+        ax.xaxis.tick_top()
+        ax.xaxis.set_label_position("top")
+        ax.tick_params(axis="x", bottom=False)
+
+        fig.colorbar(im, fraction=0.026, pad=0.01)
+        fig.tight_layout()
+
+        # --------------------------------------------------
+        # Export
+        # --------------------------------------------------
+        buf_png = io.BytesIO()
+        buf_pdf = io.BytesIO()
+        fig.savefig(buf_png, format="png", dpi=140)
+        fig.savefig(buf_pdf, format="pdf")
+        plt.close(fig)
+
+        png_b64 = base64.b64encode(buf_png.getvalue()).decode()
+        pdf_b64 = base64.b64encode(buf_pdf.getvalue()).decode()
+
+        heat_html = f"""
+        <div style='position:relative'>
+          <a href='data:application/pdf;base64,{pdf_b64}' download='attention_{head_idx+1}.pdf'
+             style='position:absolute;top:10px;right:10px;
+                    background:#111;color:#fff;padding:8px 12px;
+                    border-radius:10px;font-size:.85rem;text-decoration:none'>
+             Download PDF
+          </a>
+          <img src='data:image/png;base64,{png_b64}' />
+        </div>
+        """
+
+        return heat_html, ranges_html
+
+
+
+
+# ───── Gradio callbacks ─────────────────────────────────────────
+ROOT = os.path.dirname(os.path.abspath(__file__))
+FOLDSEEK_BIN = os.path.join(ROOT, "bin", "foldseek")
+
+def extract_sequence_cb(structure_file, drug_text, protein_text):
+    """
+    Unified “Extract / Convert” callback.
+
+    Behaviour:
+      • If a .pdb / .cif file is uploaded and successfully parsed:
+          → overwrite protein textbox with extracted SA sequence
+      • If ligand textbox contains SMILES:
+          → convert to SELFIES
+      • Otherwise:
+          → keep existing protein / ligand unchanged
+    """
+
+    msgs = []
+
+    # --------------------------------------------------
+    # Default: ALWAYS keep current contents
+    # --------------------------------------------------
+    prot_seq_out = (protein_text or "").strip()
+    drug_seq_out = (drug_text or "").strip()
+
+    # --------------------------------------------------
+    # Structure file → SA sequence (overwrite protein ONLY on success)
+    # --------------------------------------------------
+    if structure_file is not None and os.path.exists(structure_file.name):
+        try:
+            parsed = get_struc_seq(
+                FOLDSEEK_BIN,
+                structure_file.name,
+                None,
+                plddt_mask=False,
+            )
+            first_chain = next(iter(parsed))
+            _, _, struct_seq = parsed[first_chain]
+
+            if struct_seq:
+                prot_seq_out = struct_seq
+                msgs.append(
+                    "<li>✅ Extracted <b>SA sequence</b> from the uploaded structure.</li>"
+                )
+            else:
+                msgs.append(
+                    "<li>❌ Structure parsed, but no valid protein sequence found.</li>"
+                )
+
+        except Exception as e:
+            msgs.append(
+                f"<li>❌ Failed to extract SA sequence: <b>{e}</b></li>"
+            )
+
+    # --------------------------------------------------
+    # SMILES → SELFIES (do NOT touch protein)
+    # --------------------------------------------------
+    if drug_seq_out:
+        lig_type = detect_ligand_type(drug_seq_out)
+        if lig_type == "smiles":
+            try:
+                conv = smiles_to_selfies(drug_seq_out)
+                if conv is None:
+                    msgs.append(
+                        "<li>❌ SMILES → SELFIES failed: <b>invalid SMILES</b>.</li>"
+                    )
+                else:
+                    drug_seq_out = conv
+                    msgs.append(
+                        "<li>✅ Converted <b>SMILES</b> to <b>SELFIES</b>.</li>"
+                    )
+            except Exception as e:
+                msgs.append(
+                    f"<li>❌ SMILES → SELFIES error: <b>{e}</b></li>"
+                )
+
+    # --------------------------------------------------
+    # Status message box
+    # --------------------------------------------------
+    if msgs:
+        status_html = (
+            "<div style='margin:10px 0;padding:10px 12px;"
+            "border:1px solid #e5e7eb;border-radius:10px;"
+            "background:#f8fafc;color:#0f172a'>"
+            "<ul style='margin:0 0 0 18px;padding:0'>"
+            f"{''.join(msgs)}"
+            "</ul></div>"
+        )
+    else:
+        status_html = ""
+
+    return prot_seq_out, drug_seq_out, status_html
+
+
+def _choose_ckpt_by_types(prot_seq: str, ligand_text: str) -> Tuple[str, str, str]:
+    """Return (folder_name, protein_type, ligand_type) for checkpoint routing."""
+    ptype = detect_protein_type(prot_seq)
+    ltype = detect_ligand_type(ligand_text)
+    folder = f"{ptype}_{ltype}"  # sa_selfies / fasta_selfies / sa_smiles / fasta_smiles
+    return folder, ptype, ltype
+
+
+def inference_cb(prot_seq, drug_seq, head_choice, topk_choice, mode_choice):
+    """
+    Inference callback supporting two Top-K modes:
+      - Top-K interaction pairs
+      - Top-K residues
+    """
+
+    # ------------------------------
+    # Input validation
+    # ------------------------------
+    if not prot_seq or not prot_seq.strip():
+        return "<p style='color:red'>Please extract or enter a protein sequence first.</p>", ""
+
+    if not drug_seq or not drug_seq.strip():
+        return "<p style='color:red'>Please enter a ligand sequence (SELFIES or SMILES).</p>", ""
+
+    prot_seq = prot_seq.strip()
+    drug_seq_in = drug_seq.strip()
+
+    # ------------------------------
+    # Detect types & checkpoint routing
+    # ------------------------------
+    folder, ptype, ltype = _choose_ckpt_by_types(prot_seq, drug_seq_in)
+
+    # Ligand normalisation: always tokenise as SELFIES
+    if ltype == "smiles":
+        conv = smiles_to_selfies(drug_seq_in)
+        if conv is None:
+            return (
+                "<p style='color:red'>SMILES→SELFIES conversion failed. "
+                "The SMILES appears invalid.</p>",
+                "",
+            )
+        drug_seq_for_tokenizer = conv
+    else:
+        drug_seq_for_tokenizer = drug_seq_in
+    
+    # 🔒 强制统一类型
+    ltype = "selfies"
+    ligand_type_flag = "selfies"
+    raw_selfies = drug_seq_for_tokenizer
+    folder = f"{ptype}_selfies"
+
+
+    # # Ligand normalisation: always tokenise as SELFIES
+    # if ltype == "smiles":
+    #     conv = smiles_to_selfies(drug_seq_in)
+    #     if conv is None:
+    #         return (
+    #             "<p style='color:red'>SMILES→SELFIES conversion failed. "
+    #             "The SMILES appears invalid.</p>",
+    #             "",
+    #         )
+    #     drug_seq_for_tokenizer = conv
+    #     ligand_type_flag = "selfies"
+    # else:
+    #     drug_seq_for_tokenizer = drug_seq_in
+    #     ligand_type_flag = "selfies"
+
+    # raw_selfies = drug_seq_for_tokenizer if ligand_type_flag == "selfies" else None
+
+    # ------------------------------
+    # Load encoders
+    # ------------------------------
+    prot_tok, prot_m, drug_tok, drug_m, encoding = load_encoders(ptype, ltype, args)
+
+    loader = DataLoader(
+        [(prot_seq, drug_seq_for_tokenizer, 1)],
+        batch_size=1,
+        collate_fn=make_collate_fn(prot_tok, drug_tok),
+    )
+
+    feats = get_case_feature(encoding, loader)
+
+    # ------------------------------
+    # Load trained checkpoint (if exists)
+    # ------------------------------
+    ckpt = os.path.join(args.save_path_prefix, folder, "best_model.ckpt")
+    model = ExplainBind(1280, 768, args=args).to(DEVICE)
+
+    if os.path.isfile(ckpt):
+        model.load_state_dict(torch.load(ckpt, map_location=DEVICE))
+        warn_html = (
+            "<div style='margin:8px 0 14px;padding:8px 10px;"
+            "border-left:4px solid #10b981;background:#ecfdf5'>"
+            f"<b>Loaded model:</b> <code>{folder}/best_model.ckpt</code></div>"
+        )
+    else:
+        warn_html = (
+            "<div style='margin:8px 0 14px;padding:8px 10px;"
+            "border-left:4px solid #f59e0b;background:#fffbeb'>"
+            "<b>Warning:</b> checkpoint not found "
+            f"<code>{folder}/best_model.ckpt</code>. "
+            "Using randomly initialised weights for visualisation.</div>"
+        )
+
+    # ------------------------------
+    # Parse interaction head
+    # ------------------------------
+    sel = str(head_choice).strip()
+    if sel in INTERACTION_NAMES:
+        head_idx = INTERACTION_NAMES.index(sel)
+    else:
+        try:
+            n = int(sel.split(".", 1)[0])
+            head_idx = max(0, min(len(INTERACTION_NAMES) - 1, n - 1))
+        except Exception:
+            head_idx = len(INTERACTION_NAMES) - 1  # Combined Interaction
+
+    # ------------------------------
+    # Parse Top-K value
+    # ------------------------------
+    try:
+        topk = int(str(topk_choice).strip())
+    except Exception:
+        topk = 1
+
+    topk = max(1, topk)
+
+    # ------------------------------
+    # Parse mode (pair vs residue)
+    # ------------------------------
+    mode_choice = str(mode_choice).lower()
+
+    if "residue" in mode_choice:
+        mode = "residue"
+        topk_pairs = 1
+        topk_residues = min(20, topk)
+    else:
+        mode = "pair"
+        topk_pairs = topk
+        topk_residues = 1
+
+    # ------------------------------
+    # Visualisation
+    # ------------------------------
+    heat_html, table_html = visualize_attention_and_ranges(
+        model,
+        feats,
+        head_idx,
+        mode=mode,
+        topk_pairs=topk_pairs,
+        topk_residues=topk_residues,
+        prot_tokenizer=prot_tok,
+        drug_tokenizer=drug_tok,
+        ligand_type=ligand_type_flag,
+        raw_selfies=raw_selfies,
+    )
+
+    # ------------------------------
+    # Info box
+    # ------------------------------
+    # info_html = (
+    #     "<div style='margin:10px 0;padding:10px 12px;"
+    #     "border:1px solid #e5e7eb;border-radius:10px;"
+    #     "background:#f8fafc;color:#0f172a'>"
+    #     "<div style='margin-bottom:6px'>"
+    #     f"<b>Detected types:</b> "
+    #     f"Protein=<code>{ptype.upper()}</code>, "
+    #     f"Ligand=<code>{ltype.upper()}</code></div>"
+    #     f"<div><b>Top-K mode:</b> "
+    #     f"<code>{'Interaction pairs' if mode == 'pair' else 'Residues'}</code>, "
+    #     f"K=<code>{topk}</code></div>"
+    #     f"{warn_html}</div>"
+    # )
+
+    # table_html = info_html + table_html
+
+    return table_html, heat_html
+
+
+def clear_cb():
+    return "", "", "", "", None, ""
+# protein, drug, table, heat, file, status
+
+
+# ───── Gradio interface definition ───────────────────────────────
+css = """
+:root{
+  --bg:#f8fafc; --card:#f8fafc; --text:#0f172a;
+  --muted:#6b7280; --border:#e5e7eb; --shadow:0 6px 24px rgba(2,6,23,.06);
+  --radius:14px; --icon-size:20px;
+}
+
+*{box-sizing:border-box}
+html,body{background:#fff!important;color:var(--text)!important}
+.gradio-container{max-width:1120px;margin:0 auto}
+
+/* Title and subtitle */
+h1{
+  font-family:Inter,ui-sans-serif;letter-spacing:.2px;font-weight:700;
+  font-size:32px;margin:22px 0 12px;text-align:center
+}
+.subtle{color:var(--muted);font-size:14px;text-align:center;margin:-6px 0 18px}
+
+/* Card style */
+.card{
+  background:var(--card); border:1px solid var(--border); border-radius:var(--radius);
+  box-shadow:var(--shadow); padding:22px;
+}
+
+/* Top links */
+.link-row{display:flex;justify-content:center;gap:14px;margin:0 auto 18px;flex-wrap:wrap}
+
+/* Two-column grid: left=input, right=controls */
+.grid-2{display:grid;grid-template-columns:1.4fr .9fr;gap:16px}
+.grid-2 .col{display:flex;flex-direction:column;gap:12px}
+
+/* Buttons */
+.gr-button{border-radius:12px !important;font-weight:700 !important;letter-spacing:.2px}
+#extract-btn{background:linear-gradient(90deg,#EFAFB2,#EFAFB2); color:#0f172a}
+#inference-btn{background:linear-gradient(90deg,#B2CBDF,#B2CBDF); color:#0f172a}
+#clear-btn{background:#FFE2B5; color:#0A0A0A; border:1px solid var(--border)}
+
+/* Result spacing */
+#result-table{margin-bottom:16px}
+
+/* Figure container */
+.figure-wrap{border:1px solid var(--border);border-radius:12px;overflow:hidden;box-shadow:var(--shadow)}
+.figure-wrap img{display:block;width:100%;height:auto}
+
+/* Right pane: vertical radio layout and full-width controls (kept for button styling) */
+.right-pane .gr-button{
+  width:100% !important;
+  height:48px !important;
+  border-radius:12px !important;
+  font-weight:700 !important;
+  letter-spacing:.2px;
+}
+/* ───────── Publication links (Bulma-like) ───────── */
+
+.publication-links {
+  display: flex;
+  justify-content: center;
+  gap: 14px;
+  flex-wrap: wrap;
+  margin: 6px 0 18px;
+}
+
+.link-block a {
+  display: inline-flex;
+  align-items: center;
+  gap: 8px;
+  padding: 10px 18px;
+  font-size: 14px;
+  font-weight: 600;
+  border-radius: 9999px;
+  text-decoration: none;
+  transition: all 0.15s ease-in-out;
+}
+
+/* colour variants */
+.btn-danger  { background:#e2e8f0; color:#0f172a; }
+.btn-dark    { background:#e2e8f0; color:#0f172a; }
+.btn-link    { background:#e2e8f0; color:#0f172a; }
+.btn-warning { background:#e2e8f0; color:#0f172a; }
+
+.link-block a:hover {
+  filter: brightness(0.95);
+  transform: translateY(-1px);
+}
+
+.loscalzo-block img {
+  height: 100px;          
+  width: auto;     
+  object-fit: contain;
+}
+
+.loscalzo-block {
+  display: flex;
+  align-items: center;
+  gap: 10px;
+
+  margin: 0 auto;
+  justify-content: center;
+}
+
+
+}
+"""
+with gr.Blocks(
+    theme=gr.themes.Default(),
+    css=css,
+    head="""
+    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.5.1/css/all.min.css">
+    """
+) as demo:
+    gr.Markdown("<h1>ExplainBind: Token-level Protein–Ligand Interaction Visualiser</h1>")
+    gr.Markdown(
+        '<p class="subtle">'
+        'Upload structure → Extract SA sequence · Paste SMILES/SELFIES · '
+        'Choose interaction channel and Top-K mode'
+        '</p>'
+    )
+
+
+    # ===== Loscalzo logo + group name =====
+    gr.HTML(f"""
+    <div class="loscalzo-block">
+      <img src="data:image/png;base64,{LOSCAZLO_B64}"
+           alt="Loscalzo Research Group logo" />
+      <a class="loscalzo-name"
+         href="https://ogephd.hms.harvard.edu/people/joseph-loscalzo"
+         target="_blank" rel="noopener">
+      </a>
+    </div>
+    """)
+    # ───────────────────────────────
+    # Top links
+    # ───────────────────────────────
+    gr.Markdown("""
+    <div class="publication-links">
+    
+      <span class="link-block">
+        <a href="https://arxiv.org/abs/2406.01651" target="_blank" class="btn-danger">
+          <i class="fa-solid fa-file-lines"></i> Paper
+        </a>
+      </span>
+    
+      <span class="link-block">
+        <a href="https://github.com/ZhaohanM/ExplainBind" target="_blank" class="btn-dark">
+          <i class="fa-brands fa-github"></i> Source Code
+        </a>
+      </span>
+    
+      <span class="link-block">
+        <a href="ZhaohanM.github.io/ExplainBind/" target="_blank" class="btn-link">
+          <i class="fa-solid fa-globe"></i> Project Page
+        </a>
+      </span>
+    </div>
+    """)
+
+      #     <span class="link-block">
+      #   <a href="https://huggingface.co/spaces/Zhaohan-Meng/InteractBind" target="_blank" class="btn-warning">
+      #     🗄 Datasets
+      #   </a>
+      # </span>
+    # gr.Markdown("""
+    #     <div class="publication-links" style="text-align:center; margin-top: 10px;">
+        
+    #     <!-- ───── Project Badges ───── -->
+    #     <a href="https://zhaohanm.github.io/ExplainBind/" target="_blank">
+    #       <img src="https://img.shields.io/badge/Project-Page-4285F4?style=for-the-badge&logo=googlelens&logoColor=4285F4">
+    #     </a>
+        
+    #     <a href="https://doi.org/10.1101/2022.09.16.508229" target="_blank">
+    #       <img src="https://img.shields.io/badge/bioRxiv-10.1101%2F2022.09.16.508229-orange?style=for-the-badge">
+    #     </a>
+        
+    #     <a href="https://github.com/ZhaohanM/ExplainBind/blob/main/LICENSE" target="_blank">
+    #       <img src="https://img.shields.io/badge/License-MIT-green?style=for-the-badge">
+    #     </a>
+        
+    #     <a href="https://visitorbadge.io/status?path=https%3A%2F%2Fgithub.com%2FZhaohanM%2FExplainBind" target="_blank">
+    #       <img src="https://api.visitorbadge.io/api/combined?path=https%3A%2F%2Fgithub.com%2FZhaohanM%2FExplainBind&label=Views&countColor=%23f36f43&style=for-the-badge">
+    #     </a>
+        
+    #     </div>
+    #     """)
+
+
+    # ───────────────────────────────
+    # Guidelines
+    # ───────────────────────────────
+    with gr.Accordion("Guidelines for Users", open=True, elem_classes=["card"]):
+        gr.HTML("""
+        <ol style="font-size:1rem;line-height:1.6;margin-left:22px;">
+          <li>
+            <strong>Input types:</strong>
+            The model supports <em> structure-aware (SA)</em> or <em>FASTA</em> protein sequences,
+            and <em>SMILES</em> or <em>SELFIES</em> ligand inputs.
+          </li>
+          <li>
+            <strong>Extract sequence:</strong>
+            (1) Converts <em>SMILES</em> to <em>SELFIES</em>;
+            (2) Extracts an <em>SA</em> sequence from uploaded
+            <code>.pdb</code> or <code>.cif</code> files.
+          </li>
+          <li>
+            <strong>Top-K mode:</strong>
+            <ul style="margin-top:6px;">
+              <li>
+                <em>Top-K residues-atom pairs</em>:
+                ranks individual protein-residue and ligand-atom pairs by attention score.
+              </li>
+              <li>
+                <em>Top-K residues</em>:
+                ranks protein residues by attention aggregated over all ligand tokens.
+              </li>
+            </ul>
+          </li>
+          <li>
+            <strong>Inference output:</strong>
+            Generates a token-level attention heat map
+            and a corresponding results table
+            based on the selected Top-K mode.
+          </li>
+        </ol>
+        """)
+
+
+    # ───────────────────────────────
+    # Inputs (left) + Controls (right)
+    # ───────────────────────────────
+    with gr.Row():
+        with gr.Column(elem_classes=["card", "grid-2"]):
+            # ── Left: Inputs ──
+            with gr.Column(elem_id="left"):
+                protein_seq = gr.Textbox(
+                    label="Protein structure-aware / FASTA sequence",
+                    lines=3,
+                    placeholder="Paste SA/FASTA sequence or click Extract…",
+                    elem_id="protein-seq",
+                )
+
+                drug_seq = gr.Textbox(
+                    label="Ligand (SELFIES / SMILES)",
+                    lines=3,
+                    placeholder="Paste SELFIES or SMILES",
+                    elem_id="drug-seq",
+                )
+
+                structure_file = gr.File(
+                    label="Upload protein structure (.pdb / .cif)",
+                    file_types=[".pdb", ".cif"],
+                    elem_id="structure-file",
+                )
+
+                gr.Examples(
+                    examples=[[
+                        "MTLSILVAHDLQRVIGFENQLPWHLPNDLKHVKKLSTGHTLVMGRKTFESIGKPLPNRRNVVLTSDTSFNVEGVDVIHSIEDIYQLPGHVFIFGGQTLFEEMIDKVDDMYITVIEGKFRGDTFFPPYTFEDWEVASSVEGKLDEKNTIPHTFLHLIRKK",
+                        "[C][O][C][=C][C][Branch1][=C][C][C][=C][N][=C][Branch1][C][N][N][=C][Ring1][#Branch1][N][=C][C][Branch1][Ring1][O][C][=C][Ring1][P][O][C]"
+                    ]],
+                    inputs=[protein_seq, drug_seq],
+                    label="Click to load an example",
+                )
+
+            # ── Right: Controls ──
+            with gr.Column(elem_id="right", elem_classes=["right-pane"]):
+                head_dd = gr.Dropdown(
+                    label="Interaction Type/Overall",
+                    choices=INTERACTION_NAMES,
+                    value="Overall Interaction",
+                    interactive=True,
+                )
+
+                mode_dd = gr.Dropdown(
+                    label="Top-K selection mode",
+                    choices=[
+                        "Top-K residues-atom pairs",
+                        "Top-K residues",
+                    ],
+                    value="Top-K residues-atom pairs",
+                    interactive=True,
+                )
+
+                top_k_dd = gr.Dropdown(
+                    label="Top-K value",
+                    choices=[str(i) for i in range(1, 21)],
+                    value="1",
+                    interactive=True,
+                )
+
+                btn_extract = gr.Button("Extract / Convert sequences", elem_id="extract-btn")
+                btn_infer   = gr.Button("Inference",        elem_id="inference-btn")
+                clear_btn   = gr.Button("Clear",            elem_id="clear-btn")
+
+    # ───────────────────────────────
+    # Outputs
+    # ───────────────────────────────
+    with gr.Column(elem_classes=["card"]):
+        status_box   = gr.HTML(elem_id="status-box")
+        output_table = gr.HTML(elem_id="result-table")
+        output_heat  = gr.HTML(elem_id="result-heat")
+
+    # ───────────────────────────────
+    # Wiring
+    # ───────────────────────────────
+    btn_extract.click(
+        fn=extract_sequence_cb,
+        inputs=[
+            structure_file,
+            drug_seq,
+            protein_seq,
+        ],
+        outputs=[
+            protein_seq,
+            drug_seq,
+            status_box,
+        ],
+    )
+
+    btn_infer.click(
+        fn=inference_cb,
+        inputs=[protein_seq, drug_seq, head_dd, top_k_dd, mode_dd],
+        outputs=[output_table, output_heat],
+    )
+
+    clear_btn.click(
+        fn=clear_cb,
+        inputs=[],
+        outputs=[
+            protein_seq,
+            drug_seq,
+            output_table,
+            output_heat,
+            structure_file,
+            status_box,
+        ],
+    )
+
+if __name__ == "__main__":
+    demo.launch(
+        server_name="127.0.0.1",
+        server_port=7860,
+        share=True,
+        inbrowser=False,
+        show_error=True,
+    )

requirements.txt

@@ -0,0 +1,13 @@
+flask
+torch
+biopython
+transformers
+selfies
+rdkit-pypi
+gradio==6.5.1
+matplotlib
+scipy<1.9.0
+numpy<1.23.0
+scikit-learn
+pandas
+ipython

save_model_ckp/fasta_selfies/best_model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19fbf033e333cc71a47931880eef70fc404101cc0b887ef1e123b1f9f8fe4624
+size 35855458

save_model_ckp/sa_selfies/best_model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3402a8f2f1ac484ca42ca4f54e96ad67e5e1c764a02e70a76dda42b6f8a1c2d5
+size 35855458

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

@@ -0,0 +1,319 @@
+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 ExplainBind(nn.Module):
+    def __init__(self, prot_out_dim, drug_out_dim, args):
+        super(ExplainBind, self).__init__()
+        self.fusion = args.fusion
+        self.drug_reg = nn.Linear(drug_out_dim, 768)
+        self.prot_reg = nn.Linear(prot_out_dim, 768)
+
+        if self.fusion == "CAN":
+            self.can_layer = CAN_Layer(hidden_dim=768, num_heads=8, args=args)
+            self.mlp_classifier = MlPdecoder_CAN(input_dim=1536)
+        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)
+
+            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, 
+            output_hidden_states=True,  # Request hidden states
+            return_dict=True
+        ).hidden_states[-1]
+        # 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()
+            
+
+        query_embed = torch.cat([prot_embed, drug_embed], dim=1)
+
+        return query_embed, alpha_dp
+
+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']