Spaces:
Running
Running
| import json | |
| import os | |
| import subprocess | |
| from pathlib import Path | |
| import uuid | |
| from datetime import datetime | |
| import gradio as gr | |
| import pandas as pd | |
| from huggingface_hub import snapshot_download | |
| from rdkit import Chem | |
| from rdkit.Chem import Descriptors, Crippen, Lipinski, rdMolDescriptors, QED | |
| BUNDLE_REPO = os.getenv("BUNDLE_REPO", "abdullahsamra/egfr-ai-bundle") | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| CACHE_DIR = Path("bundle_cache") | |
| EXAMPLE_JSON = None | |
| API_SPEC_TEXT = None | |
| BACKEND_CONTRACT_TEXT = None | |
| BUNDLE_ROOT = None | |
| def download_bundle_light(): | |
| global EXAMPLE_JSON, API_SPEC_TEXT, BACKEND_CONTRACT_TEXT, BUNDLE_ROOT | |
| try: | |
| local_dir = snapshot_download( | |
| repo_id=BUNDLE_REPO, | |
| repo_type="dataset", | |
| token=HF_TOKEN, | |
| allow_patterns=[ | |
| "examples/generate_response.example.json", | |
| "docs/API_RESPONSE_SPEC.md", | |
| "configs/backend_contract.json", | |
| "configs/sampling_model.bundle.toml", | |
| "models/generator/egfr_generator.chkpt", | |
| "README.md", | |
| ], | |
| local_dir=str(CACHE_DIR), | |
| ) | |
| root = Path(local_dir) | |
| BUNDLE_ROOT = root | |
| example_path = root / "examples" / "generate_response.example.json" | |
| spec_path = root / "docs" / "API_RESPONSE_SPEC.md" | |
| contract_path = root / "configs" / "backend_contract.json" | |
| if example_path.exists(): | |
| EXAMPLE_JSON = json.loads(example_path.read_text(encoding="utf-8")) | |
| if spec_path.exists(): | |
| API_SPEC_TEXT = spec_path.read_text(encoding="utf-8") | |
| if contract_path.exists(): | |
| BACKEND_CONTRACT_TEXT = contract_path.read_text(encoding="utf-8") | |
| return True, f"Loaded docs/examples from {BUNDLE_REPO}" | |
| except Exception as e: | |
| return False, f"Could not load private bundle. Check HF_TOKEN secret and BUNDLE_REPO variable. Error: {e}" | |
| def canonicalize(smiles: str): | |
| mol = Chem.MolFromSmiles(str(smiles).strip()) | |
| if mol is None: | |
| return None, None | |
| return mol, Chem.MolToSmiles(mol, isomericSmiles=True) | |
| def try_sa_score(mol): | |
| try: | |
| from rdkit.Contrib.SA_Score import sascorer | |
| return round(float(sascorer.calculateScore(mol)), 4) | |
| except Exception: | |
| return None | |
| def compute_row(smiles: str): | |
| mol, can = canonicalize(smiles) | |
| if mol is None: | |
| return { | |
| "SMILES": smiles, | |
| "valid": False, | |
| "canonical_smiles": None, | |
| "mw": None, | |
| "logp": None, | |
| "tpsa": None, | |
| "hbd": None, | |
| "hba": None, | |
| "rot_bonds": None, | |
| "qed": None, | |
| "sa_score": None, | |
| "pred_pAff_mean": None, | |
| "docking_score": None, | |
| "docking_status": "not_run", | |
| "docking_pose_file": None, | |
| "note": "Invalid SMILES", | |
| } | |
| return { | |
| "SMILES": smiles, | |
| "valid": True, | |
| "canonical_smiles": can, | |
| "mw": round(float(Descriptors.MolWt(mol)), 4), | |
| "logp": round(float(Crippen.MolLogP(mol)), 4), | |
| "tpsa": round(float(rdMolDescriptors.CalcTPSA(mol)), 4), | |
| "hbd": int(Lipinski.NumHDonors(mol)), | |
| "hba": int(Lipinski.NumHAcceptors(mol)), | |
| "rot_bonds": int(Lipinski.NumRotatableBonds(mol)), | |
| "qed": round(float(QED.qed(mol)), 4), | |
| "sa_score": try_sa_score(mol), | |
| "pred_pAff_mean": None, | |
| "docking_score": None, | |
| "docking_status": "not_run", | |
| "docking_pose_file": None, | |
| "note": "CPU demo: DeepPurpose and real docking are disabled in this Space.", | |
| } | |
| def score_smiles(smiles_text): | |
| smiles_list = [] | |
| for part in str(smiles_text).replace(",", "\n").splitlines(): | |
| part = part.strip() | |
| if part: | |
| smiles_list.append(part) | |
| if not smiles_list: | |
| return pd.DataFrame(), { | |
| "status": "error", | |
| "message": "Please enter at least one SMILES.", | |
| } | |
| rows = [compute_row(s) for s in smiles_list] | |
| df = pd.DataFrame(rows) | |
| response = { | |
| "status": "completed", | |
| "mode": "score_smiles_cpu_demo", | |
| "docking_mode": "off", | |
| "summary": { | |
| "num_input": len(smiles_list), | |
| "num_valid": int(df["valid"].sum()), | |
| "num_returned": len(df), | |
| }, | |
| "columns": list(df.columns), | |
| "results": rows, | |
| "warnings": [ | |
| "This Space is a CPU demo.", | |
| "DeepPurpose pred_pAff_mean is not computed in this Space version.", | |
| "Real docking is not run in this CPU Space; docking_status is not_run.", | |
| "Outputs are computational predictions only.", | |
| ], | |
| } | |
| return df, response | |
| def get_example_table(): | |
| if EXAMPLE_JSON and "results" in EXAMPLE_JSON: | |
| return pd.DataFrame(EXAMPLE_JSON["results"]) | |
| return pd.DataFrame([ | |
| {"message": "Example response not loaded. Check HF_TOKEN secret or BUNDLE_REPO variable."} | |
| ]) | |
| def get_example_json(): | |
| if EXAMPLE_JSON: | |
| return EXAMPLE_JSON | |
| return { | |
| "status": "not_loaded", | |
| "message": "Example response not loaded. Check HF_TOKEN secret or BUNDLE_REPO variable.", | |
| } | |
| def get_api_spec(): | |
| return API_SPEC_TEXT or "API_RESPONSE_SPEC.md not loaded. Check HF_TOKEN secret or BUNDLE_REPO variable." | |
| def get_contract(): | |
| if BACKEND_CONTRACT_TEXT: | |
| try: | |
| return json.loads(BACKEND_CONTRACT_TEXT) | |
| except Exception: | |
| return {"raw": BACKEND_CONTRACT_TEXT} | |
| return { | |
| "status": "not_loaded", | |
| "message": "backend_contract.json not loaded. Check HF_TOKEN secret or BUNDLE_REPO variable.", | |
| } | |
| def _write_reinvent_sampling_config(job_id, num_molecules): | |
| if BUNDLE_ROOT is None: | |
| raise RuntimeError("Bundle root is not loaded. Check HF_TOKEN and BUNDLE_REPO.") | |
| checkpoint = Path(BUNDLE_ROOT) / "models" / "generator" / "egfr_generator.chkpt" | |
| if not checkpoint.exists(): | |
| raise FileNotFoundError(f"Generator checkpoint not found: {checkpoint}") | |
| output_dir = Path("outputs") | |
| output_dir.mkdir(exist_ok=True) | |
| raw_csv = output_dir / f"{job_id}_raw_sampling.csv" | |
| log_file = output_dir / f"{job_id}_reinvent.log" | |
| resolved_json = output_dir / f"{job_id}_sampling_resolved.json" | |
| config_path = output_dir / f"{job_id}_sampling.toml" | |
| config_text = f""" | |
| run_type = "sampling" | |
| device = "cpu" | |
| json_out_config = "{resolved_json.resolve()}" | |
| [parameters] | |
| model_file = "{checkpoint.resolve()}" | |
| output_file = "{raw_csv.resolve()}" | |
| num_smiles = {int(num_molecules)} | |
| unique_molecules = true | |
| randomize_smiles = true | |
| """ | |
| config_path.write_text(config_text.strip() + "\n", encoding="utf-8") | |
| return config_path, raw_csv, log_file | |
| def _run_reinvent_sampling(job_id, num_molecules): | |
| config_path, raw_csv, log_file = _write_reinvent_sampling_config(job_id, num_molecules) | |
| cmd = [ | |
| "reinvent", | |
| "-l", | |
| str(log_file.resolve()), | |
| str(config_path.resolve()), | |
| ] | |
| proc = subprocess.run( | |
| cmd, | |
| text=True, | |
| stdout=subprocess.PIPE, | |
| stderr=subprocess.STDOUT, | |
| timeout=300, | |
| ) | |
| if proc.returncode != 0: | |
| raise RuntimeError( | |
| "REINVENT sampling failed.\n" | |
| f"Command: {' '.join(cmd)}\n" | |
| f"Output:\n{proc.stdout[-4000:]}" | |
| ) | |
| if not raw_csv.exists(): | |
| raise RuntimeError(f"REINVENT finished but output CSV was not found: {raw_csv}") | |
| return raw_csv, log_file | |
| def _enrich_reinvent_csv(raw_csv, return_top_k): | |
| raw_df = pd.read_csv(raw_csv) | |
| if "SMILES" not in raw_df.columns: | |
| raise ValueError(f"REINVENT output missing SMILES column. Columns: {list(raw_df.columns)}") | |
| rows = [] | |
| for _, r in raw_df.iterrows(): | |
| row = compute_row(str(r["SMILES"])) | |
| if "SMILES_state" in raw_df.columns: | |
| try: | |
| row["SMILES_state"] = int(r["SMILES_state"]) | |
| except Exception: | |
| row["SMILES_state"] = None | |
| else: | |
| row["SMILES_state"] = None | |
| if "NLL" in raw_df.columns: | |
| try: | |
| row["NLL"] = float(r["NLL"]) | |
| except Exception: | |
| row["NLL"] = None | |
| else: | |
| row["NLL"] = None | |
| # Keep these stable for frontend/backend filters. | |
| row["pred_pAff_mean"] = None | |
| row["docking_score"] = None | |
| row["docking_status"] = "not_run" | |
| row["docking_pose_file"] = None | |
| rows.append(row) | |
| df = pd.DataFrame(rows) | |
| preferred_cols = [ | |
| "SMILES", | |
| "SMILES_state", | |
| "NLL", | |
| "valid", | |
| "canonical_smiles", | |
| "mw", | |
| "logp", | |
| "tpsa", | |
| "hbd", | |
| "hba", | |
| "rot_bonds", | |
| "qed", | |
| "sa_score", | |
| "pred_pAff_mean", | |
| "docking_score", | |
| "docking_status", | |
| "docking_pose_file", | |
| "note", | |
| ] | |
| cols = [c for c in preferred_cols if c in df.columns] + [c for c in df.columns if c not in preferred_cols] | |
| df = df[cols] | |
| return_top_k = int(return_top_k) | |
| if return_top_k > 0: | |
| df = df.head(return_top_k) | |
| return df | |
| def generate_job_demo(num_molecules, return_top_k, docking_mode, generation_backend): | |
| """ | |
| Generation endpoint for the Space. | |
| Demo mode uses packaged example rows. | |
| Live mode tries to run REINVENT CPU sampling from the packaged checkpoint. | |
| """ | |
| job_id = f"gen_{uuid.uuid4().hex[:8]}" | |
| created_at = datetime.utcnow().isoformat() + "Z" | |
| num_molecules = int(num_molecules) | |
| return_top_k = int(return_top_k) | |
| warnings = [] | |
| try: | |
| if generation_backend == "live_reinvent_cpu": | |
| raw_csv, log_file = _run_reinvent_sampling(job_id, num_molecules) | |
| df = _enrich_reinvent_csv(raw_csv, return_top_k) | |
| mode = "live_reinvent_cpu_sampling" | |
| status = "completed" | |
| warnings.extend([ | |
| "Live REINVENT CPU sampling was executed in this Space.", | |
| "DeepPurpose scoring is not enabled in this Space version; pred_pAff_mean is null.", | |
| "Real docking is not run in this CPU Space; docking_status is not_run.", | |
| ]) | |
| else: | |
| if EXAMPLE_JSON and "results" in EXAMPLE_JSON: | |
| rows = EXAMPLE_JSON["results"] | |
| else: | |
| rows = [compute_row("CCO"), compute_row("c1ccccc1")] | |
| selected = rows[: max(1, min(return_top_k, len(rows)))] | |
| clean_rows = [] | |
| for row in selected: | |
| clean = dict(row) | |
| clean["docking_score"] = None | |
| clean["docking_status"] = "not_run" | |
| clean["docking_pose_file"] = None | |
| if "pred_pAff_mean" not in clean: | |
| clean["pred_pAff_mean"] = None | |
| clean_rows.append(clean) | |
| df = pd.DataFrame(clean_rows) | |
| mode = "generate_job_cpu_demo" | |
| status = "completed" | |
| warnings.extend([ | |
| "This is a CPU demo generation response.", | |
| "Live REINVENT generation was not executed because demo mode was selected.", | |
| "The response shape matches the planned backend contract.", | |
| "Real docking is not run in this CPU Space; docking_status is not_run.", | |
| ]) | |
| except Exception as e: | |
| df = pd.DataFrame() | |
| mode = "live_reinvent_cpu_sampling" | |
| status = "failed" | |
| warnings.append(str(e)) | |
| output_dir = Path("outputs") | |
| output_dir.mkdir(exist_ok=True) | |
| csv_path = output_dir / f"{job_id}_results.csv" | |
| df.to_csv(csv_path, index=False) | |
| job_summary = { | |
| "job_id": job_id, | |
| "status": status, | |
| "mode": mode, | |
| "created_at": created_at, | |
| "preset": "egfr_generator", | |
| "docking_mode": docking_mode, | |
| "num_requested": num_molecules, | |
| "num_returned": int(len(df)), | |
| "download_csv_file": str(csv_path), | |
| "live_generation": generation_backend == "live_reinvent_cpu", | |
| } | |
| response = { | |
| "job_id": job_id, | |
| "status": status, | |
| "mode": mode, | |
| "created_at": created_at, | |
| "preset": "egfr_generator", | |
| "docking_mode": docking_mode, | |
| "summary": { | |
| "num_requested": num_molecules, | |
| "num_returned": int(len(df)), | |
| }, | |
| "columns": list(df.columns), | |
| "results": df.to_dict(orient="records"), | |
| "download_csv_file": str(csv_path), | |
| "warnings": warnings, | |
| } | |
| return job_summary, df, response, str(csv_path) | |
| loaded_ok, loaded_msg = download_bundle_light() | |
| with gr.Blocks(title="EGFR AI Demo") as demo: | |
| gr.Markdown( | |
| """ | |
| # EGFR AI Demo | |
| Private CPU demo for the EGFR AI backend workflow. | |
| This demo shows the backend response shape for generation, scoring, filtering, and docking integration. | |
| Real AutoDock-GPU docking is disabled in this CPU Space, but docking fields are present so backend/frontend teams can implement database columns and filters. | |
| """ | |
| ) | |
| gr.Markdown(f"Bundle load status: {loaded_msg}") | |
| with gr.Tab("Score SMILES"): | |
| gr.Markdown( | |
| """ | |
| Enter one or more SMILES. | |
| This CPU demo calculates molecular properties and returns the same response columns expected by the backend. | |
| Important: | |
| - pred_pAff_mean is an integration field in this Space version. | |
| - real docking is disabled in this CPU Space. | |
| - docking_status is returned as not_run. | |
| """ | |
| ) | |
| smiles_input = gr.Textbox( | |
| label="SMILES input", | |
| value="CCO\nc1ccccc1", | |
| lines=5, | |
| ) | |
| score_btn = gr.Button("Score SMILES") | |
| score_table = gr.Dataframe(label="Results table") | |
| score_json = gr.JSON(label="Backend-style JSON response") | |
| score_btn.click( | |
| score_smiles, | |
| inputs=smiles_input, | |
| outputs=[score_table, score_json], | |
| ) | |
| with gr.Tab("Generate Job Demo"): | |
| gr.Markdown( | |
| """ | |
| This tab demonstrates generation job behavior. | |
| Generation backend options: | |
| - demo_rows: uses packaged example rows. Fast and stable. | |
| - live_reinvent_cpu: tries to run real REINVENT sampling from the packaged checkpoint on CPU. | |
| The output is split into: | |
| 1. Job summary: job_id, status, preset, mode, and metadata. | |
| 2. Generated molecule results table. | |
| 3. Full backend-style generation JSON. | |
| 4. Download CSV. | |
| Notes: | |
| - DeepPurpose live scoring is not enabled yet in this Space version. | |
| - Real docking is disabled in this CPU Space, so docking_status is returned as not_run. | |
| """ | |
| ) | |
| generation_backend = gr.Dropdown( | |
| label="generation_backend", | |
| choices=["demo_rows", "live_reinvent_cpu"], | |
| value="demo_rows", | |
| ) | |
| gen_num_molecules = gr.Number( | |
| label="num_molecules", | |
| value=20, | |
| precision=0, | |
| ) | |
| gen_return_top_k = gr.Number( | |
| label="return_top_k", | |
| value=20, | |
| precision=0, | |
| ) | |
| gen_docking_mode = gr.Dropdown( | |
| label="docking_mode", | |
| choices=["off", "top_k", "all"], | |
| value="off", | |
| ) | |
| gen_btn = gr.Button("Run Generation Job") | |
| gen_summary = gr.JSON(label="Job summary") | |
| gen_table = gr.Dataframe(label="Generated molecule results table") | |
| gen_json = gr.JSON(label="Full backend-style generation JSON") | |
| gen_file = gr.File(label="Download CSV") | |
| gen_btn.click( | |
| generate_job_demo, | |
| inputs=[gen_num_molecules, gen_return_top_k, gen_docking_mode, generation_backend], | |
| outputs=[gen_summary, gen_table, gen_json, gen_file], | |
| ) | |
| with gr.Tab("Example Generate Response"): | |
| gr.Markdown( | |
| """ | |
| This example is loaded from the private Hugging Face Dataset bundle. | |
| It shows the expected response shape after: | |
| Generate -> Enrich -> Optional Docking | |
| """ | |
| ) | |
| example_btn = gr.Button("Load example response") | |
| example_table = gr.Dataframe(label="Example results table") | |
| example_json = gr.JSON(label="Example JSON response") | |
| example_btn.click( | |
| lambda: (get_example_table(), get_example_json()), | |
| outputs=[example_table, example_json], | |
| ) | |
| with gr.Tab("Docking Contract"): | |
| gr.Markdown( | |
| """ | |
| ## Docking modes | |
| The full backend supports: | |
| - off | |
| - top_k | |
| - all | |
| In this CPU Space, real AutoDock-GPU docking is disabled. | |
| Returned docking fields in CPU demo: | |
| - docking_score: null | |
| - docking_status: not_run | |
| - docking_pose_file: null | |
| Recommended database columns: | |
| - docking_score FLOAT NULL | |
| - docking_status VARCHAR(30) DEFAULT 'not_run' | |
| - docking_pose_file TEXT NULL | |
| - docking_mode VARCHAR(20) DEFAULT 'off' | |
| - docking_error TEXT NULL | |
| Recommended frontend and filter behavior: | |
| - Show all valid molecules: valid == true | |
| - Show docked only: docking_status == completed | |
| - Show not-yet-docked molecules: docking_status == not_run | |
| - Sort docking results ascending, because more negative docking scores are usually better | |
| """ | |
| ) | |
| contract_btn = gr.Button("Load backend contract") | |
| contract_json = gr.JSON(label="backend_contract.json") | |
| contract_btn.click( | |
| get_contract, | |
| outputs=contract_json, | |
| ) | |
| with gr.Tab("API Spec"): | |
| gr.Markdown( | |
| """ | |
| This tab loads the API response specification from the private artifact bundle. | |
| It explains the expected backend/frontend response format for: | |
| - generated molecules | |
| - scored SMILES | |
| - docking fields | |
| - frontend filters | |
| """ | |
| ) | |
| spec_btn = gr.Button("Load API response spec") | |
| spec_output = gr.Markdown(label="API_RESPONSE_SPEC.md") | |
| spec_btn.click( | |
| get_api_spec, | |
| outputs=spec_output, | |
| ) | |
| demo.launch() | |