"""
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 re
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 torch
import gradio as gr
import spaces
warnings.filterwarnings("ignore")
# ── Paths ─────────────────────────────────────────────────────────────────────
BASE_DIR = Path(__file__).parent
# ── 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")
# α-synuclein (used as built-in example)
EXAMPLE_SEQ = (
"MDVFMKGLSKAKEGVVAAAEKTKQGVAEAAGKTKEGVLYVGSKTKEGVVHGVATVAEKTK"
"EQVTNVGGAVVTGVTAVAQKTVEGAGSIAAATGFVKKDQLGKNEEGAPQEGILEDMPVDP"
"DNEAYEMPSEEGYQDYEPEA"
)
# ── Lazy-loaded singletons ────────────────────────────────────────────────────
_t5_tokenizer = None
_t5_model = None
_llps_model = None
def get_device() -> str:
return "cuda" if torch.cuda.is_available() else "cpu"
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
def load_t5():
global _t5_tokenizer, _t5_model
if _t5_model is None:
from transformers import AutoTokenizer, T5EncoderModel
import transformers.utils.import_utils as _hf_utils
# Rostlab/prot_t5_xl_half_uniref50-enc is only available as .bin (no safetensors).
# transformers 5.x blocks torch.load on torch < 2.6 due to CVE-2025-32434.
# We bypass that gate for this specific trusted checkpoint from the official
# HuggingFace Hub. The check lives in two places — import_utils AND the
# locally-imported name in modeling_utils — so both must be patched.
import transformers.modeling_utils as _modeling_utils
_noop = lambda: None
_orig_hf = _hf_utils.check_torch_load_is_safe
_orig_mdl = _modeling_utils.check_torch_load_is_safe
_hf_utils.check_torch_load_is_safe = _noop
_modeling_utils.check_torch_load_is_safe = _noop
try:
# Pre-fetched into the HF cache at Docker build time (see Dockerfile's
# snapshot_download step), so ZeroGPU calls never hit the network here —
# a cold download would blow past the GPU-allocation time budget and
# get aborted. local_files_only=True enforces that guarantee.
repo_id = "Rostlab/prot_t5_xl_half_uniref50-enc"
dev = get_device()
dtype = torch.float16 if dev == "cuda" else torch.float32
_t5_tokenizer = AutoTokenizer.from_pretrained(
repo_id, do_lower_case=False, local_files_only=True
)
_t5_model = (
T5EncoderModel.from_pretrained(
repo_id, torch_dtype=dtype, local_files_only=True
)
.to(dev)
.eval()
)
_t5_model.requires_grad_(False)
finally:
_hf_utils.check_torch_load_is_safe = _orig_hf # always restore
_modeling_utils.check_torch_load_is_safe = _orig_mdl
return _t5_tokenizer, _t5_model
# ── Feature extraction ────────────────────────────────────────────────────────
@spaces.GPU
def extract_t5_feature(sequence: str) -> np.ndarray:
"""
Return mean-pooled ProtT5-XL embedding for a single sequence.
Output shape: (1024,) — matches the feature dimension LLPSense was trained on.
"""
tok, mdl = load_t5()
dev = get_device()
# Replace ambiguous residues with X (same as original pipeline)
clean = re.sub(r"[UZOB]", "X", sequence.strip().upper())
spaced = " ".join(list(clean))
enc = tok([spaced], add_special_tokens=True, padding="longest", return_tensors="pt")
input_ids = enc["input_ids"].to(dev)
attention_mask = enc["attention_mask"].to(dev)
with torch.no_grad():
hidden = mdl(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
# Mean-pool over sequence positions (ignoring padding tokens)
feat = (hidden * attention_mask[..., None]).sum(dim=1) / \
attention_mask.sum(dim=1, keepdim=True)
return feat[0].float().cpu().numpy() # (1024,)
# ── 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])
# ── 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) -> plt.Figure:
"""Line graph for condition screening result."""
xlabel = SCREEN_LABELS[screen_name]
LLPS_COLOR = "#e74c3c"
NON_COLOR = "#2980b9"
fig, ax = plt.subplots(figsize=(9, 5))
ax.plot(xvals, probs, lw=2.5, color=NON_COLOR, label="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, probs, 0.5,
where=(probs >= 0.5), alpha=0.22,
color=LLPS_COLOR, label="LLPS region", zorder=2)
ax.fill_between(xvals, probs, 0.5,
where=(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 = """
ProtT5-XL feature extraction in progress…
First run loads the model onto the GPU — this may take a moment.
"""
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''
f'✅ Feature extracted'
f'Length: {seq_len} AA'
f'
')
def _status_warn(msg: str) -> str:
return f'⚠️ {msg}
'
def _status_err(msg: str) -> str:
return f'❌ {msg}
'
# ── 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: {''.join(sorted(invalid))}"), None, "", None, ""
return
# ── invalidate old feature/results immediately, then show spinner ─────────
yield None, _SPINNER_HTML, None, "", None, ""
try:
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):
"""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)
fig = screening_figure(xvals, probs, screen_name)
peak_idx = probs.argmax()
xlabel = SCREEN_LABELS[screen_name]
txt = (f"Screening completed. \n"
f"Peak probability **{probs[peak_idx]:.4f}** "
f"at {xlabel} = **{xvals[peak_idx]:.1f}** \n"
f"LLPS-positive range: "
f"**{(probs >= 0.5).sum()}** / {len(probs)} points ≥ 0.5")
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=EXAMPLE_SEQ,
elem_id="seq-input",
)
gr.Examples(
examples=[[EXAMPLE_SEQ]],
inputs=[seq_box],
label="Example: α-synuclein (SNCA)",
)
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 (%)")
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,
],
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],
)
# ── 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)