Update app.py

app.py

@@ -1,5 +1,5 @@
 # app.py — VeloBind HF Spaces inference app
-import os, warnings, time, base64
+import warnings, time, base64
 import numpy as np
 import pandas as pd
 import streamlit as st
@@ -18,7 +18,7 @@ MODEL_CACHE   = Path("/tmp/velobind_models")
 SEEDS         = [42, 123, 456]
 MODEL_TYPES   = ["lgbm", "cb", "xgb"]
 N_FOLDS       = 5
-VALID_AA      = set("ACDEFGHIKLMNPQRSTVWYacdefghiklmnpqrstvwyX")
+VALID_AA      = set("ACDEFGHIKLMNPQRSTVWYX")
 
 import sys
 sys.path.append(str(Path(__file__).parent))
@@ -28,6 +28,13 @@ from src.models.ensemble  import TargetScaler
 from src.config import config
 
 
+# ══════════════════════════════════════════════════════════════════════
+# Session state — theme
+# ══════════════════════════════════════════════════════════════════════
+if "dark_mode" not in st.session_state:
+    st.session_state.dark_mode = True
+
+
 # ══════════════════════════════════════════════════════════════════════
 # Validation
 # ══════════════════════════════════════════════════════════════════════
@@ -35,12 +42,8 @@ def validate_sequence(raw: str):
     raw = raw.strip()
     if not raw:
         return None, "Please enter a sequence."
-    
-    # Strip FASTA header(s)
-    lines = raw.splitlines()
-    seq_lines = [l.strip() for l in lines if not l.startswith(">")]
-    seq = "".join(seq_lines).upper().replace(" ", "")
-
+    lines   = raw.splitlines()
+    seq     = "".join(l.strip() for l in lines if not l.startswith(">")).upper().replace(" ", "")
     if len(seq) < 10:
         return None, "Sequence too short (minimum 10 residues)."
     invalid = set(seq) - VALID_AA
@@ -52,21 +55,18 @@ def validate_sequence(raw: str):
 # ══════════════════════════════════════════════════════════════════════
 # Model loading
 # ══════════════════════════════════════════════════════════════════════
-@st.cache_resource(show_spinner="Downloading and loading VeloBind models (first run ~30s)...")
+@st.cache_resource(show_spinner="Loading VeloBind models (first run ~30s)...")
 def load_all_models():
     from huggingface_hub import hf_hub_download
     MODEL_CACHE.mkdir(parents=True, exist_ok=True)
-
     model_files = (
         [f"fold_model_s{s}_{t}_f{f}.pkl"
          for s in SEEDS for t in MODEL_TYPES for f in range(N_FOLDS)]
         + ["meta_type_casf16.pkl", "target_scaler.pkl", "ligand_scaler.pkl"]
     )
-
     bar = st.progress(0, text="Loading models...")
     for i, fname in enumerate(model_files):
-        local = MODEL_CACHE / fname
-        if not local.exists():
+        if not (MODEL_CACHE / fname).exists():
             hf_hub_download(repo_id=HF_MODEL_REPO, filename=fname,
                             local_dir=str(MODEL_CACHE))
         bar.progress((i + 1) / len(model_files), text=f"Loading {fname}...")
@@ -80,14 +80,13 @@ def load_all_models():
                 joblib.load(MODEL_CACHE / f"fold_model_s{s}_{t}_f{f}.pkl")
                 for f in range(N_FOLDS)
             ]
-
     meta   = joblib.load(MODEL_CACHE / "meta_type_casf16.pkl")
     scaler = joblib.load(MODEL_CACHE / "target_scaler.pkl")
     lig_sc = joblib.load(MODEL_CACHE / "ligand_scaler.pkl")
     return fold_models, meta, scaler, lig_sc
 
 
-@st.cache_resource(show_spinner="Loading ESM-2 protein language model...")
+@st.cache_resource(show_spinner="Loading ESM-2...")
 def load_esm_model():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     tokenizer, esm_model = load_esm(config.ESM_MODEL, device)
@@ -98,8 +97,7 @@ def load_esm_model():
 def load_ad_centroid():
     for p in [Path("output/models/deployment"), Path("output/models")]:
         if (p / "ad_centroid.npy").exists():
-            return (np.load(p / "ad_centroid.npy"),
-                    float(np.load(p / "ad_threshold.npy")))
+            return np.load(p / "ad_centroid.npy"), float(np.load(p / "ad_threshold.npy"))
     for fname in ["ad_centroid.npy", "ad_threshold.npy"]:
         local = MODEL_CACHE / fname
         if not local.exists():
@@ -109,33 +107,25 @@ def load_ad_centroid():
                                 local_dir=str(MODEL_CACHE))
             except Exception:
                 return None, None
-    return (np.load(MODEL_CACHE / "ad_centroid.npy"),
-            float(np.load(MODEL_CACHE / "ad_threshold.npy")))
+    return np.load(MODEL_CACHE / "ad_centroid.npy"), float(np.load(MODEL_CACHE / "ad_threshold.npy"))
 
 
-def ad_check(esm_mean_vec, centroid, threshold):
+def ad_check(esm_vec, centroid, threshold):
     if centroid is None:
         return "UNKNOWN", float("nan")
-    dist = float(np.linalg.norm(esm_mean_vec - centroid))
+    dist = float(np.linalg.norm(esm_vec - centroid))
     return ("IN DOMAIN" if dist <= threshold else "OUT OF DOMAIN"), dist
 
 
 # ══════════════════════════════════════════════════════════════════════
-# Feature extraction
+# Features + prediction
 # ══════════════════════════════════════════════════════════════════════
 def assemble_from_parts(esm_mean, esm_var, esm_attn, seq_feat, lig_feats):
     return np.concatenate([
-        esm_mean[:, -480:],
-        seq_feat,
-        lig_feats["ecfp"],
-        lig_feats["ecfp2"],
-        lig_feats["ecfp6"],
-        lig_feats["fcfp"],
-        lig_feats["estate"],
-        lig_feats["maccs"],
-        lig_feats["atom_pair"],
-        lig_feats["torsion"],
-        lig_feats["phys"],
+        esm_mean[:, -480:], seq_feat,
+        lig_feats["ecfp"], lig_feats["ecfp2"], lig_feats["ecfp6"], lig_feats["fcfp"],
+        lig_feats["estate"], lig_feats["maccs"], lig_feats["atom_pair"],
+        lig_feats["torsion"], lig_feats["phys"],
     ], axis=1)
 
 
@@ -145,45 +135,34 @@ def extract_features(sequence, smiles_list, tokenizer, esm_model, device, lig_sc
         config.ESM_LAYERS, config.MAX_SEQ_LEN, config.HALF_SEQ_LEN,
         batch_size=1, device=device,
     )
-    seq_feat = np.array([sequence_features(sequence)])
-
+    seq_feat   = np.array([sequence_features(sequence)])
     lig_feats, valid_mask, _ = extract_ligand_features(
-        smiles_list, scaler=lig_scaler, fit_scaler=False
-    )
+        smiles_list, scaler=lig_scaler, fit_scaler=False)
     valid_mask = np.array(valid_mask)
     if valid_mask.dtype != bool:
-        bool_mask = np.zeros(len(smiles_list), dtype=bool)
-        bool_mask[valid_mask] = True
-        valid_mask = bool_mask
-
-    n_valid    = int(valid_mask.sum())
-    esm_mean_t = np.tile(esm_mean, (n_valid, 1))
-    esm_var_t  = np.tile(esm_var,  (n_valid, 1))
-    esm_attn_t = np.tile(esm_attn, (n_valid, 1))
-    seq_feat_t = np.tile(seq_feat, (n_valid, 1))
-
-    X = assemble_from_parts(esm_mean_t, esm_var_t, esm_attn_t, seq_feat_t, lig_feats)
+        bm = np.zeros(len(smiles_list), dtype=bool)
+        bm[valid_mask] = True
+        valid_mask = bm
+    n = int(valid_mask.sum())
+    X = assemble_from_parts(
+        np.tile(esm_mean, (n, 1)), np.tile(esm_var, (n, 1)),
+        np.tile(esm_attn, (n, 1)), np.tile(seq_feat, (n, 1)), lig_feats)
     return X, valid_mask, esm_mean[0]
 
 
-# ══════════════════════════════════════════════════════════════════════
-# Prediction
-# ══════════════════════════════════════════════════════════════════════
 def predict(X, fold_models, meta, scaler):
     type_avgs = []
     for s in SEEDS:
         for t in MODEL_TYPES:
-            fold_preds = np.stack([
-                scaler.inverse(fold_models[s][t][f].predict(X))
-                for f in range(N_FOLDS)
-            ], axis=1)
-            type_avgs.append(fold_preds.mean(axis=1))
-
+            fp = np.stack([scaler.inverse(fold_models[s][t][f].predict(X))
+                           for f in range(N_FOLDS)], axis=1)
+            type_avgs.append(fp.mean(axis=1))
     preds_all = np.stack(type_avgs, axis=1)
-    lgbm_avg  = preds_all[:, [0, 3, 6]].mean(axis=1)
-    cb_avg    = preds_all[:, [1, 4, 7]].mean(axis=1)
-    xgb_avg   = preds_all[:, [2, 5, 8]].mean(axis=1)
-    preds     = meta.predict(np.column_stack([lgbm_avg, cb_avg, xgb_avg]))
+    preds = meta.predict(np.column_stack([
+        preds_all[:, [0,3,6]].mean(1),
+        preds_all[:, [1,4,7]].mean(1),
+        preds_all[:, [2,5,8]].mean(1),
+    ]))
     return preds, preds_all
 
 
@@ -193,73 +172,76 @@ def uncertainty_interval(preds_all, z=1.96):
 
 
 def format_ki(pkd):
-    """Format Ki with appropriate unit (nM, uM, mM)."""
     ki_nM = 10 ** (9 - pkd)
-    if ki_nM < 1000:
-        return f"{ki_nM:.1f} nM"
-    elif ki_nM < 1_000_000:
-        return f"{ki_nM/1000:.2f} uM"
-    else:
-        return f"{ki_nM/1_000_000:.2f} mM"
+    if ki_nM < 1000:       return f"{ki_nM:.1f} nM"
+    elif ki_nM < 1_000_000: return f"{ki_nM/1000:.2f} uM"
+    else:                   return f"{ki_nM/1_000_000:.2f} mM"
 
 
 # ══════════════════════════════════���═══════════════════════════════════
-# Plots
+# Plot
 # ══════════════════════════════════════════════════════════════════════
 def bar_chart(names, preds, lo, hi, title, dark=True):
-    bg    = "#1e2a38" if dark else "#f8f9fa"
-    fg    = "#ffffff" if dark else "#111111"
-    grid  = "#2d3f55" if dark else "#cccccc"
-
-    fig, ax = plt.subplots(figsize=(max(6, len(names) * 0.9), 4),
-                           facecolor=bg)
+    bg, fg, gc = ("#1a2332", "#e8edf2", "#2d3f55") if dark else ("#f8fafc", "#1a202c", "#cbd5e0")
+    fig, ax = plt.subplots(figsize=(max(6, len(names) * 0.9), 4), facecolor=bg)
     ax.set_facecolor(bg)
     x   = np.arange(len(names))
     err = [preds - lo, hi - preds]
-    bars = ax.bar(x, preds, color="#4C72B0", alpha=0.85, width=0.6,
+    bars = ax.bar(x, preds, color="#3b82f6", alpha=0.9, width=0.6,
                   yerr=err, capsize=5, error_kw=dict(ecolor=fg, lw=1.5))
     ax.set_xticks(x)
     ax.set_xticklabels(names, rotation=30, ha='right', fontsize=10, color=fg)
     ax.set_ylabel("Predicted pKd", fontsize=11, color=fg)
     ax.set_title(title, fontsize=12, fontweight='bold', color=fg)
     ax.tick_params(colors=fg)
-    ax.spines[:].set_color(grid)
-    ax.grid(True, axis='y', alpha=0.25, color=grid)
+    for spine in ax.spines.values():
+        spine.set_color(gc)
+    ax.grid(True, axis='y', alpha=0.3, color=gc)
     for bar, val in zip(bars, preds):
         ax.text(bar.get_x() + bar.get_width() / 2,
                 bar.get_height() + 0.05, f"{val:.2f}",
-                ha='center', va='bottom', fontsize=9,
-                fontweight='bold', color=fg)
+                ha='center', va='bottom', fontsize=9, fontweight='bold', color=fg)
     plt.tight_layout()
     return fig
 
 
 # ══════════════════════════════════════════════════════════════════════
-# Page setup
+# Page layout
 # ══════════════════════════════════════════════════════════════════════
 st.set_page_config(page_title="VeloBind", layout="wide")
 
-# ── Theme toggle ──────────────────────────────────────────────────────
-with st.sidebar:
-    st.markdown("### Display")
-    dark_mode = st.toggle("Dark mode", value=True)
-
-if dark_mode:
-    header_bg    = "linear-gradient(135deg, #1a3a5c, #1e6091, #2980b9)"
-    card_bg      = "#1e2a38"
-    card_border  = "#2d3f55"
-    val_color    = "#4fc3f7"
-    lab_color    = "#aaa"
-    page_bg      = "#0e1117"
-    text_color   = "#ffffff"
+dark = st.session_state.dark_mode
+
+# ── Theme-aware CSS (only custom elements, never .stApp) ──────────────
+if dark:
+    card_bg, card_border = "#1a2332", "#2d4a6b"
+    val_col, lab_col     = "#60a5fa", "#94a3b8"
+    banner_grad          = "linear-gradient(135deg, #0f172a 0%, #1e3a5f 50%, #1e40af 100%)"
+    banner_sub           = "#93c5fd"
+    logo_bg              = "rgba(255,255,255,0.12)"
+    logo_border          = "rgba(255,255,255,0.2)"
+    toggle_bg            = "#1e3a5f"
+    toggle_knob          = "#60a5fa"
+    toggle_label         = "#93c5fd"
 else:
-    header_bg    = "linear-gradient(135deg, #2980b9, #5dade2, #85c1e9)"
-    card_bg      = "#f0f4f8"
-    card_border  = "#b0c4de"
-    val_color    = "#1a5276"
-    lab_color    = "#555"
-    page_bg      = "#ffffff"
-    text_color   = "#111111"
+    card_bg, card_border = "#f0f7ff", "#bfdbfe"
+    val_col, lab_col     = "#1d4ed8", "#475569"
+    banner_grad          = "linear-gradient(135deg, #1d4ed8 0%, #2563eb 50%, #3b82f6 100%)"
+    banner_sub           = "#dbeafe"
+    logo_bg              = "rgba(255,255,255,0.85)"
+    logo_border          = "rgba(255,255,255,0.9)"
+    toggle_bg            = "#93c5fd"
+    toggle_knob          = "#1d4ed8"
+    toggle_label         = "#dbeafe"
+
+ad_css = """
+.ad-in  { background:#064e3b; border:1px solid #059669; color:#34d399;
+          border-radius:20px; padding:0.3rem 1rem; font-weight:700; display:inline-block; font-size:0.9rem; }
+.ad-out { background:#450a0a; border:1px solid #dc2626; color:#f87171;
+          border-radius:20px; padding:0.3rem 1rem; font-weight:700; display:inline-block; font-size:0.9rem; }
+.ad-unk { background:#1e293b; border:1px solid #475569; color:#94a3b8;
+          border-radius:20px; padding:0.3rem 1rem; font-weight:700; display:inline-block; font-size:0.9rem; }
+"""
 
 def load_svg_b64(path):
     with open(path, "rb") as f:
@@ -269,93 +251,158 @@ logo_b64 = load_svg_b64("logo.svg")
 
 st.markdown(f"""
 <style>
-    .stApp {{ background-color: {page_bg}; color: {text_color}; }}
-    .header-wrap {{
-        display: flex; align-items: center; gap: 1.5rem;
-        margin-bottom: 1.5rem;
-    }}
-    .logo-box {{
-        background: #ffffff; border-radius: 12px;
-        padding: 0.75rem; flex-shrink: 0;
-    }}
-    .logo-box img {{ height: 130px; width: auto; display: block; }}
-    .header-text {{
-        background: {header_bg};
-        padding: 1.5rem 2rem; border-radius: 12px; flex: 1;
-    }}
-    .header-text h1 {{ color: #fff; font-size: 2.2rem; margin: 0; }}
-    .header-text p  {{ color: #d6eaf8; margin: 0.3rem 0 0; font-size: 1rem; }}
-    .metric-card {{
-        background: {card_bg}; border: 1px solid {card_border};
-        border-radius: 10px; padding: 1rem; text-align: center;
-    }}
-    .metric-val {{ font-size: 2rem; font-weight: 700; color: {val_color}; }}
-    .metric-lab {{ font-size: 0.8rem; color: {lab_color}; margin-top: 0.2rem; }}
-    .ad-in  {{ background:#1b4332; border:1px solid #2d6a4f; color:#40916c;
-              border-radius:8px; padding:0.4rem 1rem; font-weight:700; display:inline-block; }}
-    .ad-out {{ background:#4a1c24; border:1px solid #9b2335; color:#e74c3c;
-              border-radius:8px; padding:0.4rem 1rem; font-weight:700; display:inline-block; }}
-    .ad-unk {{ background:#2d2d2d; border:1px solid #555; color:#aaa;
-              border-radius:8px; padding:0.4rem 1rem; font-weight:700; display:inline-block; }}
+{ad_css}
+.vb-banner {{
+    background: {banner_grad};
+    border-radius: 16px;
+    padding: 1.2rem 1.8rem;
+    display: flex;
+    align-items: center;
+    gap: 1.5rem;
+    margin-bottom: 0.5rem;
+    box-shadow: 0 4px 24px rgba(0,0,0,0.18);
+    position: relative;
+}}
+.vb-logo-wrap {{
+    background: {logo_bg};
+    border: 1px solid {logo_border};
+    border-radius: 14px;
+    padding: 0.6rem;
+    backdrop-filter: blur(8px);
+    flex-shrink: 0;
+}}
+.vb-logo-wrap img {{
+    height: 110px;
+    width: auto;
+    display: block;
+}}
+.vb-title-wrap {{
+    flex: 1;
+}}
+.vb-title-wrap h1 {{
+    color: #ffffff;
+    font-size: 2.4rem;
+    font-weight: 800;
+    margin: 0 0 0.3rem 0;
+    letter-spacing: -0.5px;
+}}
+.vb-title-wrap p {{
+    color: {banner_sub};
+    font-size: 0.92rem;
+    margin: 0;
+    line-height: 1.6;
+}}
+.vb-toggle-wrap {{
+    position: absolute;
+    top: 1rem;
+    right: 1.2rem;
+    display: flex;
+    align-items: center;
+    gap: 0.5rem;
+}}
+.vb-toggle-label {{
+    color: {toggle_label};
+    font-size: 0.78rem;
+    font-weight: 600;
+    letter-spacing: 0.03em;
+}}
+.metric-card {{
+    background: {card_bg};
+    border: 1px solid {card_border};
+    border-radius: 12px;
+    padding: 1.1rem;
+    text-align: center;
+    transition: box-shadow 0.2s;
+}}
+.metric-card:hover {{
+    box-shadow: 0 4px 16px rgba(59,130,246,0.15);
+}}
+.metric-val {{
+    font-size: 1.9rem;
+    font-weight: 700;
+    color: {val_col};
+    line-height: 1.2;
+}}
+.metric-lab {{
+    font-size: 0.75rem;
+    color: {lab_col};
+    margin-top: 0.35rem;
+    line-height: 1.4;
+}}
 </style>
-<div class="header-wrap">
-    <div class="logo-box">
-        <img src="data:image/svg+xml;base64,{logo_b64}" alt="VeloBind logo"/>
+""", unsafe_allow_html=True)
+
+# ── Banner ────────────────────────────────────────────────────────────
+toggle_icon = "☀" if dark else "☾"
+toggle_text = "Light mode" if dark else "Dark mode"
+
+st.markdown(f"""
+<div class="vb-banner">
+    <div class="vb-logo-wrap">
+        <img src="data:image/svg+xml;base64,{logo_b64}" alt="VeloBind"/>
     </div>
-    <div class="header-text">
+    <div class="vb-title-wrap">
         <h1>VeloBind</h1>
-        <p>Structure-free protein-ligand binding affinity prediction &nbsp;·&nbsp;
-           Sequence + SMILES only &nbsp;·&nbsp;
-           Pearson R = 0.8469 on CASF-2016 &nbsp;·&nbsp;
-           45-model ensemble (LGBM + CatBoost + XGBoost)</p>
+        <p>
+            Structure-free protein-ligand binding affinity prediction &nbsp;&middot;&nbsp;
+            Sequence + SMILES &nbsp;&middot;&nbsp;
+            Pearson R = 0.8469 on CASF-2016 &nbsp;&middot;&nbsp;
+            45-model ensemble (LGBM + CatBoost + XGBoost)
+        </p>
     </div>
 </div>
 """, unsafe_allow_html=True)
 
-# ── Load everything ───────────────────────────────────────────────────
+# Theme toggle — just below banner, right-aligned
+_, tcol = st.columns([6, 1])
+with tcol:
+    if st.button(f"{toggle_icon}  {toggle_text}", use_container_width=True):
+        st.session_state.dark_mode = not st.session_state.dark_mode
+        st.rerun()
+
+# ── Load models ───────────────────────────────────────────────────────
 fold_models, meta, target_scaler, lig_scaler = load_all_models()
 tokenizer, esm_model, device                 = load_esm_model()
 ad_centroid, ad_threshold                    = load_ad_centroid()
 n_loaded = sum(len(fold_models[s][t]) for s in SEEDS for t in MODEL_TYPES)
-st.success(f"✓ {n_loaded} fold models loaded  |  Device: {device.upper()}")
-
-# ── Mode selector ─────────────────────────────────────────────────────
-mode = st.radio(
-    "Select mode",
-    ["Single query", "Batch screening (CSV)", "One compound vs. multiple targets"],
-    horizontal=True,
-)
-st.markdown("---")
+st.success(f"{n_loaded} fold models loaded  |  Device: {device.upper()}")
+
+# ── Mode tabs ─────────────────────────────────────────────────────────
+tab1, tab2, tab3 = st.tabs([
+    "Single query",
+    "Batch screening (CSV)",
+    "One compound vs. multiple targets",
+])
 
 
 # ══════════════════════════════════════════════════════════════════════
-# MODE 1 — Single query
+# TAB 1 — Single query
 # ══════════════════════════════════════════════════════════════════════
-if mode == "Single query":
-
+with tab1:
     col_p, col_l = st.columns(2)
     with col_p:
         st.subheader("Protein")
         seq_raw = st.text_area(
-            "Amino acid sequence (single-letter FASTA, no header)",
-            height=150,
-            placeholder="MKTAYIAKQRQISFVK...",
-            help="Only standard amino acid letters accepted (A C D E F G H I K L M N P Q R S T V W Y)."
+            "Amino acid sequence (plain or FASTA format)",
+            height=160,
+            placeholder=">ProteinName\nMKTAYIAKQRQISFVK...",
+            help="Plain sequence or FASTA with >header line. Only standard amino acid letters (A-Z subset).",
+            key="sq_seq"
         )
     with col_l:
         st.subheader("Ligand")
-        smi = st.text_input("SMILES", placeholder="CC(=O)Oc1ccccc1C(=O)O")
+        smi = st.text_input("SMILES", placeholder="CC(=O)Oc1ccccc1C(=O)O", key="sq_smi")
         examples = {
             "Aspirin":       "CC(=O)Oc1ccccc1C(=O)O",
             "Imatinib":      "Cc1ccc(NC(=O)c2ccc(CN3CCN(C)CC3)cc2)cc1Nc1nccc(-c2cccnc2)n1",
             "Gefitinib":     "COc1cc2ncnc(Nc3ccc(F)c(Cl)c3)c2cc1OCCCN1CCOCC1",
             "Staurosporine": "C[C@@H]1CCCN2C(=O)c3[nH]c4ccccc4c3C2=N1",
         }
-        chosen = st.selectbox("Load example SMILES", ["—"] + list(examples))
+        chosen = st.selectbox("Load example SMILES", ["—"] + list(examples), key="sq_ex")
         if chosen != "—":
             smi = examples[chosen]
 
-    if st.button("Predict", type="primary", use_container_width=True):
+    if st.button("Predict", type="primary", use_container_width=True, key="sq_btn"):
         seq, err = validate_sequence(seq_raw)
         if err:
             st.error(err)
@@ -366,18 +413,17 @@ if mode == "Single query":
                 t0 = time.time()
                 try:
                     X, valid, esm_vec = extract_features(
-                        seq, [smi.strip()],
-                        tokenizer, esm_model, device, lig_scaler
-                    )
+                        seq, [smi.strip()], tokenizer, esm_model, device, lig_scaler)
                     if not valid.any():
-                        st.error("RDKit could not parse this SMILES. Please check the input.")
+                        st.error("RDKit could not parse this SMILES.")
                     else:
                         preds, preds_all = predict(X, fold_models, meta, target_scaler)
                         lo, hi  = uncertainty_interval(preds_all)
                         elapsed = time.time() - t0
                         pkd     = float(preds[0])
+                        ad_label, _ = ad_check(esm_vec[-480:], ad_centroid, ad_threshold)
 
-                        st.markdown("### Results")
+                        st.markdown("#### Results")
                         c1, c2, c3, c4 = st.columns(4)
                         with c1:
                             st.markdown(f"""<div class="metric-card">
@@ -387,25 +433,25 @@ if mode == "Single query":
                         with c2:
                             st.markdown(f"""<div class="metric-card">
                                 <div class="metric-val">[{lo[0]:.2f}, {hi[0]:.2f}]</div>
-                                <div class="metric-lab">95% model interval (±1.96σ, 45 models)</div>
+                                <div class="metric-lab">95% model interval<br>(±1.96σ · 45 models)</div>
                             </div>""", unsafe_allow_html=True)
                         with c3:
                             st.markdown(f"""<div class="metric-card">
                                 <div class="metric-val">{format_ki(pkd)}</div>
-                                <div class="metric-lab">Estimated Ki (pKd ≈ pKi assumed)</div>
+                                <div class="metric-lab">Estimated Ki<br>(pKd ≈ pKi assumed)</div>
                             </div>""", unsafe_allow_html=True)
                         with c4:
-                            ad_label, _ = ad_check(esm_vec[-480:], ad_centroid, ad_threshold)
                             ad_cls = ("ad-in"  if ad_label == "IN DOMAIN" else
                                       "ad-out" if ad_label == "OUT OF DOMAIN" else "ad-unk")
                             st.markdown(f"""<div class="metric-card">
-                                <div class="{ad_cls}">{ad_label}</div>
-                                <div class="metric-lab">Applicability domain</div>
+                                <div style="padding-top:0.4rem">
+                                    <span class="{ad_cls}">{ad_label}</span>
+                                </div>
+                                <div class="metric-lab" style="margin-top:0.6rem">Applicability domain</div>
                             </div>""", unsafe_allow_html=True)
 
                         if ad_label == "OUT OF DOMAIN":
-                            st.warning("Protein is outside the training distribution. "
-                                       "Predictions may be unreliable.")
+                            st.warning("Protein is outside the training distribution. Predictions may be unreliable.")
 
                         st.caption(
                             f"Inference time: {elapsed:.2f}s  |  "
@@ -419,8 +465,8 @@ if mode == "Single query":
                                 labels, preds_all[0],
                                 preds_all[0] - preds_all[0].std(),
                                 preds_all[0] + preds_all[0].std(),
-                                "Seed x type predictions (fold-averaged)",
-                                dark=dark_mode,
+                                "Per-seed and model-type predictions (fold-averaged)",
+                                dark=dark,
                             )
                             st.pyplot(fig, use_container_width=True)
                             plt.close(fig)
@@ -431,26 +477,26 @@ if mode == "Single query":
 
 
 # ══════════════════════════════════════════════════════════════════════
-# MODE 2 — Batch screening
+# TAB 2 — Batch screening
 # ══════════════════════════════════════════════════════════════════════
-elif mode == "Batch screening (CSV)":
-
+with tab2:
     st.subheader("Batch Screening")
-    st.markdown("One protein, many compounds. Upload a CSV with a `smiles` column "
-                "(and optionally `name`). Results are ranked by predicted pKd.")
+    st.markdown("Screen a library of compounds against one target. "
+                "Upload a CSV with a `smiles` column (and optionally `name`). "
+                "Results are ranked by predicted pKd.")
 
     col_seq, col_csv = st.columns(2)
     with col_seq:
-        batch_seq_raw = st.text_area("Target protein sequence", height=180,
-                                     placeholder="Paste UniProt sequence...")
+        batch_seq_raw = st.text_area("Target protein sequence (plain or FASTA)", height=180,
+                                     placeholder=">Target\nMKTAYIAKQRQISFVK...", key="bs_seq")
     with col_csv:
-        uploaded = st.file_uploader("Compound CSV (smiles, name)", type=["csv"])
+        uploaded = st.file_uploader("Compound CSV (smiles, name)", type=["csv"], key="bs_up")
         st.code("smiles,name\nCC(=O)Oc1ccccc1C(=O)O,Aspirin", language="csv")
 
-    max_cpds = st.slider("Max compounds", 10, 500, 100,
+    max_cpds = st.slider("Max compounds", 10, 500, 100, key="bs_max",
                          help="~1s per compound on CPU free tier.")
 
-    if st.button("Run batch screening", type="primary", use_container_width=True):
+    if st.button("Run batch screening", type="primary", use_container_width=True, key="bs_btn"):
         batch_seq, err = validate_sequence(batch_seq_raw)
         if err:
             st.error(err)
@@ -470,9 +516,7 @@ elif mode == "Batch screening (CSV)":
             with st.spinner(f"Screening {len(smiles_list)} compounds..."):
                 t0 = time.time()
                 X, valid, esm_vec = extract_features(
-                    batch_seq, smiles_list,
-                    tokenizer, esm_model, device, lig_scaler
-                )
+                    batch_seq, smiles_list, tokenizer, esm_model, device, lig_scaler)
                 ad_label, _ = ad_check(esm_vec[-480:], ad_centroid, ad_threshold)
                 preds, preds_all = predict(X, fold_models, meta, target_scaler)
                 lo, hi  = uncertainty_interval(preds_all)
@@ -482,6 +526,9 @@ elif mode == "Batch screening (CSV)":
             valid_smiles = [smiles_list[i] for i in range(len(smiles_list)) if valid[i]]
             n_invalid    = int((~valid).sum())
 
+            if ad_label == "OUT OF DOMAIN":
+                st.warning("Protein is outside the training distribution. Predictions may be unreliable.")
+
             results_df = pd.DataFrame({
                 'name':      valid_names,
                 'smiles':    valid_smiles,
@@ -494,25 +541,19 @@ elif mode == "Batch screening (CSV)":
             }).sort_values('pKd_pred', ascending=False).reset_index(drop=True)
             results_df.insert(0, 'rank', range(1, len(results_df) + 1))
 
-            if ad_label == "OUT OF DOMAIN":
-                st.warning("Protein is outside the training distribution. "
-                           "Predictions may be unreliable.")
-
             st.success(
-                f"✓ {len(results_df)} compounds in {elapsed:.1f}s "
+                f"{len(results_df)} compounds screened in {elapsed:.1f}s "
                 f"({elapsed / max(len(results_df), 1):.2f}s/compound)"
                 + (f"  |  {n_invalid} invalid SMILES skipped" if n_invalid else "")
             )
 
             top_n  = min(20, len(results_df))
-            top_df = results_df.head(top_n)
-            fig    = bar_chart(
-                top_df['name'].tolist(),
-                top_df['pKd_pred'].values,
-                top_df['CI_lo'].values,
-                top_df['CI_hi'].values,
-                f"Top {top_n} hits",
-                dark=dark_mode,
+            fig = bar_chart(
+                results_df.head(top_n)['name'].tolist(),
+                results_df.head(top_n)['pKd_pred'].values,
+                results_df.head(top_n)['CI_lo'].values,
+                results_df.head(top_n)['CI_hi'].values,
+                f"Top {top_n} hits by predicted pKd", dark=dark,
             )
             st.pyplot(fig, use_container_width=True)
             plt.close(fig)
@@ -524,52 +565,46 @@ elif mode == "Batch screening (CSV)":
             st.download_button(
                 "Download ranked CSV",
                 results_df.to_csv(index=False).encode(),
-                file_name="velobind_screening.csv",
-                mime="text/csv",
+                file_name="velobind_screening.csv", mime="text/csv",
             )
 
 
 # ══════════════════════════════════════════════════════════════════════
-# MODE 3 — One compound vs. multiple targets
+# TAB 3 — Selectivity profiling
 # ══════════════════════════════════════════════════════════════════════
-elif mode == "One compound vs. multiple targets":
-
+with tab3:
     st.subheader("Selectivity Profiling")
-    st.markdown("One SMILES, multiple proteins — ranked by predicted pKd. "
-                "Format: `TargetName: SEQUENCE` (name optional).")
+    st.markdown("One compound, multiple targets — ranked by predicted pKd. "
+                "Format: `TargetName: SEQUENCE` (name optional). "
+                "Accepts plain sequence or FASTA per line.")
 
-    multi_smi  = st.text_input("Compound SMILES",
-                               placeholder="Cc1ccc(NC(=O)...)cc1Nc1nccc(...)n1")
+    multi_smi  = st.text_input("Compound SMILES", placeholder="Cc1ccc(...)cc1Nc1nccc(...)n1", key="sp_smi")
     multi_seqs = st.text_area(
         "Target proteins (one per line)",
         height=250,
         placeholder="ABL1: MGPSENDPNLFVALY...\nEGFR: MRPSGTAGAALLALL...\nCDK2: MENFQKVEKIGEGTY...",
+        key="sp_seqs",
     )
 
-    if st.button("Run selectivity profiling", type="primary", use_container_width=True):
+    if st.button("Run selectivity profiling", type="primary", use_container_width=True, key="sp_btn"):
         if not multi_smi.strip() or not multi_seqs.strip():
             st.error("Please enter a SMILES and at least one protein sequence.")
         else:
-            targets = {}
-            parse_errors = []
+            targets, parse_errors = {}, []
             for i, line in enumerate(multi_seqs.strip().splitlines()):
                 line = line.strip()
                 if not line:
                     continue
-                if ":" in line:
-                    name, raw_seq = line.split(":", 1)
-                    name = name.strip()
-                else:
-                    name, raw_seq = f"Target_{i+1}", line
-                seq, err = validate_sequence(raw_seq)
+                name, raw_seq = (line.split(":", 1) if ":" in line
+                                 else (f"Target_{i+1}", line))
+                seq, err = validate_sequence(raw_seq if isinstance(raw_seq, str) else raw_seq)
                 if err:
-                    parse_errors.append(f"{name}: {err}")
+                    parse_errors.append(f"{name.strip()}: {err}")
                 else:
-                    targets[name] = seq
+                    targets[name.strip()] = seq
 
-            if parse_errors:
-                for e in parse_errors:
-                    st.warning(f"Skipped — {e}")
+            for e in parse_errors:
+                st.warning(f"Skipped — {e}")
             if not targets:
                 st.error("No valid sequences found.")
                 st.stop()
@@ -578,18 +613,16 @@ elif mode == "One compound vs. multiple targets":
             for idx, (name, seq) in enumerate(targets.items()):
                 try:
                     X, valid, esm_vec = extract_features(
-                        seq, [multi_smi.strip()],
-                        tokenizer, esm_model, device, lig_scaler
-                    )
+                        seq, [multi_smi.strip()], tokenizer, esm_model, device, lig_scaler)
                     if valid.any():
                         preds, preds_all = predict(X, fold_models, meta, target_scaler)
                         lo, hi = uncertainty_interval(preds_all)
                         ad_label, _ = ad_check(esm_vec[-480:], ad_centroid, ad_threshold)
                         results.append({
                             'Target':    name,
-                            'pKd_pred':  round(float(preds[0]),  3),
-                            'CI_lo':     round(float(lo[0]),     3),
-                            'CI_hi':     round(float(hi[0]),     3),
+                            'pKd_pred':  round(float(preds[0]), 3),
+                            'CI_lo':     round(float(lo[0]), 3),
+                            'CI_hi':     round(float(hi[0]), 3),
                             'Ki_est':    format_ki(float(preds[0])),
                             'model_std': round(float(preds_all.std()), 3),
                             'AD':        ad_label,
@@ -599,21 +632,16 @@ elif mode == "One compound vs. multiple targets":
                 progress.progress((idx + 1) / len(targets))
 
             progress.empty()
-            res_df = (
-                pd.DataFrame(results)
-                .sort_values('pKd_pred', ascending=False)
-                .reset_index(drop=True)
-            )
+            res_df = (pd.DataFrame(results)
+                      .sort_values('pKd_pred', ascending=False)
+                      .reset_index(drop=True))
             res_df.insert(0, 'rank', range(1, len(res_df) + 1))
 
             st.success(f"Profiled {len(res_df)} targets.")
             fig = bar_chart(
-                res_df['Target'].tolist(),
-                res_df['pKd_pred'].values,
-                res_df['CI_lo'].values,
-                res_df['CI_hi'].values,
-                "Selectivity profile — predicted pKd by target",
-                dark=dark_mode,
+                res_df['Target'].tolist(), res_df['pKd_pred'].values,
+                res_df['CI_lo'].values, res_df['CI_hi'].values,
+                "Selectivity profile — predicted pKd by target", dark=dark,
             )
             st.pyplot(fig, use_container_width=True)
             plt.close(fig)
@@ -622,19 +650,18 @@ elif mode == "One compound vs. multiple targets":
             st.download_button(
                 "Download selectivity CSV",
                 res_df.to_csv(index=False).encode(),
-                file_name="velobind_selectivity.csv",
-                mime="text/csv",
+                file_name="velobind_selectivity.csv", mime="text/csv",
             )
 
 
 # ── Footer ────────────────────────────────────────────────────────────
 st.markdown("---")
 st.markdown(f"""
-<div style="color:{lab_color};font-size:0.8rem;text-align:center;padding:0.5rem">
-    VeloBind &nbsp;·&nbsp; Structure-free binding affinity &nbsp;·&nbsp;
-    ESM-2 + GBM ensemble &nbsp;·&nbsp;
-    Trained on LP-PDBBind &nbsp;·&nbsp;
-    Evaluated on CASF-2016/2013 &nbsp;·&nbsp;
-    <b>Not for clinical use.</b>
+<div style="color:{lab_col};font-size:0.78rem;text-align:center;padding:0.4rem 0 0.8rem">
+    VeloBind &nbsp;&middot;&nbsp; Structure-free binding affinity &nbsp;&middot;&nbsp;
+    ESM-2 + gradient-boosted ensemble &nbsp;&middot;&nbsp;
+    Trained on LP-PDBBind &nbsp;&middot;&nbsp;
+    Evaluated on CASF-2016 and CASF-2013 &nbsp;&middot;&nbsp;
+    <b>Not for clinical use</b>
 </div>
 """, unsafe_allow_html=True)