File size: 4,485 Bytes
8bd2e63 d9a68e9 8bd2e63 e4298fb 8bd2e63 e4298fb 8bd2e63 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """Pre-generate example cache files for judge-proofing.
Run this script once on a machine with GPU access (HF Space or local with CUDA)
AFTER the Space has warmed up so the model is resident. The output JSON files
are written to data/example_cache/<key>.json and are committed to the repo so
that example-button clicks are instant even on a cold ZeroGPU worker.
Usage (from the project root):
python scripts/cache_examples.py
Each file contains a serialised HandoffOutput (Pydantic .model_dump()). The app
loads these at click time via app._load_cached_example() and renders them through
the normal render_handoff_html / render_safety_html path — validated, not raw.
NOTE: Do NOT run this script inside a ZeroGPU Space (SPACES_ZERO_GPU is set).
ZeroGPU only allocates GPU within a live request context (@spaces.GPU window),
so model.generate() calls made at the top level will run on CPU and be very slow
or time out. Run on a regular GPU box (A10/A100 Space, local CUDA machine, or
a standard HF Space with a persistent GPU).
"""
from __future__ import annotations
import json
import os
import sys
from pathlib import Path
# Make sure the project root is on the path when run as a script.
sys.path.insert(0, str(Path(__file__).parent.parent))
# Parse this before importing app.py. Otherwise --no-model still triggers the
# import-time 4B model load, defeating the purpose of the flag.
USE_MODEL = "--no-model" not in sys.argv
if not USE_MODEL:
os.environ["DENTAL_SOAP_USE_MODEL"] = "0"
from app import build_outputs, EXAMPLES # noqa: E402 (after sys.path fix)
from examples import STEP2_CHECKS, STEP3_CHECKS, STEP4_CHECKS # noqa: E402
from schema import HandoffOutput # noqa: E402
CACHE_DIR = Path(__file__).parent.parent / "data" / "example_cache"
CACHE_DIR.mkdir(parents=True, exist_ok=True)
def _split_checks(checks: list[str]):
return (
[c for c in checks if c in STEP2_CHECKS],
[c for c in checks if c in STEP3_CHECKS],
[c for c in checks if c in STEP4_CHECKS],
)
def cache_example(key: str, use_model: bool = True) -> None:
ex = EXAMPLES[key]
checks_dental, checks_jaw, checks_body = _split_checks(ex["checks"])
print(f"Generating cache for '{key}' (use_model={use_model}) ...")
_safety, _handoff, raw_json, status, _pdf, _email, _workflow = build_outputs(
ex["story"],
ex.get("name", ""),
ex["age"],
ex["language"],
ex["chief_concern"],
ex["duration"],
ex["tooth_or_area"],
ex["recent_dental_work"],
ex["pain_score"],
ex["meds"],
ex["allergies"],
ex["goals"],
checks_dental,
checks_jaw,
checks_body,
use_model,
)
print(f" status: {status}")
if use_model and not status.startswith("AI model path used (ZeroGPU):"):
raise RuntimeError(
f"Refusing to cache '{key}': model-backed generation was requested "
f"but the app returned: {status}"
)
output = HandoffOutput.model_validate_json(raw_json)
out_path = CACHE_DIR / f"{key}.json"
tmp_path = out_path.with_suffix(".json.tmp")
tmp_path.write_text(
json.dumps(output.model_dump(), ensure_ascii=False, indent=2),
encoding="utf-8",
)
tmp_path.replace(out_path)
print(f" written: {out_path}")
if __name__ == "__main__":
# Guard: ZeroGPU only allocates GPU inside a request context, so running the
# model generate() calls at script top-level will silently fall back to CPU
# and produce unusably slow or broken output. Bail out early with a clear
# message so the operator knows to re-run on a non-ZeroGPU GPU machine.
if os.getenv("SPACE_ID") and os.getenv("SPACES_ZERO_GPU"):
sys.exit(
"[cache_examples] ERROR: this script is running inside a ZeroGPU Space "
"(SPACE_ID and SPACES_ZERO_GPU are both set). ZeroGPU only allocates GPU "
"within a live @spaces.GPU request window — model calls here run on CPU "
"and will be extremely slow or time out. Re-run on a regular GPU box "
"(a standard HF Space with a persistent GPU, a local CUDA machine, or an "
"A10/A100 HF Space)."
)
# Pass --no-model to generate from the deterministic template path only.
for example_key in EXAMPLES:
cache_example(example_key, use_model=USE_MODEL)
print("Done. Commit data/example_cache/*.json to the repo.")
|