Upload folder using huggingface_hub

app.py

@@ -1,5 +1,6 @@
 import json
 import os
+import subprocess
 from pathlib import Path
 import uuid
 from datetime import datetime
@@ -18,10 +19,11 @@ 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
+    global EXAMPLE_JSON, API_SPEC_TEXT, BACKEND_CONTRACT_TEXT, BUNDLE_ROOT
 
     try:
         local_dir = snapshot_download(
@@ -32,12 +34,15 @@ def download_bundle_light():
                 "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"
@@ -187,92 +192,239 @@ def get_contract():
         "message": "backend_contract.json not loaded. Check HF_TOKEN secret or BUNDLE_REPO variable.",
     }
 
-def generate_job_demo(num_molecules, return_top_k, docking_mode):
-    """
-    CPU demo generation endpoint.
+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.")
 
-    This does not run live REINVENT generation.
-    It returns a backend-style generation job response using packaged example rows,
-    so frontend/filter/backend teams can test JSON + CSV + job history behavior.
-    """
-    job_id = f"gen_{uuid.uuid4().hex[:8]}"
-    created_at = datetime.utcnow().isoformat() + "Z"
+    checkpoint = Path(BUNDLE_ROOT) / "models" / "generator" / "egfr_generator.chkpt"
+    if not checkpoint.exists():
+        raise FileNotFoundError(f"Generator checkpoint not found: {checkpoint}")
 
-    if EXAMPLE_JSON and "results" in EXAMPLE_JSON:
-        rows = EXAMPLE_JSON["results"]
-    else:
-        rows = [
-            compute_row("CCO"),
-            compute_row("c1ccccc1"),
-        ]
+    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]
 
-    num_molecules = int(num_molecules)
     return_top_k = int(return_top_k)
+    if return_top_k > 0:
+        df = df.head(return_top_k)
+
+    return df
 
-    selected = rows[: max(1, min(return_top_k, len(rows)))]
 
-    clean_rows = []
-    for row in selected:
-        clean = dict(row)
+def generate_job_demo(num_molecules, return_top_k, docking_mode, generation_backend):
+    """
+    Generation endpoint for the Space.
 
-        clean["docking_score"] = None
-        clean["docking_status"] = "not_run"
-        clean["docking_pose_file"] = None
+    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)
 
-        if "pred_pAff_mean" not in clean:
-            clean["pred_pAff_mean"] = None
+    warnings = []
 
-        clean_rows.append(clean)
+    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.",
+            ])
 
-    df = pd.DataFrame(clean_rows)
+    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": "completed",
-        "mode": "generate_job_cpu_demo",
+        "status": status,
+        "mode": mode,
         "created_at": created_at,
         "preset": "egfr_generator",
         "docking_mode": docking_mode,
         "num_requested": num_molecules,
-        "num_generated": len(rows),
-        "num_returned": len(clean_rows),
+        "num_returned": int(len(df)),
         "download_csv_file": str(csv_path),
-        "live_generation": False,
-        "note": "CPU demo uses packaged example rows. Live REINVENT generation is not executed in this Space version."
+        "live_generation": generation_backend == "live_reinvent_cpu",
     }
 
     response = {
         "job_id": job_id,
-        "status": "completed",
-        "mode": "generate_job_cpu_demo",
+        "status": status,
+        "mode": mode,
         "created_at": created_at,
         "preset": "egfr_generator",
         "docking_mode": docking_mode,
         "summary": {
             "num_requested": num_molecules,
-            "num_generated": len(rows),
-            "num_returned": len(clean_rows),
+            "num_returned": int(len(df)),
         },
         "columns": list(df.columns),
-        "results": clean_rows,
+        "results": df.to_dict(orient="records"),
         "download_csv_file": str(csv_path),
-        "warnings": [
-            "This is a CPU demo generation response.",
-            "Live REINVENT generation is not executed in this Space version.",
-            "The response shape matches the planned backend contract.",
-            "Real docking is not run in this CPU Space; docking_status is not_run.",
-        ],
+        "warnings": warnings,
     }
 
     return job_summary, df, response, str(csv_path)
 
 
+
+
 loaded_ok, loaded_msg = download_bundle_light()
 
 
@@ -323,28 +475,35 @@ Important:
     with gr.Tab("Generate Job Demo"):
         gr.Markdown(
             """
-This tab demonstrates the expected backend behavior for generation jobs.
+This tab demonstrates generation job behavior.
 
-Important:
-This CPU Space does not run live REINVENT generation. It uses packaged example rows to let frontend, filters, and backend test the full response format.
+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 4 parts:
+The output is split into:
 
-1. Job summary: job-level metadata such as job_id, status, preset, and summary.
-2. Generated results table: molecule-level rows only.
-3. Backend-style generation JSON: full response containing job metadata + results.
-4. Download CSV: downloadable file for the same molecule-level results.
+1. Job summary: job_id, status, preset, mode, and metadata.
+2. Generated molecule results table.
+3. Full backend-style generation JSON.
+4. Download CSV.
 
-In the real backend:
-- The frontend should consume the JSON.
-- The user can optionally download the CSV.
-- The backend should store job history using job_id/status/created_at.
+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=100,
+            value=20,
             precision=0,
         )
 
@@ -360,7 +519,7 @@ In the real backend:
             value="off",
         )
 
-        gen_btn = gr.Button("Generate Job Demo")
+        gen_btn = gr.Button("Run Generation Job")
 
         gen_summary = gr.JSON(label="Job summary")
         gen_table = gr.Dataframe(label="Generated molecule results table")
@@ -369,7 +528,7 @@ In the real backend:
 
         gen_btn.click(
             generate_job_demo,
-            inputs=[gen_num_molecules, gen_return_top_k, gen_docking_mode],
+            inputs=[gen_num_molecules, gen_return_top_k, gen_docking_mode, generation_backend],
             outputs=[gen_summary, gen_table, gen_json, gen_file],
         )
 

requirements.txt

@@ -1,3 +1,5 @@
 pandas
 huggingface_hub
 rdkit
+toml
+git+https://github.com/MolecularAI/REINVENT4.git