app.py

"""
LLPSense Gradio Demo
Condition-dependent protein LLPS prediction using ProtT5 + XGBoost

Pipeline:
  (1) User inputs amino-acid sequence
  (2) Click "Extract Feature" → mean-pool ProtT5-XL embedding (1024-dim)
  (3) Tab 1 – Predict LLPS Probability: adjust temp / conc / pH sliders → predict
  (4) Tab 2 – Condition Screening: pick one condition to vary, fix the rest → plot
"""

import sys
import warnings
from pathlib import Path
from copy import deepcopy

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import joblib
import gradio as gr
import spaces
from huggingface_hub import snapshot_download

from examples import EXAMPLES, feature_path, find_example_by_seq

# Only surface examples that actually have a pre-computed T5 feature cached
# in assets/ — an entry added via preprocess/add_example.py but not yet run
# through preprocess/extract_example_feat.py would otherwise show up in the
# picker and silently fall back to a slow/GPU extraction on first click.
AVAILABLE_EXAMPLES = [example for example in EXAMPLES if feature_path(example["id"]).exists()]
from t5_utils import (
    T5_REPO_ID,
    extract_t5_feature as _extract_t5_feature_core,
    preload_t5_cpu,
    read_feature_h5,
)

warnings.filterwarnings("ignore")

# ── Paths ─────────────────────────────────────────────────────────────────────
BASE_DIR = Path(__file__).parent

# ── ProtT5 prefetch ────────────────────────────────────────────────────────────
# This SDK-gradio Space builds from requirements.txt, not the repo's Dockerfile,
# so there is no build-time prefetch step. Download the weights here instead,
# at module import (app startup, plain CPU context) — before any @spaces.GPU
# call — so extract_t5_feature() never blocks on a network download while
# holding a ZeroGPU allocation (which has a short time budget).
snapshot_download(T5_REPO_ID)

# Also deserialize the weights into CPU memory here, still outside any
# @spaces.GPU context. Without this, T5EncoderModel.from_pretrained() (loading
# ~2.9GB from disk) would run lazily inside the first GPU-decorated call and
# could eat enough of the ZeroGPU time budget to abort the task. With this,
# the first GPU call only needs a fast .to("cuda") transfer.
preload_t5_cpu()

# ── Physical constants (from preprocess/misc.py) ──────────────────────────────
MAX_TEMP   = 60.0
MAX_CONC   = 1000.0
MAX_PH     = 14.0
MAX_MGCL2  = 50.0
MAX_NACL   = 2000.0
MAX_KCL    = 1000.0
MAX_CAGENT = 50.0
MAX_GLYC   = 10.0

# ── Screening ranges ─────────────────────────────────────────────────────────
SCREEN_RANGES = {
    "Temperature": np.arange(0.0, 61.0, 1.0),        # 0–60 °C
    "Concentration": np.arange(0.0, 1010.0, 10.0),   # 0–1000 µM
    "pH": np.arange(4.0, 12.1, 0.1),                 # 4.0–12.0
}
SCREEN_KEYS = {
    "Temperature": "temp",
    "Concentration": "conc",
    "pH": "pH",
}
SCREEN_LABELS = {
    "Temperature": "Temperature (°C)",
    "Concentration": "Concentration (µM)",
    "pH": "pH",
}

VALID_AA = set("ACDEFGHIKLMNPQRSTVWY")
ALLOW_AA = VALID_AA | set("XBJOUZ")

# ── Lazy-loaded singletons ────────────────────────────────────────────────────
_llps_model = None


def load_llps_model():
    global _llps_model
    if _llps_model is None:
        model_path = BASE_DIR / "models" / "LLPSense.pkl"
        if not model_path.exists():
            raise FileNotFoundError(
                f"Model file not found: {model_path}\n"
                "Please place 'LLPSense.pkl' inside the 'models/' directory."
            )
        d = joblib.load(model_path)
        mdl = d["model"]
        # Force XGBoost to run on CPU to avoid device-mismatch warnings
        mdl.set_params(device="cpu")
        _llps_model = mdl
    return _llps_model


# ── Feature extraction ────────────────────────────────────────────────────────
# The actual tokenize/forward/mean-pool logic lives in t5_utils.py, shared
# with preprocess/extract_example_feat.py so the offline-cached example
# features in assets/ always match what this would compute live. Only the
# @spaces.GPU wrapping (ZeroGPU allocation) is app-specific, and cb_extract()
# below skips calling this entirely when a cached feature is available.
@spaces.GPU(duration=180)
def extract_t5_feature(sequence: str) -> np.ndarray:
    return _extract_t5_feature_core(sequence)


# ── Condition vector builder ──────────────────────────────────────────────────
def build_cond(temp, conc, pH,
               nacl=160.0, mgcl2=0.0, kcl=0.0, glyc=0.0,
               peg1=0.0, peg2=0.0, peg3=0.0,
               ficoll=0.0, dext40=0.0, dext70=0.0) -> np.ndarray:
    """
    Normalise environmental parameters and return a 13-dim condition vector.
    Order: [temp, conc, pH, PEG300-1k, PEG3k-6k, PEG8k-20k,
            Ficoll, Dextran≤40, Dextran≥70, MgCl2, NaCl, KCl, Glycerol]
    """
    return np.array([
        temp   / MAX_TEMP,
        conc   / MAX_CONC,
        pH     / MAX_PH,
        peg1   / MAX_CAGENT,
        peg2   / MAX_CAGENT,
        peg3   / MAX_CAGENT,
        ficoll / MAX_CAGENT,
        dext40 / MAX_CAGENT,
        dext70 / MAX_CAGENT,
        mgcl2  / MAX_MGCL2,
        nacl   / MAX_NACL,
        kcl    / MAX_KCL,
        glyc   / MAX_GLYC,
    ], dtype=np.float32)


def model_predict(feat: np.ndarray, cond: np.ndarray) -> float:
    model = load_llps_model()
    x = np.concatenate([feat, cond]).reshape(1, -1)
    return float(model.predict_proba(x)[0, 1])


# ── Smoothing (same as LLPSXG.py's moving_average) ────────────────────────────
def moving_average(y: np.ndarray, window_size: int) -> np.ndarray:
    if window_size % 2 == 0:
        raise ValueError("Window size should be odd to ensure symmetry.")
    window = np.ones(int(window_size)) / float(window_size)
    y_padded = np.pad(y, (window_size // 2, window_size // 2), mode="edge")
    return np.convolve(y_padded, window, "valid")


# ── Matplotlib helpers ────────────────────────────────────────────────────────
def prob_gauge_figure(prob: float) -> plt.Figure:
    """Horizontal probability bar gauge."""
    LLPS_COLOR  = "#e74c3c"
    NON_COLOR   = "#2980b9"
    color = LLPS_COLOR if prob >= 0.5 else NON_COLOR
    label = "Phase Separating" if prob >= 0.5 else "Non-Phase Separating"

    fig, ax = plt.subplots(figsize=(7, 2.8))
    ax.barh([0], [prob],     height=0.55, color=color,    alpha=0.88, zorder=3)
    ax.barh([0], [1 - prob], height=0.55, left=prob,
            color="#ecf0f1", alpha=0.9, zorder=2)
    ax.axvline(0.5, color="#2c3e50", lw=1.8, ls="--", label="Threshold 0.5", zorder=4)
    ax.set_xlim(0, 1)
    ax.set_ylim(-0.55, 0.55)
    ax.set_yticks([])
    ax.set_xlabel("LLPS Probability", fontsize=12)
    ax.set_title(f"{label}   |   Probability: {prob:.4f}",
                 fontsize=14, fontweight="bold", color=color, pad=10)
    ax.legend(fontsize=10, loc="lower right")
    ax.spines[["top", "right", "left"]].set_visible(False)
    plt.tight_layout()
    return fig


def screening_figure(xvals: np.ndarray, probs: np.ndarray,
                     screen_name: str, probs_smooth: np.ndarray = None) -> plt.Figure:
    """Line graph for condition screening result.

    If probs_smooth differs from probs (smoothing window > 1), the raw curve
    is drawn as a faint dotted reference line and the smoothed curve becomes
    the main plotted/filled line.
    """
    xlabel = SCREEN_LABELS[screen_name]
    LLPS_COLOR = "#e74c3c"
    NON_COLOR  = "#2980b9"

    smoothed = probs_smooth is not None and not np.array_equal(probs, probs_smooth)
    plot_probs = probs_smooth if smoothed else probs

    fig, ax = plt.subplots(figsize=(9, 5))
    if smoothed:
        ax.plot(xvals, probs, lw=1.2, color=NON_COLOR, alpha=0.35, ls=":",
                label="Raw", zorder=2)
    ax.plot(xvals, plot_probs, lw=2.5, color=NON_COLOR,
            label="Smoothed LLPS Probability" if smoothed else "LLPS Probability",
            zorder=3)
    ax.axhline(0.5, color=LLPS_COLOR, lw=1.8, ls="--",
               label="Threshold 0.5", zorder=4)
    ax.fill_between(xvals, plot_probs, 0.5,
                    where=(plot_probs >= 0.5), alpha=0.22,
                    color=LLPS_COLOR, label="LLPS region", zorder=2)
    ax.fill_between(xvals, plot_probs, 0.5,
                    where=(plot_probs < 0.5),  alpha=0.15,
                    color=NON_COLOR, label="Non-LLPS region", zorder=2)
    ax.set_xlim(xvals[0], xvals[-1])
    ax.set_ylim(0, 1)
    ax.set_xlabel(xlabel, fontsize=13)
    ax.set_ylabel("LLPS Probability", fontsize=13)
    ax.set_title(f"Condition Screening — {xlabel}", fontsize=14, fontweight="bold")
    ax.legend(fontsize=11, loc="upper right")
    ax.grid(True, alpha=0.3)
    ax.spines[["top", "right"]].set_visible(False)
    plt.tight_layout()
    return fig


# ── Status HTML templates ─────────────────────────────────────────────────────
_SPINNER_HTML = """
<div style="display:flex;align-items:center;gap:14px;padding:10px 4px;">
  <div style="
    width:28px;height:28px;flex-shrink:0;
    border:3px solid #fecaca;
    border-top-color:#ef4444;
    border-radius:50%;
    animation:llps-spin 0.85s linear infinite;
  "></div>
  <span style="color:#555;font-size:14px;font-weight:500;line-height:1.5;">
    ProtT5-XL feature extraction in progress…<br>
    <span style="font-size:12px;color:#999;font-weight:400;">
      First run loads the model onto the GPU — this may take a moment.
    </span>
  </span>
</div>
<style>@keyframes llps-spin{to{transform:rotate(360deg)}}</style>
"""

def _status_ok(seq_len: int, feat_dim: int) -> str:
    pill = ("background:#f0fdf4;color:#15803d;border:1px solid #bbf7d0;"
            "border-radius:9999px;padding:2px 10px;font-size:12px;font-weight:600;"
            "display:inline-block;margin:0 4px 0 0;")
    return (f'<div style="color:#16a34a;font-weight:600;padding:6px 0;display:flex;align-items:center;gap:6px;">'
            f'<span>✅ Feature extracted</span>'
            f'<span style="{pill}">Length: {seq_len} AA</span>'
            f'</div>')

def _status_warn(msg: str) -> str:
    return f'<div style="color:#d97706;padding:6px 0;">⚠️ {msg}</div>'

def _status_err(msg: str) -> str:
    return f'<div style="color:#dc2626;padding:6px 0;">❌ {msg}</div>'


# ── Gradio callback functions ─────────────────────────────────────────────────
def cb_extract(sequence: str):
    """Step 2: Extract ProtT5 feature from sequence (generator → streams status).

    Every yield also clears feat_state and the Step-3 result panels so a
    stale feature/result from a previously extracted sequence can never
    remain visible or be used once a new extraction starts.
    """
    seq = sequence.strip().upper()
    if not seq:
        yield None, _status_warn("Please enter a protein sequence."), None, "", None, ""
        return

    invalid = set(seq) - ALLOW_AA
    if invalid:
        yield None, _status_warn(f"Invalid characters: <code>{''.join(sorted(invalid))}</code>"), None, "", None, ""
        return

    # ── invalidate old feature/results immediately, then show spinner ─────────
    yield None, _SPINNER_HTML, None, "", None, ""

    try:
        # Known example sequence with a pre-computed feature? Skip the T5
        # model/GPU call entirely and load it straight from assets/.
        example = find_example_by_seq(seq)
        cached_path = feature_path(example["id"]) if example else None
        if cached_path and cached_path.exists():
            feat = read_feature_h5(cached_path)
        else:
            feat = extract_t5_feature(seq)
        yield feat, _status_ok(len(seq), feat.shape[0]), None, "", None, ""
    except Exception as e:
        yield None, _status_err(str(e)), None, "", None, ""


def cb_predict(feat,
               temp, conc, pH,
               nacl, mgcl2, kcl, glyc,
               peg1, peg2, peg3, ficoll, dext40, dext70):
    """Tab 1: Predict LLPS probability for a single condition point."""
    if feat is None:
        return None, "⚠️ Please extract the T5 feature first (Step 2)."
    try:
        cond = build_cond(temp, conc, pH, nacl, mgcl2, kcl, glyc,
                          peg1, peg2, peg3, ficoll, dext40, dext70)
        prob = model_predict(feat, cond)
        fig  = prob_gauge_figure(prob)
        label = "**Phase Separating** 🔴" if prob >= 0.5 else "**Non-Phase Separating** 🔵"
        txt   = f"{label}  \nLLPS Probability: **{prob:.4f}**"
        return fig, txt
    except Exception as e:
        return None, f"❌ Prediction failed: {e}"


def cb_screen(feat, screen_name,
              fix_temp, fix_conc, fix_pH,
              nacl, mgcl2, kcl, glyc,
              peg1, peg2, peg3, ficoll, dext40, dext70,
              smooth_window):
    """Tab 2: Screen LLPS across a range of one condition."""
    if feat is None:
        return None, "⚠️ Please extract the T5 feature first (Step 2)."
    if screen_name is None:
        return None, "⚠️ Please select a condition to screen."
    try:
        xvals = SCREEN_RANGES[screen_name]
        probs = []
        for v in xvals:
            t = float(v) if screen_name == "Temperature"   else fix_temp
            c = float(v) if screen_name == "Concentration" else fix_conc
            p = float(v) if screen_name == "pH"            else fix_pH
            cond = build_cond(t, c, p, nacl, mgcl2, kcl, glyc,
                              peg1, peg2, peg3, ficoll, dext40, dext70)
            probs.append(model_predict(feat, cond))

        probs = np.array(probs)

        # Slider step keeps this odd (1, 3, 5, ...); window=1 is a no-op average.
        window = int(smooth_window)
        probs_plot = moving_average(probs, window) if window > 1 else probs

        fig = screening_figure(xvals, probs, screen_name, probs_plot)

        peak_idx = probs_plot.argmax()
        xlabel   = SCREEN_LABELS[screen_name]
        txt = (f"Screening completed.  \n"
               f"Peak probability **{probs_plot[peak_idx]:.4f}** "
               f"at {xlabel} = **{xvals[peak_idx]:.1f}**  \n"
               f"LLPS-positive range: "
               f"**{(probs_plot >= 0.5).sum()}** / {len(probs_plot)} points ≥ 0.5")
        if window > 1:
            txt += f"  \n*(Smoothed with moving-average window size {window})*"
        return fig, txt
    except Exception as e:
        return None, f"❌ Screening failed: {e}"




# ── Gradio UI ─────────────────────────────────────────────────────────────────
DESCRIPTION = """
# 🧬 LLPSense — Condition-Dependent LLPS Prediction

**LLPSense** predicts whether a protein undergoes liquid-liquid phase separation (LLPS)
under user-defined environmental conditions.

> Bae†, Kang†, et al. *"A machine learning framework for predicting and modulating
> condition-dependent protein phase separation."* bioRxiv 2025.
---
## Workflow
1. **Paste your sequence** in the text box below.
2. Click **Extract Feature** to compute the ProtT5-XL embedding.
3. Use **Predict LLPS Probability** to query a specific condition point, or
   **Condition Screening** to sweep one condition across its full range.
---
"""

CUSTOM_CSS = """
/* Import fonts from Google Fonts */
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600&family=JetBrains+Mono:wght@400;500&display=swap');

/* Sequence textbox: monospace font for clean AA letter display */
#seq-input textarea,
#seq-input input {
    font-family: 'JetBrains Mono', 'Source Code Pro', 'Courier New', monospace !important;
    font-size: 14px !important;
    line-height: 1.7 !important;
    letter-spacing: 0.04em !important;
}

/* Slider thumb: orange */
input[type="range"]::-webkit-slider-thumb {
    background: #ef4444 !important;
    border-color: #ef4444 !important;
}
input[type="range"]::-moz-range-thumb {
    background: #ef4444 !important;
    border-color: #ef4444 !important;
}

/* Slider numeric value input */
input[type="number"] {
    font-size: 15px !important;
    font-weight: 600 !important;
}

/* Initial state: hide Temperature slider column (default screen condition) */
#s-temp-col { display: none; }

/* Tab buttons */
button[role="tab"] {
    background: #ffffff !important;
    border: 1.5px solid #d1d5db !important;
    border-radius: 8px 8px 0 0 !important;
    color: #4b5563 !important;
    font-weight: 500 !important;
    transition: background 0.15s, color 0.15s !important;
}
button[role="tab"]:hover {
    background: #eff6ff !important;
    color: #1d4ed8 !important;
    border-color: #93c5fd !important;
}
button[role="tab"][aria-selected="true"] {
    background: #2563eb !important;
    color: #ffffff !important;
    border-color: #2563eb !important;
    font-weight: 600 !important;
}
"""

with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="blue",
        font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
        font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
    ),
    css=CUSTOM_CSS,
    title="LLPSense Demo",
) as demo:

    feat_state = gr.State(None)

    gr.Markdown(DESCRIPTION)

    # ────────────────────────────────────────────────────────────────
    # Step 1 + 2: Sequence input & feature extraction
    # ────────────────────────────────────────────────────────────────
    with gr.Group():
        gr.Markdown("## Step 1 — Enter Protein Sequence")
        seq_box = gr.Textbox(
            label="Amino Acid Sequence (1-letter code)",
            placeholder="Paste your protein sequence here (e.g. MDVFMKGLSK…)",
            lines=5,
            value=AVAILABLE_EXAMPLES[0]["seq"] if AVAILABLE_EXAMPLES else "",
            elem_id="seq-input",
        )
        # Clicking an example fills seq_box; the matching Step-2 extraction
        # is chained onto it further down (once the Step-3 output components
        # exist) so selecting an example runs extraction automatically.
        # Skipped entirely if no example has a cached feature yet.
        example_picker = None
        if AVAILABLE_EXAMPLES:
            with gr.Accordion("🧪 Examples  (We provide preprocessed T5 feature)", open=False):
                example_picker = gr.Examples(
                    examples=[[example["seq"]] for example in AVAILABLE_EXAMPLES],
                    example_labels=[example["name"] for example in AVAILABLE_EXAMPLES],
                    inputs=[seq_box],
                )

    with gr.Group():
        gr.Markdown("## Step 2 — Extract ProtT5 Feature")
        gr.Markdown(
            "Runs the [**ProtT5-XL**](https://github.com/agemagician/ProtTrans) encoder to produce a 1024-dim mean-pool embedding.  \n"
            "⏳ *First call loads the model onto the GPU and may take a moment.*"
        )
        extract_btn    = gr.Button("🔬 Extract Feature", variant="primary", size="lg")
        extract_status = gr.HTML("")

    # ────────────────────────────────────────────────────────────────
    # Step 3: Prediction & Screening
    # ────────────────────────────────────────────────────────────────
    with gr.Group():
        gr.Markdown("## Step 3 — Run Demo")
        gr.Markdown(
            "Predict LLPS probability for a specific condition, "
            "or sweep one condition across its full range."
        )

        with gr.Tabs():

            # ── Tab 1: Single-point prediction ───────────────────────────
            with gr.Tab("🔮 Predict LLPS Probability"):
                gr.Markdown(
                    "Set the environmental conditions with the sliders below, "
                    "then click **Predict** to obtain the LLPS probability."
                )

                # Primary conditions
                with gr.Row():
                    p_temp = gr.Slider(0, 60,   value=25.0,  step=0.5,
                                       label="Temperature (°C)")
                    p_conc = gr.Slider(0, 1000, value=100.0, step=5.0,
                                       label="Concentration (µM)")
                    p_pH   = gr.Slider(0, 14,   value=7.3,   step=0.1,
                                       label="pH")

                # Advanced conditions
                with gr.Accordion("⚙️ Advanced Conditions  (Salts & Crowding Agents)", open=False):
                    gr.Markdown("Default values represent a common physiological buffer (160 mM NaCl).")
                    with gr.Row():
                        p_nacl  = gr.Slider(0, 2000, value=160.0, step=10.0, label="NaCl (mM)")
                        p_mgcl2 = gr.Slider(0, 50,   value=0.0,   step=1.0,  label="MgCl₂ (mM)")
                        p_kcl   = gr.Slider(0, 1000, value=0.0,   step=10.0, label="KCl (mM)")
                        p_glyc  = gr.Slider(0, 10,   value=0.0,   step=0.5,  label="Glycerol (%)")
                    with gr.Row():
                        p_peg1   = gr.Slider(0, 50, value=0, step=1, label="PEG 300–1000 (%)")
                        p_peg2   = gr.Slider(0, 50, value=0, step=1, label="PEG 3k–6k (%)")
                        p_peg3   = gr.Slider(0, 50, value=0, step=1, label="PEG 8k–20k (%)")
                    with gr.Row():
                        p_ficoll = gr.Slider(0, 50, value=0, step=1, label="Ficoll (%)")
                        p_dext40 = gr.Slider(0, 50, value=0, step=1, label="Dextran ≤40 kDa (%)")
                        p_dext70 = gr.Slider(0, 50, value=0, step=1, label="Dextran ≥70 kDa (%)")

                pred_btn  = gr.Button("⚡ Predict LLPS Probability", variant="primary")
                pred_plot = gr.Plot(label="Prediction Result")
                pred_text = gr.Markdown("")

                pred_btn.click(
                    fn=cb_predict,
                    inputs=[
                        feat_state,
                        p_temp, p_conc, p_pH,
                        p_nacl, p_mgcl2, p_kcl, p_glyc,
                        p_peg1, p_peg2, p_peg3, p_ficoll, p_dext40, p_dext70,
                    ],
                    outputs=[pred_plot, pred_text],
                )

            # ── Tab 2: Condition Screening ────────────────────────────────
            with gr.Tab("📈 Condition Screening"):
                gr.Markdown(
                    "Select **one condition** to screen across its full physiological range.  \n"
                    "The remaining conditions are held fixed at the values you specify below."
                )

                screen_radio = gr.Radio(
                    choices=["Temperature", "Concentration", "pH"],
                    value="Temperature",
                    label="Condition to Screen",
                    info="This condition will be swept across its full range; its slider value below is hidden.",
                )

                # Fixed-value sliders — JS hides the swept condition's column (no Gradio re-render)
                with gr.Row():
                    with gr.Column(elem_id="s-temp-col"):
                        s_temp = gr.Slider(0, 60,   value=25.0,  step=0.5,
                                           label="Temperature (°C)  [fixed]")
                    with gr.Column(elem_id="s-conc-col"):
                        s_conc = gr.Slider(0, 1000, value=100.0, step=5.0,
                                           label="Concentration (µM)  [fixed]")
                    with gr.Column(elem_id="s-ph-col"):
                        s_pH   = gr.Slider(0, 14,   value=7.3,   step=0.1,
                                           label="pH  [fixed]")

                # Pure JS toggle — bypasses Gradio server update, so slider fills are preserved
                screen_radio.change(
                    fn=None,
                    inputs=[screen_radio],
                    outputs=[],
                    js="""(screen_name) => {
                        const map = {Temperature: 's-temp-col', Concentration: 's-conc-col', pH: 's-ph-col'};
                        for (const [cond, id] of Object.entries(map)) {
                            const el = document.getElementById(id);
                            if (el) el.style.display = (cond === screen_name) ? 'none' : 'flex';
                        }
                    }""",
                )

                # Advanced conditions (always fixed during screening)
                with gr.Accordion("⚙️ Advanced Conditions  (Salts & Crowding Agents)", open=False):
                    gr.Markdown("Default values represent a common physiological buffer (160 mM NaCl).")
                    with gr.Row():
                        s_nacl  = gr.Slider(0, 2000, value=160.0, step=10.0, label="NaCl (mM)")
                        s_mgcl2 = gr.Slider(0, 50,   value=0.0,   step=1.0,  label="MgCl₂ (mM)")
                        s_kcl   = gr.Slider(0, 1000, value=0.0,   step=10.0, label="KCl (mM)")
                        s_glyc  = gr.Slider(0, 10,   value=0.0,   step=0.5,  label="Glycerol (%)")
                    with gr.Row():
                        s_peg1   = gr.Slider(0, 50, value=0, step=1, label="PEG 300–1000 (%)")
                        s_peg2   = gr.Slider(0, 50, value=0, step=1, label="PEG 3k–6k (%)")
                        s_peg3   = gr.Slider(0, 50, value=0, step=1, label="PEG 8k–20k (%)")
                    with gr.Row():
                        s_ficoll = gr.Slider(0, 50, value=0, step=1, label="Ficoll (%)")
                        s_dext40 = gr.Slider(0, 50, value=0, step=1, label="Dextran ≤40 kDa (%)")
                        s_dext70 = gr.Slider(0, 50, value=0, step=1, label="Dextran ≥70 kDa (%)")

                s_smooth = gr.Slider(
                    1, 21, value=15, step=2,
                    label="Smoothing Window Size",
                    info="Odd window size for moving-average smoothing of the screening curve (1 = no smoothing).",
                )

                screen_btn  = gr.Button("📊 Run Condition Screening", variant="primary")
                screen_plot = gr.Plot(label="Screening Result")
                screen_text = gr.Markdown("")

                screen_btn.click(
                    fn=cb_screen,
                    inputs=[
                        feat_state, screen_radio,
                        s_temp, s_conc, s_pH,
                        s_nacl, s_mgcl2, s_kcl, s_glyc,
                        s_peg1, s_peg2, s_peg3, s_ficoll, s_dext40, s_dext70,
                        s_smooth,
                    ],
                    outputs=[screen_plot, screen_text],
                )

    # Registered here (after Step 3 components exist) since cb_extract also
    # clears the Predict/Screening panels so a re-extracted sequence can
    # never leave a stale result from the previous sequence on screen.
    extract_btn.click(
        fn=cb_extract,
        inputs=[seq_box],
        outputs=[feat_state, extract_status, pred_plot, pred_text, screen_plot, screen_text],
    )

    # Selecting a built-in example runs extraction automatically (no manual
    # Step-2 click needed) — cache-hit examples resolve near-instantly since
    # cb_extract reads their pre-computed feature from assets/ instead of
    # calling the T5 model. example_picker is None when no example has a
    # cached feature yet (see AVAILABLE_EXAMPLES above).
    if example_picker is not None:
        example_picker.load_input_event.then(
            fn=cb_extract,
            inputs=[seq_box],
            outputs=[feat_state, extract_status, pred_plot, pred_text, screen_plot, screen_text],
        )

    # ── Footer ────────────────────────────────────────────────────────────────
    gr.Markdown("""
---
**Authors:** Jangwon Bae†, Minjun Kang†, Donghyuk Lee, Kuk-Jin Yoon*, Yongwon Jung*
**Paper:** [bioRxiv 2025.12.28.696755](https://doi.org/10.64898/2025.12.28.696755)
**GitHub:** [NearNiah/LLPSense](https://github.com/NearNiah/LLPSense)
""")


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)