ChiralFold-App / app.py
The-Philosopher's picture
Deploy ChiralFold web Space from CI
aca9d91 verified
Raw
History Blame Contribute Delete
16.7 kB
#!/usr/bin/env python3
"""
ChiralFold Hugging Face Space entrypoint.
Self-contained Gradio UI that depends on the `chiralfold` package
(installed via requirements.txt from GitHub). Theme/CSS is loaded from
``web.theme`` when available, otherwise from bundled ``theme.css``.
"""
from __future__ import annotations
import json
import os
import shutil
import tempfile
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any, Dict, Optional, Tuple
import gradio as gr
from chiralfold import __version__
from chiralfold.af3_correct import correct_af3_output, detect_chirality_violations
from chiralfold.auditor import audit_pdb, format_report
from chiralfold.pdb_pipeline import mirror_pdb
try:
from chiralfold.fetch import FetchError, fetch_structure
from chiralfold.structure_io import SUPPORTED_UPLOAD_SUFFIXES
except ImportError: # Space image may briefly lag behind master
FetchError = ValueError # type: ignore
SUPPORTED_UPLOAD_SUFFIXES = (".pdb", ".ent", ".cif", ".mmcif", ".fasta", ".fa", ".faa")
fetch_structure = None # type: ignore
MAX_UPLOAD_BYTES = int(os.environ.get("CHIRALFOLD_MAX_UPLOAD_MB", "25")) * 1024 * 1024
RCSB_DOWNLOAD = "https://files.rcsb.org/download/{pdb_id}.pdb"
AFDB_PDB_URL = "https://alphafold.ebi.ac.uk/files/AF-{uid}-F1-model_v4.pdb"
SPACE_ROOT = Path(__file__).resolve().parent
EXAMPLE_TOY = SPACE_ROOT / "examples" / "toy_ubiquitin_fragment.pdb"
EXAMPLE_INV = SPACE_ROOT / "examples" / "synthetic_l_ala_inverted.pdb"
def _load_css() -> str:
"""Prefer bundled theme.css so Docker COPY updates always win.
Loading ``web.theme`` first was wrong for Spaces: ``pip install`` from
``requirements.txt`` is Docker-layer-cached when the pin is unchanged, so
an old package CSS would override a freshly uploaded ``theme.css``.
"""
css_path = SPACE_ROOT / "theme.css"
if css_path.is_file():
return css_path.read_text()
try:
from web.theme import CUSTOM_CSS as css # type: ignore
return css
except Exception:
return "/* missing theme */"
def _make_theme():
try:
from web.theme import make_theme
return make_theme()
except Exception:
return gr.themes.Soft(
primary_hue="teal",
secondary_hue="slate",
neutral_hue="slate",
font=["IBM Plex Sans", "Segoe UI", "Helvetica", "Arial", "sans-serif"],
)
def _header_html(version: str) -> str:
try:
from web.theme import header_html
return header_html(version)
except Exception:
return f"<div id='cf-header'><h1>ChiralFold</h1><p>v{version}</p></div>"
try:
from web.ui_format import audit_result_markdown, error_markdown, footer_html
except ImportError:
import html as _html
import traceback as _traceback
def report_pre_html(text: str) -> str:
return f'<pre class="cf-report">{_html.escape(text)}</pre>'
def audit_result_markdown(kpis_html: str, filename: str, summary: str) -> str:
return (
f"## Stereochemistry audit\n\n{kpis_html}\n\n"
f"**File:** `{filename}`\n\n"
f"{report_pre_html(summary)}\n\n"
"Download the full JSON report below."
)
def error_markdown(exc: BaseException) -> str:
tb = _html.escape(_traceback.format_exc())
return (
f"**Error:** {_html.escape(str(exc))}\n\n"
f'<pre class="cf-error">{tb}</pre>'
)
def footer_html(version: str, *, include_hf_link: bool = False) -> str:
return (
f'<div id="cf-footer"><p class="cf-footer-line">'
f"<strong>ChiralFold v{_html.escape(version)}</strong> · "
f'<a href="https://github.com/Tommaso-R-Marena/ChiralFold">GitHub</a> · '
f'<code class="cf-install-cmd">pip install chiralfold</code></p>'
f'<p class="cf-footer-note">Processed in-session; nothing stored permanently.</p></div>'
)
CUSTOM_CSS = _load_css()
def _as_path(uploaded) -> str:
if uploaded is None:
raise gr.Error(
"Upload a PDB/mmCIF/FASTA file, fetch an ID, or load an example first."
)
path = uploaded if isinstance(uploaded, str) else getattr(uploaded, "name", None)
if not path:
raise gr.Error("Could not read the uploaded file.")
path = str(path)
size = Path(path).stat().st_size
if size > MAX_UPLOAD_BYTES:
raise gr.Error(
f"File too large ({size / 1e6:.1f} MB). "
f"Limit is {MAX_UPLOAD_BYTES / 1e6:.0f} MB."
)
if size == 0:
raise gr.Error("Uploaded file is empty.")
if fetch_structure is not None:
try:
return fetch_structure(path, max_bytes=MAX_UPLOAD_BYTES).path
except FetchError as exc:
raise gr.Error(str(exc)) from exc
# Legacy fallback: PDB/ENT only until package with fetch is installed.
if not path.lower().endswith((".pdb", ".ent")):
raise gr.Error(
"Upgrade chiralfold for mmCIF/FASTA support, or upload a .pdb file."
)
return path
def _copy_with_name(src: str, stem: str, suffix: str) -> str:
out_dir = Path(tempfile.mkdtemp(prefix="chiralfold_dl_"))
dest = out_dir / f"{stem}{suffix}"
shutil.copy2(src, dest)
return str(dest)
def _stem_from_path(path: str) -> str:
name = Path(path).stem
for junk in ("_corrected", "_mirror", "_audit"):
name = name.replace(junk, "")
return name or "structure"
def _kpi_html(cells):
parts = ['<div class="cf-kpi">']
for label, value, cls in cells:
parts.append(
f'<div class="cell {cls}"><div class="label">{label}</div>'
f'<div class="value">{value}</div></div>'
)
parts.append("</div>")
return "".join(parts)
def _legacy_fetch(query: str) -> Tuple[Optional[str], str]:
"""Minimal RCSB / AFDB fetch when chiralfold.fetch is unavailable."""
q = (query or "").strip()
if len(q) == 4 and q.isalnum():
pid = q.upper()
url = RCSB_DOWNLOAD.format(pdb_id=pid)
try:
with urllib.request.urlopen(url, timeout=30) as resp:
data = resp.read()
except Exception as exc: # noqa: BLE001
return None, f"Could not fetch {pid}: {exc}"
if not data or (b"ATOM" not in data and b"HETATM" not in data):
return None, f"Downloaded file for {pid} does not look like a PDB."
out = Path(tempfile.mkdtemp(prefix="chiralfold_rcsb_")) / f"{pid}.pdb"
out.write_bytes(data)
return str(out), f"Loaded **{pid}** from RCSB ({len(data) / 1024:.0f} KB)."
uid = q.upper()
if uid.startswith("AF-") and "-F" in uid:
uid = uid.split("-")[1]
url = AFDB_PDB_URL.format(uid=uid)
try:
with urllib.request.urlopen(url, timeout=45) as resp:
data = resp.read()
except Exception as exc: # noqa: BLE001
return None, f"Could not fetch AFDB/{uid}: {exc}"
if not data or (b"ATOM" not in data and b"HETATM" not in data):
return None, f"AFDB download for {uid} does not look like a PDB."
out = Path(tempfile.mkdtemp(prefix="chiralfold_afdb_")) / f"AF-{uid}-F1.pdb"
out.write_bytes(data)
return str(out), f"Loaded **AF-{uid}-F1** from AlphaFold DB ({len(data) / 1024:.0f} KB)."
def fetch_pdb_id(query: str) -> Tuple[Optional[str], str]:
q = (query or "").strip()
if not q:
return (
None,
"Enter a PDB ID (1UBQ), UniProt accession (P04637), or AFDB id (AF-P04637-F1).",
)
if fetch_structure is not None:
try:
result = fetch_structure(q, max_bytes=MAX_UPLOAD_BYTES)
except FetchError as exc:
return None, str(exc)
except Exception as exc: # noqa: BLE001
return None, f"Could not fetch {q}: {exc}"
kb = (result.bytes_downloaded or Path(result.path).stat().st_size) / 1024
return (
result.path,
f"Loaded **{result.label}** from `{result.source}` ({kb:.0f} KB).",
)
return _legacy_fetch(q)
def load_example(which: str) -> Tuple[Optional[str], str]:
mapping = {
"toy": (EXAMPLE_TOY, "Toy ubiquitin fragment (10 residues) — bundled, offline."),
"inverted": (
EXAMPLE_INV,
"Synthetic inverted L-Ala — demonstrates chirality correction.",
),
}
path, msg = mapping.get(which, (None, "Unknown example."))
if path is None or not Path(path).is_file():
return None, f"Example not found: {which}"
dest = Path(tempfile.mkdtemp(prefix="chiralfold_ex_")) / path.name
shutil.copy2(path, dest)
return str(dest), msg
def run_audit(uploaded) -> Tuple[str, Optional[str]]:
try:
src = _as_path(uploaded)
report = audit_pdb(src)
chir = report["chirality"]
rama = report["ramachandran"]
n_wrong = chir.get("n_wrong", 0)
score = report["overall_score"]
score_cls = "ok" if score >= 80 else ("warn" if score >= 60 else "bad")
chir_cls = "ok" if n_wrong == 0 else "bad"
clash = report["clashes"]["clash_score"]
clash_cls = "ok" if clash < 30 else ("warn" if clash < 60 else "bad")
kpis = _kpi_html(
[
("Overall score", f"{score:.1f}", score_cls),
("Cα correct %", f"{chir['pct_correct']:.1f}", chir_cls),
("Chirality violations", str(n_wrong), chir_cls),
("Ramachandran favored %", f"{rama['pct_favored']:.1f}", "ok"),
("Clashscore", f"{clash:.1f}", clash_cls),
]
)
summary = format_report(report)
md = audit_result_markdown(kpis, Path(src).name, summary)
with tempfile.NamedTemporaryFile(
mode="w", suffix=".json", delete=False, prefix="chiralfold_audit_"
) as fp:
json.dump(report, fp, indent=2, default=float)
raw_json = fp.name
return md, _copy_with_name(raw_json, _stem_from_path(src), "_audit.json")
except gr.Error:
raise
except Exception as exc: # noqa: BLE001
return error_markdown(exc), None
def run_correct(uploaded) -> Tuple[str, Optional[str]]:
try:
src = _as_path(uploaded)
before = detect_chirality_violations(src)
out_raw = tempfile.NamedTemporaryFile(
suffix="_corrected.pdb", delete=False, prefix="chiralfold_"
)
out_raw.close()
result = correct_af3_output(src, out_raw.name)
corr = result.get("correction", {})
after = result.get("after", {})
n_before = before.get(
"n_violations", result.get("before", {}).get("n_violations", 0)
)
n_after = after.get("n_violations", 0)
n_fixed = corr.get("n_corrected", max(0, n_before - n_after))
before_cls = "ok" if n_before == 0 else "bad"
after_cls = "ok" if n_after == 0 else "bad"
kpis = _kpi_html(
[
("Before", str(n_before), before_cls),
("Fixed", str(n_fixed), "ok"),
("After (residual)", str(n_after), after_cls),
]
)
md = (
f"## Chirality correction\n\n{kpis}\n\n"
f"**File:** `{Path(src).name}`\n\n"
"**Download the corrected PDB below.**"
)
return md, _copy_with_name(out_raw.name, _stem_from_path(src), "_corrected.pdb")
except gr.Error:
raise
except Exception as exc: # noqa: BLE001
return error_markdown(exc), None
def run_mirror(uploaded, axis: str, rename: bool) -> Tuple[str, Optional[str]]:
try:
src = _as_path(uploaded)
out_raw = tempfile.NamedTemporaryFile(
suffix="_mirror.pdb", delete=False, prefix="chiralfold_"
)
out_raw.close()
info: Dict[str, Any] = mirror_pdb(
src, out_raw.name, axis=axis.lower(), rename_residues=rename
)
kpis = _kpi_html(
[
("Atoms mirrored", str(info.get("n_atoms", "?")), "ok"),
("Residues", str(info.get("n_residues", "?")), "ok"),
("Axis", axis.upper(), "ok"),
("Rename D↔L", "yes" if rename else "no", "ok"),
]
)
md = (
f"## Mirror-image (L↔D) transform\n\n{kpis}\n\n"
"Exact coordinate reflection for mirror-image binder design.\n\n"
"**Download the mirrored PDB below.**"
)
return md, _copy_with_name(out_raw.name, _stem_from_path(src), "_mirror.pdb")
except gr.Error:
raise
except Exception as exc: # noqa: BLE001
return error_markdown(exc), None
def build_app() -> gr.Blocks:
with gr.Blocks(
title=f"ChiralFold Web v{__version__}",
css=CUSTOM_CSS,
theme=_make_theme(),
) as demo:
gr.HTML(_header_html(__version__))
with gr.Row():
with gr.Column(scale=5, elem_classes=["cf-panel"]):
gr.Markdown("### 1 · Load a structure")
pdb_in = gr.File(
label="Upload PDB / mmCIF / FASTA",
file_types=list(SUPPORTED_UPLOAD_SUFFIXES),
type="filepath",
)
with gr.Row():
pdb_id = gr.Textbox(
label="Or fetch (RCSB / AlphaFold DB)",
placeholder="1UBQ · P04637 · AF-P04637-F1",
max_lines=1,
scale=3,
)
fetch_btn = gr.Button("Fetch", scale=1)
with gr.Row():
ex_toy = gr.Button("Example: ubiquitin fragment", size="sm")
ex_inv = gr.Button("Example: inverted Ala", size="sm")
status = gr.Markdown(
"Upload PDB/mmCIF/FASTA, fetch an ID, or load an example."
)
gr.Markdown("### 2 · Mirror options")
axis = gr.Radio(
choices=["X", "Y", "Z"], value="X", label="Reflection axis"
)
rename = gr.Checkbox(
value=True, label="Rename residue codes L↔D after mirroring"
)
with gr.Column(scale=7, elem_classes=["cf-panel"]):
gr.Markdown("### 3 · Choose an action")
with gr.Tabs():
with gr.Tab("Audit"):
gr.Markdown(
"Full stereochemistry report with downloadable JSON."
)
audit_btn = gr.Button(
"Run audit → download JSON", variant="primary", size="lg"
)
with gr.Tab("Correct chirality"):
gr.Markdown(
"Fix inverted stereocenters (AF3 / diffusion outputs)."
)
correct_btn = gr.Button(
"Correct → download PDB", variant="primary", size="lg"
)
with gr.Tab("Mirror image"):
gr.Markdown(
"Exact L↔D coordinate reflection for binder design."
)
mirror_btn = gr.Button(
"Mirror → download PDB", variant="primary", size="lg"
)
report_out = gr.Markdown(
value="Results appear here after you run an action."
)
download_out = gr.File(label="Download result")
fetch_btn.click(fetch_pdb_id, inputs=pdb_id, outputs=[pdb_in, status])
ex_toy.click(lambda: load_example("toy"), outputs=[pdb_in, status])
ex_inv.click(lambda: load_example("inverted"), outputs=[pdb_in, status])
audit_btn.click(run_audit, inputs=pdb_in, outputs=[report_out, download_out])
correct_btn.click(
run_correct, inputs=pdb_in, outputs=[report_out, download_out]
)
mirror_btn.click(
run_mirror,
inputs=[pdb_in, axis, rename],
outputs=[report_out, download_out],
)
gr.HTML(footer_html(__version__))
return demo
demo = build_app()
if __name__ == "__main__":
demo.queue(default_concurrency_limit=4)
# Theme via header toggle or ?__theme=light|dark (Gradio).
demo.launch(
server_name="0.0.0.0",
server_port=7860,
show_error=True,
)