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()