Spaces:
Sleeping
Sleeping
| #!/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, | |
| ) | |