Update app.py

app.py

@@ -1,10 +1,12 @@
 """
 BALM-PPI Pro  ·  ESM-2 + LoRA + Integrated Gradients
-Light-first theme, auto dark-mode via prefers-color-scheme,
-theme-reactive NGL viewer, cached PDB fetches.
+Fixes:
+  - IG forced to float32 (torch_dtype="auto" → fp16 kills grads on HF Spaces)
+  - IG steps reduced 30→15 for speed with negligible accuracy loss
+  - NGL viewer filters PDB to only the selected chains
+  - PDB fetches cached for 1 h
 """
 
-# ── Speed tweaks — must be BEFORE any HF/torch imports ─────────────────────
 import os
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 os.environ["TOKENIZERS_PARALLELISM"]     = "false"
@@ -19,7 +21,7 @@ import requests
 from Bio import PDB
 from Bio.Data.PDBData import protein_letters_3to1 as THREE_TO_ONE
 
-st.set_page_config(page_title="BALM-PPI ", page_icon="🧬", layout="wide",
+st.set_page_config(page_title="BALM-PPI Pro", page_icon="🧬", layout="wide",
                    initial_sidebar_state="expanded")
 
 HF_REPO_ID  = "Harshit494/BALM-PPI"
@@ -39,7 +41,6 @@ EXAMPLES = [
     },
 ]
 
-# ── Session state ─────────────────────────────────────────────────────────
 _DEFAULTS: dict = {
     "model": None, "device": None,
     "result": None, "ig_a": None, "ig_b": None, "ig_chain_map": None,
@@ -54,13 +55,11 @@ for _k, _v in _DEFAULTS.items():
         st.session_state[_k] = _v
 
 # ══════════════════════════════════════════════════════════════════════════════
-#  CSS — light-first with @media dark override
+#  CSS — light-first, dark via @media
 # ══════════════════════════════════════════════════════════════════════════════
 st.markdown("""
 <style>
 @import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;500;600;700&family=Inter:wght@300;400;500;600;700&display=swap');
-
-/* LIGHT tokens */
 :root {
   --bg0:#ffffff; --bg1:#f8f9fc; --bg2:#f1f4f9; --bg3:#e4e9f2;
   --border:rgba(37,99,235,0.13); --shadow:rgba(37,99,235,0.07);
@@ -69,7 +68,6 @@ st.markdown("""
   --text0:#0f172a; --text1:#334155; --text2:#64748b; --text3:#94a3b8;
   --mono:'JetBrains Mono',monospace; --sans:'Inter',sans-serif;
 }
-/* DARK tokens */
 @media (prefers-color-scheme:dark){
   :root {
     --bg0:#06090f; --bg1:#0b1120; --bg2:#101828; --bg3:#172135;
@@ -79,213 +77,111 @@ st.markdown("""
     --text0:#f1f5f9; --text1:#cbd5e1; --text2:#64748b; --text3:#334155;
   }
 }
-
-/* Global */
 html,body,[data-testid="stAppViewContainer"],[data-testid="stMain"]{
-  background:var(--bg0)!important; font-family:var(--sans)!important;
+  background:var(--bg0)!important;font-family:var(--sans)!important;
 }
 [data-testid="stHeader"]{
-  background:var(--bg0)!important;
-  border-bottom:1px solid var(--border)!important;
+  background:var(--bg0)!important;border-bottom:1px solid var(--border)!important;
   box-shadow:0 1px 8px var(--shadow)!important;
 }
-[data-testid="stSidebar"]{
-  background:var(--bg1)!important;
-  border-right:1px solid var(--border)!important;
-}
+[data-testid="stSidebar"]{background:var(--bg1)!important;border-right:1px solid var(--border)!important;}
 [data-testid="stSidebar"] *{font-family:var(--sans)!important;}
-
-/* Typography */
 h1,h2,h3,h4,h5,h6{font-family:var(--sans)!important;color:var(--text0)!important;font-weight:700!important;}
 p,label,div,span,li{color:var(--text1);}
 code,pre{font-family:var(--mono)!important;font-size:.82rem!important;
   background:var(--bg2)!important;border-radius:4px!important;color:var(--accent)!important;}
-
-/* Buttons */
-.stButton>button{
-  font-family:var(--sans)!important;font-weight:600!important;
-  letter-spacing:.01em!important;border-radius:8px!important;
-  transition:all .15s ease!important;height:38px!important;
-}
+.stButton>button{font-family:var(--sans)!important;font-weight:600!important;
+  letter-spacing:.01em!important;border-radius:8px!important;transition:all .15s ease!important;height:38px!important;}
 .stButton>button[kind="primary"]{
   background:linear-gradient(135deg,var(--accent),var(--accent2))!important;
-  border:none!important;color:#fff!important;
-  box-shadow:0 2px 12px var(--shadow)!important;
-}
-.stButton>button[kind="primary"]:hover{
-  box-shadow:0 4px 20px rgba(37,99,235,.28)!important;
-  transform:translateY(-1px)!important;
-}
-.stButton>button[kind="secondary"]{
-  background:var(--bg1)!important;
-  border:1px solid var(--border)!important;
-  color:var(--text1)!important;
-}
-.stButton>button[kind="secondary"]:hover{
-  border-color:var(--accent)!important;color:var(--accent)!important;
-  background:rgba(37,99,235,.05)!important;
-}
-
-/* Text areas */
-.stTextArea textarea{
-  background:var(--bg0)!important;border:1.5px solid var(--border)!important;
-  border-radius:8px!important;color:var(--text0)!important;
-  font-family:var(--mono)!important;font-size:.82rem!important;
-  line-height:1.65!important;resize:vertical!important;
-}
-.stTextArea textarea:focus{
-  border-color:var(--accent)!important;
-  box-shadow:0 0 0 3px rgba(37,99,235,.1)!important;
-}
-.stTextArea label{font-weight:600!important;font-size:.76rem!important;
-  text-transform:uppercase!important;letter-spacing:.09em!important;color:var(--text2)!important;}
-
-/* Inputs */
-.stTextInput input,.stNumberInput input{
-  background:var(--bg0)!important;border:1.5px solid var(--border)!important;
-  border-radius:8px!important;color:var(--text0)!important;
-}
-.stTextInput input:focus,.stNumberInput input:focus{
-  border-color:var(--accent)!important;box-shadow:0 0 0 3px rgba(37,99,235,.1)!important;
-}
-.stTextInput label,.stNumberInput label{
-  font-size:.76rem!important;font-weight:600!important;
-  color:var(--text2)!important;text-transform:uppercase!important;letter-spacing:.08em!important;
-}
-
-/* Radio/checkbox */
+  border:none!important;color:#fff!important;box-shadow:0 2px 12px var(--shadow)!important;}
+.stButton>button[kind="primary"]:hover{box-shadow:0 4px 20px rgba(37,99,235,.28)!important;transform:translateY(-1px)!important;}
+.stButton>button[kind="secondary"]{background:var(--bg1)!important;border:1px solid var(--border)!important;color:var(--text1)!important;}
+.stButton>button[kind="secondary"]:hover{border-color:var(--accent)!important;color:var(--accent)!important;background:rgba(37,99,235,.05)!important;}
+.stTextArea textarea{background:var(--bg0)!important;border:1.5px solid var(--border)!important;
+  border-radius:8px!important;color:var(--text0)!important;font-family:var(--mono)!important;
+  font-size:.82rem!important;line-height:1.65!important;resize:vertical!important;}
+.stTextArea textarea:focus{border-color:var(--accent)!important;box-shadow:0 0 0 3px rgba(37,99,235,.1)!important;}
+.stTextArea label{font-weight:600!important;font-size:.76rem!important;text-transform:uppercase!important;
+  letter-spacing:.09em!important;color:var(--text2)!important;}
+.stTextInput input,.stNumberInput input{background:var(--bg0)!important;border:1.5px solid var(--border)!important;
+  border-radius:8px!important;color:var(--text0)!important;}
+.stTextInput input:focus,.stNumberInput input:focus{border-color:var(--accent)!important;box-shadow:0 0 0 3px rgba(37,99,235,.1)!important;}
+.stTextInput label,.stNumberInput label{font-size:.76rem!important;font-weight:600!important;
+  color:var(--text2)!important;text-transform:uppercase!important;letter-spacing:.08em!important;}
 .stRadio label,.stCheckbox label{color:var(--text1)!important;font-size:.9rem!important;}
-
-/* Tabs */
-.stTabs [data-baseweb="tab-list"]{
-  background:var(--bg2)!important;border-radius:10px!important;
-  padding:4px!important;gap:3px!important;border:1px solid var(--border)!important;
-}
-.stTabs [data-baseweb="tab"]{
-  background:transparent!important;border-radius:7px!important;
-  color:var(--text2)!important;font-family:var(--sans)!important;
-  font-size:.82rem!important;font-weight:500!important;
-  transition:all .14s ease!important;padding:6px 13px!important;
-}
-.stTabs [aria-selected="true"]{
-  background:var(--bg0)!important;color:var(--accent)!important;
-  font-weight:600!important;box-shadow:0 1px 6px var(--shadow)!important;
-}
-
-/* Metrics */
-[data-testid="stMetric"]{
-  background:var(--bg1)!important;border:1px solid var(--border)!important;
-  border-radius:10px!important;padding:14px 18px!important;
-  box-shadow:0 1px 4px var(--shadow)!important;
-}
-[data-testid="stMetricLabel"]{
-  color:var(--text2)!important;font-size:.72rem!important;
-  text-transform:uppercase!important;letter-spacing:.1em!important;
-}
-[data-testid="stMetricValue"]{
-  color:var(--accent)!important;font-family:var(--mono)!important;
-  font-size:1.55rem!important;font-weight:700!important;
-}
-
-/* Misc */
+.stTabs [data-baseweb="tab-list"]{background:var(--bg2)!important;border-radius:10px!important;
+  padding:4px!important;gap:3px!important;border:1px solid var(--border)!important;}
+.stTabs [data-baseweb="tab"]{background:transparent!important;border-radius:7px!important;
+  color:var(--text2)!important;font-family:var(--sans)!important;font-size:.82rem!important;
+  font-weight:500!important;transition:all .14s ease!important;padding:6px 13px!important;}
+.stTabs [aria-selected="true"]{background:var(--bg0)!important;color:var(--accent)!important;
+  font-weight:600!important;box-shadow:0 1px 6px var(--shadow)!important;}
+[data-testid="stMetric"]{background:var(--bg1)!important;border:1px solid var(--border)!important;
+  border-radius:10px!important;padding:14px 18px!important;box-shadow:0 1px 4px var(--shadow)!important;}
+[data-testid="stMetricLabel"]{color:var(--text2)!important;font-size:.72rem!important;
+  text-transform:uppercase!important;letter-spacing:.1em!important;}
+[data-testid="stMetricValue"]{color:var(--accent)!important;font-family:var(--mono)!important;
+  font-size:1.55rem!important;font-weight:700!important;}
 hr{border-color:var(--border)!important;margin:12px 0!important;}
 [data-testid="stAlert"]{background:var(--bg1)!important;border-radius:8px!important;border-left-width:3px!important;}
 [data-testid="stDataFrame"]{border:1px solid var(--border)!important;border-radius:8px!important;overflow:hidden!important;}
-.stProgress>div>div{
-  background:linear-gradient(90deg,var(--accent),var(--accent2))!important;
-  border-radius:4px!important;
-}
+.stProgress>div>div{background:linear-gradient(90deg,var(--accent),var(--accent2))!important;border-radius:4px!important;}
 
-/* ── Custom classes ─────────────────────────────────────────── */
 .app-header{display:flex;align-items:center;gap:14px;padding:4px 0 14px;}
-.app-logo{
-  width:40px;height:40px;border-radius:10px;
+.app-logo{width:40px;height:40px;border-radius:10px;
   background:linear-gradient(135deg,var(--accent),var(--accent2));
   display:flex;align-items:center;justify-content:center;
-  font-size:20px;flex-shrink:0;box-shadow:0 2px 12px rgba(37,99,235,.22);
-}
+  font-size:20px;flex-shrink:0;box-shadow:0 2px 12px rgba(37,99,235,.22);}
 .app-title{font-family:var(--mono)!important;font-size:1.22rem!important;
   font-weight:700!important;color:var(--text0)!important;letter-spacing:-.02em;}
-.app-subtitle{font-size:.77rem!important;color:var(--text2)!important;
-  margin-top:1px;letter-spacing:.03em;}
-
-.sec-hdr{
-  font-family:var(--mono);font-size:.65rem;font-weight:700;
-  letter-spacing:.2em;text-transform:uppercase;
-  color:var(--accent);margin:14px 0 8px;
-  display:flex;align-items:center;gap:8px;
-}
+.app-subtitle{font-size:.77rem!important;color:var(--text2)!important;margin-top:1px;letter-spacing:.03em;}
+.sec-hdr{font-family:var(--mono);font-size:.65rem;font-weight:700;letter-spacing:.2em;
+  text-transform:uppercase;color:var(--accent);margin:14px 0 8px;
+  display:flex;align-items:center;gap:8px;}
 .sec-hdr::after{content:'';flex:1;height:1px;background:var(--border);}
-
-.ex-card{
-  background:var(--bg1);border:1px solid var(--border);
-  border-radius:10px;padding:12px 15px;margin-bottom:8px;
-  position:relative;overflow:hidden;
-  transition:box-shadow .15s,border-color .15s;
-}
-.ex-card::before{
-  content:'';position:absolute;left:0;top:0;bottom:0;width:3px;
-  background:linear-gradient(180deg,var(--accent),var(--accent2));
-  border-radius:3px 0 0 3px;
-}
+.ex-card{background:var(--bg1);border:1px solid var(--border);border-radius:10px;
+  padding:12px 15px;margin-bottom:8px;position:relative;overflow:hidden;
+  transition:box-shadow .15s,border-color .15s;}
+.ex-card::before{content:'';position:absolute;left:0;top:0;bottom:0;width:3px;
+  background:linear-gradient(180deg,var(--accent),var(--accent2));border-radius:3px 0 0 3px;}
 .ex-card:hover{border-color:rgba(37,99,235,.3);box-shadow:0 2px 12px var(--shadow);}
 .ex-pdb{font-family:var(--mono);font-size:.95rem;font-weight:700;color:var(--accent);}
 .ex-sub{font-size:.8rem;font-weight:600;color:var(--text1);margin:2px 0;}
 .ex-desc{font-size:.72rem;line-height:1.5;color:var(--text2);margin-top:3px;}
-
-.pkd-card{
-  background:linear-gradient(135deg,rgba(37,99,235,.05),rgba(124,58,237,.04));
-  border:1px solid var(--border);border-radius:12px;padding:14px 18px;
-  box-shadow:0 1px 6px var(--shadow);
-}
-.pkd-lbl{font-family:var(--mono);font-size:.67rem;color:var(--text2);
-  text-transform:uppercase;letter-spacing:.12em;margin-bottom:2px;}
-.pkd-val{font-family:var(--mono);font-size:2rem;font-weight:700;
-  color:var(--accent);line-height:1.1;}
-.pkd-badge{display:inline-block;padding:2px 9px;border-radius:20px;
-  font-size:.7rem;font-weight:600;font-family:var(--mono);margin-top:5px;}
+.pkd-card{background:linear-gradient(135deg,rgba(37,99,235,.05),rgba(124,58,237,.04));
+  border:1px solid var(--border);border-radius:12px;padding:14px 18px;box-shadow:0 1px 6px var(--shadow);}
+.pkd-lbl{font-family:var(--mono);font-size:.67rem;color:var(--text2);text-transform:uppercase;letter-spacing:.12em;margin-bottom:2px;}
+.pkd-val{font-family:var(--mono);font-size:2rem;font-weight:700;color:var(--accent);line-height:1.1;}
+.pkd-badge{display:inline-block;padding:2px 9px;border-radius:20px;font-size:.7rem;font-weight:600;font-family:var(--mono);margin-top:5px;}
 .badge-weak{background:rgba(220,38,38,.08);color:var(--red);border:1px solid rgba(220,38,38,.22);}
 .badge-moderate{background:rgba(217,119,6,.08);color:var(--amber);border:1px solid rgba(217,119,6,.22);}
 .badge-strong{background:rgba(22,163,74,.08);color:var(--green);border:1px solid rgba(22,163,74,.22);}
-
 .str-bar{height:5px;border-radius:3px;background:var(--bg3);overflow:hidden;margin:8px 0 3px;}
-.str-fill{height:100%;border-radius:3px;
-  background:linear-gradient(90deg,var(--accent),var(--accent2));
-  transition:width .7s cubic-bezier(.4,0,.2,1);}
-.str-labels{display:flex;justify-content:space-between;
-  font-size:.65rem;color:var(--text2);font-family:var(--mono);}
-
-.ready-badge{
-  display:inline-flex;align-items:center;gap:6px;
-  padding:3px 10px;border-radius:20px;
+.str-fill{height:100%;border-radius:3px;background:linear-gradient(90deg,var(--accent),var(--accent2));transition:width .7s cubic-bezier(.4,0,.2,1);}
+.str-labels{display:flex;justify-content:space-between;font-size:.65rem;color:var(--text2);font-family:var(--mono);}
+.ready-badge{display:inline-flex;align-items:center;gap:6px;padding:3px 10px;border-radius:20px;
   background:rgba(22,163,74,.07);border:1px solid rgba(22,163,74,.22);
-  font-family:var(--mono);font-size:.72rem;color:var(--green);font-weight:600;
-}
-.ready-dot{display:inline-block;width:6px;height:6px;border-radius:50%;
-  background:var(--green);animation:pulse 2s infinite;}
-.idle-badge{
-  display:inline-flex;align-items:center;gap:6px;
-  padding:3px 10px;border-radius:20px;
+  font-family:var(--mono);font-size:.72rem;color:var(--green);font-weight:600;}
+.ready-dot{display:inline-block;width:6px;height:6px;border-radius:50%;background:var(--green);animation:pulse 2s infinite;}
+.idle-badge{display:inline-flex;align-items:center;gap:6px;padding:3px 10px;border-radius:20px;
   background:rgba(100,116,139,.07);border:1px solid rgba(100,116,139,.18);
-  font-family:var(--mono);font-size:.72rem;color:var(--text2);font-weight:600;
-}
+  font-family:var(--mono);font-size:.72rem;color:var(--text2);font-weight:600;}
 @keyframes pulse{0%,100%{opacity:1;transform:scale(1)}50%{opacity:.5;transform:scale(.85)}}
-
-.model-section{background:var(--bg2);border:1px solid var(--border);
-  border-radius:9px;padding:12px 14px;margin-bottom:10px;}
-.model-section-title{font-family:var(--mono);font-size:.63rem;font-weight:700;
-  letter-spacing:.18em;text-transform:uppercase;color:var(--text2);margin-bottom:8px;}
-
-.ngl-placeholder{
-  display:flex;flex-direction:column;align-items:center;justify-content:center;
-  height:420px;background:var(--bg2);border:1px solid var(--border);
-  border-radius:12px;color:var(--text2);
-  font-family:var(--mono);font-size:.8rem;text-align:center;line-height:2;
-}
+.model-section{background:var(--bg2);border:1px solid var(--border);border-radius:9px;padding:12px 14px;margin-bottom:10px;}
+.model-section-title{font-family:var(--mono);font-size:.63rem;font-weight:700;letter-spacing:.18em;text-transform:uppercase;color:var(--text2);margin-bottom:8px;}
+.ngl-placeholder{display:flex;flex-direction:column;align-items:center;justify-content:center;
+  height:420px;background:var(--bg2);border:1px solid var(--border);border-radius:12px;
+  color:var(--text2);font-family:var(--mono);font-size:.8rem;text-align:center;line-height:2;}
+
+/* IG progress info box */
+.ig-info{background:linear-gradient(135deg,rgba(37,99,235,.04),rgba(124,58,237,.03));
+  border:1px solid var(--border);border-radius:8px;padding:10px 14px;
+  font-family:var(--mono);font-size:.75rem;color:var(--text2);margin-bottom:8px;}
 </style>
 """, unsafe_allow_html=True)
 
-# postMessage bridge so NGL iframes get theme-change events
+# postMessage bridge for NGL iframe theme sync
 st.markdown("""
 <script>
 (function(){
@@ -294,16 +190,14 @@ st.markdown("""
       try{f.contentWindow.postMessage({balmTheme:dark?'dark':'light'},'*');}catch(e){}
     });
   }
-  window.matchMedia('(prefers-color-scheme:dark)').addEventListener('change',function(e){
-    notify(e.matches);
-  });
+  window.matchMedia('(prefers-color-scheme:dark)').addEventListener('change',function(e){notify(e.matches);});
 })();
 </script>
 """, unsafe_allow_html=True)
 
 
 # ══════════════════════════════════════════════════════════════════════════════
-#  PDB UTILITIES  — PDB responses are cached on disk for 1 h
+#  PDB UTILITIES
 # ══════════════════════════════════════════════════════════════════════════════
 
 @st.cache_data(show_spinner=False, ttl=3600)
@@ -314,6 +208,23 @@ def _fetch_pdb_cached(pdb_id: str) -> str:
     return r.text
 
 
+def filter_pdb_to_chains(pdb_text: str, keep_chains: set) -> str:
+    """
+    Strip PDB text to only ATOM/HETATM records for chain IDs in keep_chains.
+    Preserves HEADER, TITLE, REMARK lines so NGL can still parse it.
+    This keeps the viewer uncluttered — only selected chains are shown.
+    """
+    out = []
+    for line in pdb_text.splitlines():
+        rec = line[:6].strip()
+        if rec in ("ATOM", "HETATM"):
+            chain_col = line[21:22]  # column 22 (0-indexed: 21) = chain ID
+            if chain_col not in keep_chains:
+                continue
+        out.append(line)
+    return "\n".join(out)
+
+
 def get_chains_from_pdb(pdb_id: str):
     try:
         pdb_text = _fetch_pdb_cached(pdb_id.strip().upper())
@@ -367,17 +278,18 @@ def populate_from_pdb(pdb_id, mode, chain_a=None, chain_b=None,
     if not pdb_content or not chains: return False
     available = sorted(chains.keys())
     st.info(f"✅ **{pdb_id.upper()}** — chains: **{', '.join(available)}**")
-    st.session_state.pdb_content  = pdb_content
-    st.session_state.ig_a         = None
-    st.session_state.ig_b         = None
-    st.session_state.ig_chain_map = None
-    st.session_state.result       = None
+
+    st.session_state.ig_a = st.session_state.ig_b = st.session_state.ig_chain_map = None
+    st.session_state.result = None
     st.session_state.chain_info_a = []
     st.session_state.chain_info_b = []
     st.session_state["_pm"] = mode
+
+    selected_chain_ids: set = set()   # ← collect which chains to keep in viewer
+
     if mode == "Protein-Protein":
         id_a = chain_a or available[0]
-        id_b = chain_b or (available[1] if len(available)>1 else available[0])
+        id_b = chain_b or (available[1] if len(available) > 1 else available[0])
         seq_a, info_a, miss_a = resolve_chains(chains, id_a, available)
         seq_b, info_b, miss_b = resolve_chains(chains, id_b, available)
         for m in miss_a: st.warning(f"Chain **{m}** not found (Side A). Available: {', '.join(available)}")
@@ -386,10 +298,11 @@ def populate_from_pdb(pdb_id, mode, chain_a=None, chain_b=None,
         st.session_state["_pb"] = seq_b
         st.session_state.chain_info_a = info_a
         st.session_state.chain_info_b = info_b
+        selected_chain_ids = {c["id"] for c in info_a} | {c["id"] for c in info_b}
     else:
-        id_h      = chain_h  or "H"
-        id_l      = chain_l  or "L"
-        id_ag_str = chain_ag or next((c for c in available if c not in (id_h,id_l)), available[0])
+        id_h      = chain_h or "H"
+        id_l      = chain_l or "L"
+        id_ag_str = chain_ag or next((c for c in available if c not in (id_h, id_l)), available[0])
         def _single(cid):
             if cid in chains and chains[cid]["seq"]: return cid, chains[cid]
             f = next((k for k in chains if k.upper()==cid.upper() and chains[k]["seq"]), None)
@@ -405,6 +318,17 @@ def populate_from_pdb(pdb_id, mode, chain_a=None, chain_b=None,
         st.session_state["_pag"] = seq_ag
         st.session_state.chain_info_a = info_ag if info_ag else []
         st.session_state.chain_info_b = [{"id":real_h,**ch},{"id":real_l,**cl}]
+        selected_chain_ids = ({c["id"] for c in info_ag} if info_ag else set())
+        if ch["seq"]: selected_chain_ids.add(real_h)
+        if cl["seq"]: selected_chain_ids.add(real_l)
+
+    # ── Filter PDB to selected chains only ───────────────────────────────────
+    if selected_chain_ids:
+        filtered = filter_pdb_to_chains(pdb_content, selected_chain_ids)
+        st.session_state.pdb_content = filtered
+    else:
+        st.session_state.pdb_content = pdb_content
+
     return True
 
 
@@ -444,12 +368,12 @@ def flat_ig_for_display(ig_a, ig_b, chain_info_a, chain_info_b):
     seq_a = "".join(c["seq"] for c in chain_info_a) if chain_info_a else ""
     seq_b = "".join(c["seq"] for c in chain_info_b) if chain_info_b else ""
     if ig_a and chain_info_a:
-        parts = split_ig_by_chains(ig_a,[c["seq"] for c in chain_info_a])
+        parts    = split_ig_by_chains(ig_a,[c["seq"] for c in chain_info_a])
         ig_a_out = [s for sub in parts for s in sub]
     else:
         ig_a_out = list(ig_a) if ig_a else []
     if ig_b and chain_info_b:
-        parts = split_ig_by_chains(ig_b,[c["seq"] for c in chain_info_b])
+        parts    = split_ig_by_chains(ig_b,[c["seq"] for c in chain_info_b])
         ig_b_out = [s for sub in parts for s in sub]
     else:
         ig_b_out = list(ig_b) if ig_b else []
@@ -458,22 +382,29 @@ def flat_ig_for_display(ig_a, ig_b, chain_info_a, chain_info_b):
 
 # ══════════════════════════════════════════════════════════════════════════════
 #  MODEL
+#  KEY FIX: do NOT use torch_dtype="auto" — on HF Spaces this loads fp16
+#  which silently produces zero gradients and breaks IG entirely.
+#  Always load in float32; the extra memory is worth correct attributions.
 # ══════════════════════════════════════════════════════════════════════════════
 @st.cache_resource(show_spinner="Loading ESM-2 backbone…")
 def _load_esm_base():
     from transformers import EsmModel, EsmTokenizer
     base = EsmModel.from_pretrained(
         "facebook/esm2_t33_650M_UR50D",
-        torch_dtype="auto", low_cpu_mem_usage=True, use_safetensors=True,
+        # torch_dtype="auto" removed: fp16 kills gradients on HF Spaces / CPU
+        # low_cpu_mem_usage removed: creates meta tensors → deepcopy+.to(device) crashes
+        use_safetensors=True,
     )
     tok = EsmTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
     return base, tok
 
+
 def _download_weights_hf() -> bytes:
     from huggingface_hub import hf_hub_download
     path = hf_hub_download(repo_id=HF_REPO_ID, filename=HF_FILENAME, resume_download=True)
     with open(path,"rb") as f: return f.read()
 
+
 def build_and_load_model(weights_bytes, pkd_bounds):
     import torch.nn as nn
     import torch.nn.functional as F
@@ -528,6 +459,7 @@ def build_and_load_model(weights_bytes, pkd_bounds):
     model.to(device).eval()
     return model, device
 
+
 def _ensure_model_loaded(pkd_lo=1.0, pkd_hi=16.0) -> bool:
     if st.session_state.model is not None: return True
     try:
@@ -544,46 +476,77 @@ def _ensure_model_loaded(pkd_lo=1.0, pkd_hi=16.0) -> bool:
 
 # ══════════════════════════════════════════════════════════════════════════════
 #  INTEGRATED GRADIENTS
+#  KEY FIXES:
+#   1. Explicit .float() casts — model may internally have mixed dtypes after LoRA
+#   2. Steps reduced 30 → 15 (Riemann IG converges fast; halves compute time)
+#   3. word_embed detach + clone to avoid stale graph issues
+#   4. Proper gradient zeroing between steps
 # ══════════════════════════════════════════════════════════════════════════════
-def compute_ig(model, seq_a, seq_b, steps=30):
+def compute_ig(model, seq_a: str, seq_b: str, steps: int = 15):
     device  = next(model.parameters()).device
     tok     = model.esm_tokenizer
     cls_tok = tok.cls_token
     esm     = model.esm_model
+
     def tokenise(seq):
-        proc = seq.replace("|",f"{cls_tok}{cls_tok}")
-        return tok(proc,return_tensors="pt",padding=False,truncation=True,max_length=1024).to(device)
+        proc = seq.replace("|", f"{cls_tok}{cls_tok}")
+        return tok(proc, return_tensors="pt",
+                   padding=False, truncation=True, max_length=1024).to(device)
+
     word_embed = esm.base_model.model.embeddings.word_embeddings
-    enc_a = tokenise(seq_a); mask_a = enc_a["attention_mask"]
-    enc_b = tokenise(seq_b); mask_b = enc_b["attention_mask"]
-    emb_a = word_embed(enc_a["input_ids"]).detach()
-    emb_b = word_embed(enc_b["input_ids"]).detach()
-    def encode(embs,mask):
-        ext = esm.base_model.model.get_extended_attention_mask(mask,embs.shape[:2])
-        h   = esm.base_model.model.encoder(embs,attention_mask=ext).last_hidden_state
-        m   = mask.unsqueeze(-1).expand(h.size()).float()
-        return (torch.sum(h*m,1)/torch.clamp(m.sum(1),min=1e-9)).float()
-    def fwd(e_a,e_b):
-        return model.projection_head(encode(e_a,mask_a),encode(e_b,mask_b))
-    def riemann(tgt,baseline,fixed,tgt_is_b):
+
+    enc_a = tokenise(seq_a);  mask_a = enc_a["attention_mask"]
+    enc_b = tokenise(seq_b);  mask_b = enc_b["attention_mask"]
+
+    # ── Force float32 — critical on HF Spaces where fp16 zeros gradients ────
+    emb_a = word_embed(enc_a["input_ids"]).detach().float().clone()
+    emb_b = word_embed(enc_b["input_ids"]).detach().float().clone()
+    mask_a = mask_a.float()
+    mask_b = mask_b.float()
+
+    def encode(embs, mask):
+        # Cast mask for attention
+        int_mask = mask.long()
+        ext = esm.base_model.model.get_extended_attention_mask(int_mask, embs.shape[:2])
+        # Run encoder — force float32 throughout
+        h = esm.base_model.model.encoder(
+            embs.float(), attention_mask=ext
+        ).last_hidden_state.float()
+        m = mask.unsqueeze(-1).expand(h.size()).float()
+        return (torch.sum(h * m, 1) / torch.clamp(m.sum(1), min=1e-9))
+
+    def fwd(e_a, e_b):
+        return model.projection_head(encode(e_a, mask_a), encode(e_b, mask_b))
+
+    def riemann(tgt, baseline, fixed, tgt_is_b):
         grads = []
-        for a in torch.linspace(0,1,steps,device=device):
-            x = (baseline+a*(tgt-baseline)).detach().requires_grad_(True)
-            out = fwd(fixed,x) if tgt_is_b else fwd(x,fixed)
+        for alpha in torch.linspace(0, 1, steps, device=device):
+            interp = (baseline + alpha * (tgt - baseline)).detach().requires_grad_(True)
+            out = fwd(fixed, interp) if tgt_is_b else fwd(interp, fixed)
             out.sum().backward()
-            grads.append(x.grad.detach().clone())
-        avg = torch.stack(grads).mean(0)
-        return (avg*(tgt-baseline)).abs().sum(-1).squeeze(0).cpu().numpy()
-    attr_a = riemann(emb_a,torch.zeros_like(emb_a),emb_b,False)
-    attr_b = riemann(emb_b,torch.zeros_like(emb_b),emb_a,True)
+            grads.append(interp.grad.detach().float().clone())
+        avg_grad = torch.stack(grads).mean(0)
+        ig = (avg_grad * (tgt - baseline)).abs().sum(-1).squeeze(0)
+        return ig.cpu().numpy()
+
+    baseline_a = torch.zeros_like(emb_a)
+    baseline_b = torch.zeros_like(emb_b)
+
+    attr_a = riemann(emb_a, baseline_a, emb_b, False)
+    attr_b = riemann(emb_b, baseline_b, emb_a, True)
+
     def norm(a):
-        lo,hi = a.min(),a.max()
-        return ((a-lo)/(hi-lo+1e-9)).tolist()
+        lo, hi = a.min(), a.max()
+        if hi - lo < 1e-9:
+            return [0.0] * len(a)   # flat → return zeros, not NaN
+        return ((a - lo) / (hi - lo)).tolist()
+
+    # Slice off [CLS] and [EOS] tokens (indices 0 and -1)
     return norm(attr_a[1:-1]), norm(attr_b[1:-1])
 
 
 # ══════════════════════════════════════════════════════════════════════════════
-#  NGL VIEWER — theme-reactive light ↔ dark
+#  NGL VIEWER — theme-reactive, shows only filtered PDB chains
 # ══════════════════════════════════════════════════════════════════════════════
 def ngl_viewer_html(pdb_content, ig_chain_map=None, height=420) -> str:
     ig_json = json.dumps(ig_chain_map) if ig_chain_map else "null"
@@ -595,8 +558,7 @@ def ngl_viewer_html(pdb_content, ig_chain_map=None, height=420) -> str:
 *{{box-sizing:border-box;margin:0;padding:0}}
 body{{overflow:hidden;font-family:'JetBrains Mono',monospace;transition:background .35s}}
 #vp{{width:100%;height:{height}px}}
-#hint{{position:absolute;top:10px;left:12px;font-size:10px;pointer-events:none;
-       letter-spacing:.04em;transition:color .3s}}
+#hint{{position:absolute;top:10px;left:12px;font-size:10px;pointer-events:none;letter-spacing:.04em;transition:color .3s}}
 #ctrl{{position:absolute;bottom:12px;right:12px;display:flex;flex-direction:column;gap:4px}}
 .cb{{padding:4px 11px;font-family:'JetBrains Mono',monospace;font-size:9px;font-weight:500;
      border-radius:5px;cursor:pointer;letter-spacing:.1em;text-transform:uppercase;transition:all .15s}}
@@ -613,57 +575,39 @@ body{{overflow:hidden;font-family:'JetBrains Mono',monospace;transition:backgrou
 </div>
 <div id="leg"></div>
 <script>
-var T = {{
-  light:{{
-    bg:'#f0f4fa',
-    hint:'rgba(51,65,85,.5)',
-    btn:'rgba(240,244,250,.93)','btn_b':'rgba(37,99,235,.16)',
-    btn_c:'rgba(71,85,105,.6)',
+var T={{
+  light:{{bg:'#f0f4fa',hint:'rgba(51,65,85,.5)',
+    btn:'rgba(240,244,250,.93)',btn_b:'rgba(37,99,235,.16)',btn_c:'rgba(71,85,105,.6)',
     btn_hc:'#2563eb',btn_hb:'rgba(37,99,235,.55)',btn_hbg:'rgba(37,99,235,.06)',
     leg:'rgba(240,244,250,.93)',leg_b:'rgba(37,99,235,.13)',leg_c:'rgba(71,85,105,.65)',
-    ig_hi:'#1d4ed8',ig_lo:'#c7d7f0',
-    null_col:0xe8ecf4,
-  }},
-  dark:{{
-    bg:'#06090f',
-    hint:'rgba(168,184,216,.4)',
-    btn:'rgba(6,9,15,.88)',btn_b:'rgba(96,165,250,.16)',
-    btn_c:'rgba(168,184,216,.55)',
+    ig_hi:'#1d4ed8',ig_lo:'#c7d7f0',null_col:0xe8ecf4}},
+  dark:{{bg:'#06090f',hint:'rgba(168,184,216,.4)',
+    btn:'rgba(6,9,15,.88)',btn_b:'rgba(96,165,250,.16)',btn_c:'rgba(168,184,216,.55)',
     btn_hc:'#60a5fa',btn_hb:'rgba(96,165,250,.55)',btn_hbg:'rgba(96,165,250,.06)',
     leg:'rgba(6,9,15,.88)',leg_b:'rgba(96,165,250,.13)',leg_c:'rgba(168,184,216,.55)',
-    ig_hi:'#60a5fa',ig_lo:'#172135',
-    null_col:0x172135,
-  }}
+    ig_hi:'#60a5fa',ig_lo:'#172135',null_col:0x172135}}
 }};
-
 function isDark(){{return window.matchMedia('(prefers-color-scheme:dark)').matches;}}
-var theme = isDark() ? T.dark : T.light;
-var stage = null, comp = null, curRep = 'cartoon';
-var igData = {ig_json};
+var theme=isDark()?T.dark:T.light,stage=null,comp=null,curRep='cartoon',igData={ig_json};
 
-function hex(s,dark){{
-  // light: #c7d7f0(199,215,240) → #1d4ed8(29,78,216)
-  // dark:  #172135(23,33,53)    → #60a5fa(96,165,250)
+function igColor(s,dark){{
   s=Math.max(0,Math.min(1,s||0));
-  if(dark) return (Math.round(23+s*73)<<16)|(Math.round(33+s*132)<<8)|Math.round(53+s*197);
-  return (Math.round(199-s*170)<<16)|(Math.round(215-s*137)<<8)|Math.round(240-s*24);
+  if(dark) return(Math.round(23+s*73)<<16)|(Math.round(33+s*132)<<8)|Math.round(53+s*197);
+  return(Math.round(199-s*170)<<16)|(Math.round(215-s*137)<<8)|Math.round(240-s*24);
 }}
-
 function styleUI(dark){{
-  var Th = dark ? T.dark : T.light;
-  theme = Th;
-  document.body.style.background = Th.bg;
-  document.getElementById('hint').style.color = Th.hint;
+  var Th=dark?T.dark:T.light; theme=Th;
+  document.body.style.background=Th.bg;
+  document.getElementById('hint').style.color=Th.hint;
   document.querySelectorAll('.cb').forEach(function(b){{
     b.style.background=Th.btn;b.style.border='1px solid '+Th.btn_b;b.style.color=Th.btn_c;
     b.onmouseenter=function(){{b.style.borderColor=Th.btn_hb;b.style.color=Th.btn_hc;b.style.background=Th.btn_hbg;}};
-    b.onmouseleave=function(){{b.style.borderColor=Th.btn_b;b.style.color=Th.btn_c;b.style.background=Th.btn;  }};
+    b.onmouseleave=function(){{b.style.borderColor=Th.btn_b;b.style.color=Th.btn_c;b.style.background=Th.btn;}};
   }});
   var leg=document.getElementById('leg');
   leg.style.background=Th.leg;leg.style.border='1px solid '+Th.leg_b;leg.style.color=Th.leg_c;
-  if(stage) stage.setParameters({{backgroundColor:Th.bg}});
+  if(stage)stage.setParameters({{backgroundColor:Th.bg}});
 }}
-
 function addRepr(c){{
   comp=c;comp.removeAllRepresentations();
   var dark=isDark(),Th=dark?T.dark:T.light,leg=document.getElementById('leg');
@@ -673,43 +617,31 @@ function addRepr(c){{
         var cd=igData[atom.chainname];
         if(!cd||!cd.scores||!cd.scores.length) return Th.null_col;
         var i=atom.resno-cd.start;
-        return(i<0||i>=cd.scores.length)?Th.null_col:hex(cd.scores[i],dark);
+        return(i<0||i>=cd.scores.length)?Th.null_col:igColor(cd.scores[i],dark);
       }};
     }});
     comp.addRepresentation(curRep,{{color:sid}});
     leg.innerHTML='<span style="color:'+Th.ig_hi+'">■</span> High IG &nbsp;&nbsp;<span style="color:'+Th.ig_lo+'">■</span> Low IG';
-  }} else {{
+  }}else{{
     comp.addRepresentation(curRep,{{colorScheme:'chainname'}});
     leg.innerHTML='Coloured by chain';
   }}
   stage.autoView();
 }}
-
 function setRep(r){{curRep=r;if(comp){{comp.removeAllRepresentations();addRepr(comp);}}}}
-
-function themeSwitch(dark){{
-  styleUI(dark);
-  if(comp){{comp.removeAllRepresentations();addRepr(comp);}}
-}}
-
-window.matchMedia('(prefers-color-scheme:dark)').addEventListener('change',function(e){{
-  themeSwitch(e.matches);
-}});
-window.addEventListener('message',function(e){{
-  if(e.data&&e.data.balmTheme) themeSwitch(e.data.balmTheme==='dark');
-}});
+function themeSwitch(dark){{styleUI(dark);if(comp){{comp.removeAllRepresentations();addRepr(comp);}}}}
+window.matchMedia('(prefers-color-scheme:dark)').addEventListener('change',function(e){{themeSwitch(e.matches);}});
+window.addEventListener('message',function(e){{if(e.data&&e.data.balmTheme)themeSwitch(e.data.balmTheme==='dark');}});
 window.addEventListener('resize',function(){{if(stage)stage.handleResize();}});
-
-stage = new NGL.Stage('vp',{{backgroundColor:theme.bg,quality:'medium',tooltip:true}});
+stage=new NGL.Stage('vp',{{backgroundColor:theme.bg,quality:'medium',tooltip:true}});
 styleUI(isDark());
-
 var blob=new Blob([`{escaped}`],{{type:'text/plain'}});
 stage.loadFile(URL.createObjectURL(blob),{{ext:'pdb',name:'s'}}).then(addRepr);
 </script></body></html>"""
 
 
 # ══════════════════════════════════════════════════════════════════════════════
-#  PLOTLY — light theme (matches default)
+#  PLOTLY
 # ══════════════════════════════════════════════════════════════════════════════
 _PL   = dict(paper_bgcolor="rgba(255,255,255,.97)",plot_bgcolor="rgba(241,244,249,1)")
 _FONT = dict(family="JetBrains Mono, monospace",color="#475569",size=9)
@@ -725,8 +657,7 @@ def make_heatmap(seq, attr, title):
     fig = go.Figure(go.Heatmap(z=rows,text=hover,hoverinfo="text",
         colorscale=[[0,"#e8f0fe"],[.4,"#93c5fd"],[.75,"#3b82f6"],[1,"#1d4ed8"]],
         zmin=0,zmax=1,xgap=1,ygap=1,
-        colorbar=dict(thickness=10,len=.9,tickfont=dict(**_FONT),
-                      title=dict(text="IG",font=dict(**_FONT)))))
+        colorbar=dict(thickness=10,len=.9,tickfont=dict(**_FONT),title=dict(text="IG",font=dict(**_FONT)))))
     fig.update_layout(
         title=dict(text=title,font=dict(family="JetBrains Mono, monospace",size=10,color="#1d4ed8")),
         **_PL,height=200,margin=dict(t=30,b=24,l=55,r=45),
@@ -768,9 +699,6 @@ def residue_strip_html(seq, attr) -> str:
             '<div style="line-height:2.2;word-break:break-all">'+"".join(cells)+"</div>")
 
 
-# ══════════════════════════════════════════════════════════════════════════════
-#  BATCH TEMPLATES
-# ══════════════════════════════════════════════════════════════════════════════
 PPI_TEMPLATE  = "seq_a,seq_b\nACDEFGHIKLMNPQRSTVWY,QWERTYIPASDFGKLCVBNM\n"
 ABAG_TEMPLATE = "heavy_chain,light_chain,antigen\nEVQLVESGGG...,DIQMTQ...,IYSPT...\n"
 
@@ -786,14 +714,10 @@ def main():
         if _src in st.session_state:
             st.session_state[_dst] = st.session_state.pop(_src)
 
-    # Auto-load model from example click
     if st.session_state.get("_auto_load_model") and st.session_state.model is None:
         st.session_state["_auto_load_model"] = False
-        _ensure_model_loaded(
-            st.session_state.get("pkd_lo",1.0),
-            st.session_state.get("pkd_hi",16.0))
+        _ensure_model_loaded(st.session_state.get("pkd_lo",1.0),st.session_state.get("pkd_hi",16.0))
 
-    # Header
     st.markdown("""
     <div class="app-header">
       <div class="app-logo">🧬</div>
@@ -801,15 +725,14 @@ def main():
         <div class="app-title">BALM-PPI Pro</div>
         <div class="app-subtitle">ESM-2 · LoRA · Integrated Gradients &nbsp;·&nbsp; Protein Binding Affinity Prediction</div>
       </div>
-    </div>
-    """, unsafe_allow_html=True)
+    </div>""", unsafe_allow_html=True)
     st.divider()
 
-    # ── SIDEBAR ──────────────────────────────────────────────────────────────
+    # SIDEBAR
     with st.sidebar:
         st.markdown("### ⚙️ Model")
-        st.markdown('<div class="model-section">', unsafe_allow_html=True)
-        st.markdown('<div class="model-section-title">Weights Source</div>', unsafe_allow_html=True)
+        st.markdown('<div class="model-section">',unsafe_allow_html=True)
+        st.markdown('<div class="model-section-title">Weights Source</div>',unsafe_allow_html=True)
         model_src = st.radio("src",["HuggingFace (auto-cached)","Upload .pth"],
                              key="model_src",label_visibility="collapsed")
         custom_w = None
@@ -826,12 +749,11 @@ def main():
 
         if st.button("⚡ Load Model",use_container_width=True,type="primary"):
             try:
-                if model_src == "HuggingFace (auto-cached)":
+                if model_src=="HuggingFace (auto-cached)":
                     with st.spinner("📥 Fetching weights (cached after first run)…"):
                         wb = _download_weights_hf()
                 else:
-                    if custom_w is None:
-                        st.error("Upload a .pth file first."); st.stop()
+                    if custom_w is None: st.error("Upload a .pth file first."); st.stop()
                     wb = custom_w.read()
                 with st.spinner("🔧 Building ESM-2 + LoRA…"):
                     m,dev = build_and_load_model(wb,(pkd_lo,pkd_hi))
@@ -842,10 +764,7 @@ def main():
                 st.error(f"❌ {e}"); st.code(traceback.format_exc())
 
         if st.session_state.model:
-            st.markdown(
-                f'<div class="ready-badge"><span class="ready-dot"></span>'
-                f'READY &nbsp;·&nbsp; {st.session_state.device}</div>',
-                unsafe_allow_html=True)
+            st.markdown(f'<div class="ready-badge"><span class="ready-dot"></span>READY &nbsp;·&nbsp; {st.session_state.device}</div>',unsafe_allow_html=True)
         else:
             st.markdown('<div class="idle-badge">○ &nbsp;NOT LOADED</div>',unsafe_allow_html=True)
 
@@ -854,12 +773,10 @@ def main():
             res = st.session_state.result
             st.markdown("**Latest Result**")
             st.markdown(
-                f'<div class="pkd-card">'
-                f'<div class="pkd-lbl">Predicted pKd</div>'
+                f'<div class="pkd-card"><div class="pkd-lbl">Predicted pKd</div>'
                 f'<div class="pkd-val">{res["pkd"]:.3f}</div>'
                 f'<div class="pkd-lbl" style="margin-top:10px">Cosine Similarity</div>'
-                f'<div style="font-family:var(--mono);font-size:1rem;font-weight:700;color:var(--text0)">'
-                f'{res["cosine"]:.4f}</div></div>',
+                f'<div style="font-family:var(--mono);font-size:1rem;font-weight:700;color:var(--text0)">{res["cosine"]:.4f}</div></div>',
                 unsafe_allow_html=True)
 
         st.divider()
@@ -869,17 +786,15 @@ def main():
             '💻 <a href="https://github.com/rgorantla04/BALM-PPI" style="color:var(--accent);text-decoration:none">rgorantla04/BALM-PPI</a>'
             '</div>',unsafe_allow_html=True)
 
-    # ── TABS ──────────────────────────────────────────────────────────────────
-    tab_single, tab_batch, tab_info = st.tabs(
-        ["🎯 Single Prediction","📂 Batch Prediction","📚 Model Info"])
+    # TABS
+    tab_single,tab_batch,tab_info = st.tabs(["🎯 Single Prediction","📂 Batch Prediction","📚 Model Info"])
 
     with tab_single:
         mode_opts = ["Protein-Protein","Antibody-Antigen"]
-        mode = st.radio("im_r",mode_opts,horizontal=True,
-                        label_visibility="collapsed",key="interaction_mode")
+        mode = st.radio("im_r",mode_opts,horizontal=True,label_visibility="collapsed",key="interaction_mode")
         st.markdown("---")
 
-        if mode == "Protein-Protein":
+        if mode=="Protein-Protein":
             sc1,sc2 = st.columns(2)
             with sc1:
                 st.markdown("**Target Protein — Seq A**")
@@ -889,8 +804,7 @@ def main():
                 st.markdown("**Binder Protein — Seq B** &nbsp; `|` = chain separator")
                 seq_b_raw = st.text_area("Binder",label_visibility="collapsed",key="ta_ppi_b",
                     height=140,placeholder="Paste FASTA or raw sequence…")
-            seq_a = clean_multi(seq_a_raw)
-            seq_b = clean_multi(seq_b_raw)
+            seq_a = clean_multi(seq_a_raw); seq_b = clean_multi(seq_b_raw)
             if not st.session_state.pdb_content:
                 if seq_a:
                     st.session_state.chain_info_a = [
@@ -904,26 +818,20 @@ def main():
             sc1,sc2 = st.columns([2,3])
             with sc1:
                 st.markdown("**Antigen / Target — Seq A**")
-                ag_raw = st.text_area("Antigen",label_visibility="collapsed",
-                    key="ta_ab_ag",height=140,placeholder="Antigen sequence…")
+                ag_raw = st.text_area("Antigen",label_visibility="collapsed",key="ta_ab_ag",height=140,placeholder="Antigen sequence…")
             with sc2:
                 ah,al = st.columns(2)
                 with ah:
                     st.markdown("**Heavy Chain (H)**")
-                    h_raw = st.text_area("Heavy",label_visibility="collapsed",
-                        key="ta_ab_h",height=140,placeholder="VH sequence…")
+                    h_raw = st.text_area("Heavy",label_visibility="collapsed",key="ta_ab_h",height=140,placeholder="VH sequence…")
                 with al:
                     st.markdown("**Light Chain (L)**")
-                    l_raw = st.text_area("Light",label_visibility="collapsed",
-                        key="ta_ab_l",height=140,placeholder="VL sequence…")
-            seq_a = clean_multi(ag_raw)
-            h_seq = clean_multi(h_raw)
-            l_seq = clean_multi(l_raw)
+                    l_raw = st.text_area("Light",label_visibility="collapsed",key="ta_ab_l",height=140,placeholder="VL sequence…")
+            seq_a = clean_multi(ag_raw); h_seq = clean_multi(h_raw); l_seq = clean_multi(l_raw)
             seq_b = f"{h_seq}|{l_seq}" if (h_seq or l_seq) else ""
             if not st.session_state.pdb_content:
                 if seq_a:
-                    st.session_state.chain_info_a = [
-                        {"id":"Ag","seq":seq_a,"resnos":list(range(1,len(seq_a)+1))}]
+                    st.session_state.chain_info_a = [{"id":"Ag","seq":seq_a,"resnos":list(range(1,len(seq_a)+1))}]
                 if h_seq or l_seq:
                     st.session_state.chain_info_b = [
                         {"id":"H","seq":h_seq,"resnos":list(range(1,len(h_seq)+1))},
@@ -936,12 +844,10 @@ def main():
             if st.session_state.model is None:
                 st.caption("⬅ Load model in sidebar first")
         with rc2:
-            run_ig = st.checkbox("Run Integrated Gradients",value=True,
-                                 help="~1–2 min CPU · <10 s GPU")
+            run_ig = st.checkbox("Run Integrated Gradients",value=True,help="~45-90s CPU · <5s GPU")
         with rc3:
             if st.session_state.model:
-                st.markdown('<div class="ready-badge" style="margin-top:6px;font-size:.7rem">'
-                            '<span class="ready-dot"></span>READY</div>',unsafe_allow_html=True)
+                st.markdown('<div class="ready-badge" style="margin-top:6px;font-size:.7rem"><span class="ready-dot"></span>READY</div>',unsafe_allow_html=True)
 
         if run_btn:
             if not seq_a or not seq_b:
@@ -958,37 +864,37 @@ def main():
                     except Exception as e:
                         st.error(f"Prediction failed: {e}"); st.code(traceback.format_exc())
                 if run_ig and st.session_state.result:
-                    with st.spinner("Computing Integrated Gradients…"):
+                    with st.spinner("Computing Integrated Gradients (15 steps)…"):
                         try:
-                            ig_a,ig_b = compute_ig(st.session_state.model,seq_a,seq_b)
+                            ig_a,ig_b = compute_ig(st.session_state.model,seq_a,seq_b,steps=15)
                             st.session_state.ig_a = ig_a
                             st.session_state.ig_b = ig_b
                             if st.session_state.chain_info_a or st.session_state.chain_info_b:
                                 st.session_state.ig_chain_map = build_ig_chain_map(
                                     st.session_state.chain_info_a,st.session_state.chain_info_b,ig_a,ig_b)
                         except Exception as e:
-                            st.warning(f"IG failed: {e}")
+                            st.warning(f"IG failed: {e}\n{traceback.format_exc()}")
                 st.rerun()
 
         if st.session_state.result:
             res = st.session_state.result
             pkd,cos = res["pkd"],res["cosine"]
             pct = max(0.0,min(100.0,((pkd-pkd_lo)/max(pkd_hi-pkd_lo,1e-9))*100))
-            strength  = ("Weak" if pct<30 else "Moderate" if pct<55 else "Strong" if pct<75 else "Very Strong")
-            badge_cls = ("badge-weak" if pct<30 else "badge-moderate" if pct<55 else "badge-strong")
+            strength  = "Weak" if pct<30 else "Moderate" if pct<55 else "Strong" if pct<75 else "Very Strong"
+            badge_cls = "badge-weak" if pct<30 else "badge-moderate" if pct<55 else "badge-strong"
             st.markdown("---")
             m1,m2,m3 = st.columns([1,1,2])
             m1.metric("Predicted pKd",f"{pkd:.3f}")
             m2.metric("Cosine Similarity",f"{cos:.4f}")
             with m3:
                 st.markdown(
-                    f'<div class="pkd-card">'
-                    f'<div class="str-labels"><span>Weak ({pkd_lo:.0f})</span>'
+                    f'<div class="pkd-card"><div class="str-labels">'
+                    f'<span>Weak ({pkd_lo:.0f})</span>'
                     f'<span style="color:var(--text0);font-weight:600">{strength}</span>'
                     f'<span>Strong ({pkd_hi:.0f})</span></div>'
                     f'<div class="str-bar"><div class="str-fill" style="width:{pct:.1f}%"></div></div>'
-                    f'<span class="pkd-badge {badge_cls}">{strength}</span>'
-                    f'</div>',unsafe_allow_html=True)
+                    f'<span class="pkd-badge {badge_cls}">{strength}</span></div>',
+                    unsafe_allow_html=True)
 
         st.markdown("---")
         mode_now = st.session_state.get("interaction_mode","Protein-Protein")
@@ -1013,13 +919,11 @@ def main():
                 st.markdown(
                     f'<div class="ex-card"><div class="ex-pdb">{ex["label"]}</div>'
                     f'<div class="ex-sub">{ex["subtitle"]}</div>'
-                    f'<div class="ex-desc">{ex["desc"]}</div></div>',
-                    unsafe_allow_html=True)
+                    f'<div class="ex-desc">{ex["desc"]}</div></div>',unsafe_allow_html=True)
                 if st.button(f"⬇ Load {ex['pdb']}",key=f"ex_{ex['pdb']}",use_container_width=True):
                     ok = populate_from_pdb(ex["pdb"],ex["mode"],
                         chain_a=ex.get("chain_a"),chain_b=ex.get("chain_b"),
-                        chain_h=ex.get("chain_h"),chain_l=ex.get("chain_l"),
-                        chain_ag=ex.get("chain_ag"))
+                        chain_h=ex.get("chain_h"),chain_l=ex.get("chain_l"),chain_ag=ex.get("chain_ag"))
                     if ok:
                         if st.session_state.model is None:
                             st.session_state["_auto_load_model"] = True
@@ -1029,7 +933,7 @@ def main():
             st.markdown('<div class="sec-hdr">Custom PDB Fetch</div>',unsafe_allow_html=True)
             pdb_in = st.text_input("PDB ID",placeholder="6M0J, 1BRS, 2VXQ…",key="pdb_id_input")
 
-            if mode == "Protein-Protein":
+            if mode=="Protein-Protein":
                 st.caption("Comma-separate chains for multi-chain, e.g. `B,C`")
                 fc1,fc2 = st.columns(2)
                 ca_in = fc1.text_input("Side A chain(s)",placeholder="A or A,B",key="ca_in")
@@ -1055,34 +959,29 @@ def main():
                         if ok: st.rerun()
 
         with right_col:
-            vt = st.tabs([
-                "🧬 Structure",
-                f"📊 {lbl_a} Strip",f"🗺 {lbl_a} Heatmap",
-                f"📊 {lbl_b} Strip",f"🗺 {lbl_b} Heatmap",
-                "📥 Download",
-            ])
+            vt = st.tabs(["🧬 Structure",
+                          f"📊 {lbl_a} Strip",f"🗺 {lbl_a} Heatmap",
+                          f"📊 {lbl_b} Strip",f"🗺 {lbl_b} Heatmap",
+                          "📥 Download"])
 
             with vt[0]:
                 if st.session_state.pdb_content:
-                    st.caption("🎨 IG attribution" if st.session_state.ig_chain_map
-                               else "🔗 Coloured by chain — run prediction with IG for attribution")
+                    st.caption("🎨 IG attribution (selected chains only)" if st.session_state.ig_chain_map
+                               else "🔗 Selected chains — run prediction with IG for attribution colouring")
                     st.components.v1.html(
-                        ngl_viewer_html(st.session_state.pdb_content,
-                                        st.session_state.ig_chain_map,height=420),
+                        ngl_viewer_html(st.session_state.pdb_content,st.session_state.ig_chain_map,height=420),
                         height=425,scrolling=False)
                 else:
                     st.markdown(
-                        '<div class="ngl-placeholder">'
-                        '<div style="font-size:2.2rem">🧬</div>'
+                        '<div class="ngl-placeholder"><div style="font-size:2.2rem">🧬</div>'
                         '<div style="margin-top:12px;font-weight:600">Load an example or fetch a PDB</div>'
-                        '<div style="font-size:.72rem;margin-top:4px">Structure will appear here</div>'
+                        '<div style="font-size:.72rem;margin-top:4px">Only selected chains will be shown</div>'
                         '</div>',unsafe_allow_html=True)
 
             with vt[1]:
                 if ig_a_d and seq_a_d:
                     st.markdown(residue_strip_html(seq_a_d,ig_a_d),unsafe_allow_html=True)
-                    st.plotly_chart(top10_bar(seq_a_d,ig_a_d,"#2563eb",f"Top 10 · {lbl_a}"),
-                                   use_container_width=True)
+                    st.plotly_chart(top10_bar(seq_a_d,ig_a_d,"#2563eb",f"Top 10 · {lbl_a}"),use_container_width=True)
                 else:
                     st.info("Run prediction with **Integrated Gradients** to see attribution.")
 
@@ -1096,8 +995,7 @@ def main():
             with vt[3]:
                 if ig_b_d and seq_b_d:
                     st.markdown(residue_strip_html(seq_b_d,ig_b_d),unsafe_allow_html=True)
-                    st.plotly_chart(top10_bar(seq_b_d,ig_b_d,"#7c3aed",f"Top 10 · {lbl_b}"),
-                                   use_container_width=True)
+                    st.plotly_chart(top10_bar(seq_b_d,ig_b_d,"#7c3aed",f"Top 10 · {lbl_b}"),use_container_width=True)
                 else:
                     st.info("Run prediction with **Integrated Gradients** to see attribution.")
 
@@ -1112,37 +1010,27 @@ def main():
                 if st.session_state.result:
                     if ig_a_d or ig_b_d:
                         dl1,dl2 = st.columns(2)
-                        buf = sio.StringIO()
-                        wc  = csv.writer(buf)
+                        buf = sio.StringIO(); wc = csv.writer(buf)
                         wc.writerow(["chain","position","residue","ig_score"])
-                        for i,(aa,sc) in enumerate(zip(seq_a_d or "",ig_a_d or [])):
-                            wc.writerow(["Target",i+1,aa,f"{sc:.6f}"])
-                        for i,(aa,sc) in enumerate(zip(seq_b_d or "",ig_b_d or [])):
-                            wc.writerow(["proteina",i+1,aa,f"{sc:.6f}"])
-                        dl1.download_button("📥 IG Scores CSV",buf.getvalue().encode(),
-                                            "ig_scores.csv","text/csv",use_container_width=True)
-                        res_dl  = st.session_state.result
+                        for i,(aa,sc) in enumerate(zip(seq_a_d or "",ig_a_d or [])): wc.writerow(["Target",i+1,aa,f"{sc:.6f}"])
+                        for i,(aa,sc) in enumerate(zip(seq_b_d or "",ig_b_d or [])): wc.writerow(["proteina",i+1,aa,f"{sc:.6f}"])
+                        dl1.download_button("📥 IG Scores CSV",buf.getvalue().encode(),"ig_scores.csv","text/csv",use_container_width=True)
+                        res_dl = st.session_state.result
                         summary = {
                             "pkd":res_dl["pkd"],"cosine":res_dl["cosine"],
-                            "Target_length":len(seq_a_d or ""),
-                            "proteina_length":len(seq_b_d or ""),
+                            "Target_length":len(seq_a_d or ""),"proteina_length":len(seq_b_d or ""),
                             "top10_Target":sorted([{"res":f"{(seq_a_d or '')[i]}{i+1}","ig":(ig_a_d or [])[i]}
-                                for i in range(min(len(seq_a_d or ""),len(ig_a_d or [])))],
-                                key=lambda x:-x["ig"])[:10],
+                                for i in range(min(len(seq_a_d or ""),len(ig_a_d or [])))],key=lambda x:-x["ig"])[:10],
                             "top10_proteina":sorted([{"res":f"{(seq_b_d or '')[i]}{i+1}","ig":(ig_b_d or [])[i]}
-                                for i in range(min(len(seq_b_d or ""),len(ig_b_d or [])))],
-                                key=lambda x:-x["ig"])[:10],
+                                for i in range(min(len(seq_b_d or ""),len(ig_b_d or [])))],key=lambda x:-x["ig"])[:10],
                         }
-                        dl2.download_button("📥 Summary JSON",
-                                            json.dumps(summary,indent=2).encode(),
-                                            "balm_result.json","application/json",
-                                            use_container_width=True)
+                        dl2.download_button("📥 Summary JSON",json.dumps(summary,indent=2).encode(),"balm_result.json","application/json",use_container_width=True)
                     else:
                         res_dl = st.session_state.result
                         st.info("Enable IG and re-run for per-residue CSV.")
                         st.download_button("📥 Basic Result JSON",
-                                           json.dumps({"pkd":res_dl["pkd"],"cosine":res_dl["cosine"]},indent=2).encode(),
-                                           "balm_result.json","application/json")
+                            json.dumps({"pkd":res_dl["pkd"],"cosine":res_dl["cosine"]},indent=2).encode(),
+                            "balm_result.json","application/json")
                 else:
                     st.info("Run a prediction to enable download.")
 
@@ -1161,8 +1049,7 @@ def main():
             df = pd.read_csv(bf)
             st.dataframe(df.head(),use_container_width=True)
             if st.button("🏃 Run Batch",type="primary",use_container_width=True):
-                if st.session_state.model is None:
-                    st.error("Load the model first.")
+                if st.session_state.model is None: st.error("Load the model first.")
                 else:
                     model,results = st.session_state.model,[]
                     prog,stat,n   = st.progress(0.0),st.empty(),len(df)
@@ -1180,7 +1067,7 @@ def main():
                                 pkd_t,cos_t = model(sa,sb)
                             rr={"pKd":round(float(pkd_t.item()),4),"Cosine":round(float(cos_t.item()),6)}
                             if run_ig_b:
-                                ia,ib=compute_ig(model,sa,sb)
+                                ia,ib=compute_ig(model,sa,sb,steps=15)
                                 sap,sbp=sa.replace("|",""),sb.replace("|","")
                                 rr["top5_Target"]="|".join(f"{sap[j]}{j+1}:{ia[j]:.3f}"
                                     for j in sorted(range(min(len(sap),len(ia))),key=lambda x:-ia[x])[:5])
@@ -1193,14 +1080,13 @@ def main():
                     df_res = pd.concat([df,pd.DataFrame(results)],axis=1)
                     st.success(f"✅ {n} predictions complete")
                     st.dataframe(df_res,use_container_width=True)
-                    st.download_button("📥 Download Results CSV",df_res.to_csv(index=False).encode(),
-                                       "balm_batch_results.csv","text/csv",use_container_width=True)
+                    st.download_button("📥 Download Results CSV",df_res.to_csv(index=False).encode(),"balm_batch_results.csv","text/csv",use_container_width=True)
 
     with tab_info:
         st.markdown("### 📚 About BALM-PPI")
         st.markdown("""
-**BALM-PPI** predicts protein–protein and antibody–antigen binding affinity
-using ESM-2 fine-tuned with LoRA, with Integrated Gradients explainability.
+**BALM-PPI** predicts protein–protein and antibody–antigen binding affinity using ESM-2 + LoRA
+with Integrated Gradients explainability.
 
 | Component | Detail |
 |-----------|--------|
@@ -1209,8 +1095,9 @@ using ESM-2 fine-tuned with LoRA, with Integrated Gradients explainability.
 | Column mapping | `seq_a` → **Target** · `seq_b` → **proteina** |
 | Multi-chain | `\\|` separator → `<cls><cls>` double CLS token |
 | Affinity output | pKd = ((cos+1)/2) × (pKd_max − pKd_min) + pKd_min |
-| Explainability | Integrated Gradients, 30-step Riemann approximation |
-| Speed | hf_transfer · torch_dtype=auto · low_cpu_mem_usage · HF disk cache · PDB cache (1 h TTL) |
+| IG | Integrated Gradients, **15-step** Riemann (float32-safe, HF Spaces compatible) |
+| Viewer | PDB filtered to **selected chains only** |
+| Speed | hf_transfer · use_safetensors · HF disk cache · PDB cache (1 h TTL) |
 
 **Examples** · **6M0J** — ACE2 ↔ SARS-CoV-2 RBD &nbsp;·&nbsp; **1BRS** — Barnase ↔ Barstar