Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Hugging Face Space entry point for the cascade human review (private Space sociate/cascade-review). | |
| On start: downloads the web copy of the batch (images, product.json, review.html) from the private dataset BATCH_REPO and | |
| the saved verdicts (reviews.json) from VERDICT_REPO, then runs review_server on port 7860. The Space disk is wiped on | |
| restart, so after the workbook is rebuilt, reviews.json + cascade_report.xlsx are pushed back to VERDICT_REPO (at most | |
| one commit per PUSH_EVERY seconds, plus once on shutdown). Needs the HF_TOKEN secret (read on BATCH_REPO, write on VERDICT_REPO). | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import signal | |
| import sys | |
| import threading | |
| import time | |
| from functools import partial | |
| from http.server import ThreadingHTTPServer | |
| from pathlib import Path | |
| from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download, snapshot_download | |
| from huggingface_hub.errors import EntryNotFoundError | |
| import review_server | |
| BATCH_REPO = os.environ.get("BATCH_REPO", "sociate/cascade-review-batch") | |
| VERDICT_REPO = os.environ.get("VERDICT_REPO", "sociate/cascade-review-verdicts") | |
| PUSH_EVERY = float(os.environ.get("PUSH_EVERY", "30")) | |
| BATCH = Path(os.environ.get("BATCH_DIR", Path.home() / "data" / "batch")) | |
| api = HfApi() | |
| class Pusher: | |
| """Throttled push of reviews.json + cascade_report.xlsx to the verdict dataset.""" | |
| def __init__(self): | |
| self.files, self.dirty, self.lock = None, threading.Event(), threading.Lock() | |
| threading.Thread(target=self._loop, daemon=True).start() | |
| def __call__(self, reviews: Path, xlsx: Path): | |
| self.files = (reviews, xlsx) | |
| self.dirty.set() | |
| def push(self): | |
| with self.lock: | |
| if not self.files: | |
| return | |
| self.dirty.clear() | |
| reviews, xlsx = self.files | |
| try: | |
| api.create_commit(VERDICT_REPO, repo_type="dataset", commit_message="Update human review verdicts", | |
| operations=[CommitOperationAdd("reviews.json", reviews.read_bytes() if reviews.exists() else b"[]"), | |
| CommitOperationAdd("cascade_report.xlsx", xlsx.read_bytes())]) | |
| print(f"pushed verdicts to {VERDICT_REPO}", flush=True) | |
| except Exception as e: # retried on the next change / tick | |
| self.dirty.set() | |
| print("push failed:", e, flush=True) | |
| def _loop(self): | |
| while True: | |
| self.dirty.wait() | |
| self.push() | |
| time.sleep(PUSH_EVERY) | |
| def main() -> None: | |
| print(f"downloading {BATCH_REPO} ...", flush=True) | |
| snapshot_download(BATCH_REPO, repo_type="dataset", local_dir=BATCH, max_workers=16) | |
| try: | |
| hf_hub_download(VERDICT_REPO, "reviews.json", repo_type="dataset", local_dir=BATCH, force_download=True) | |
| print("restored reviews.json", flush=True) | |
| except EntryNotFoundError: | |
| print("no saved verdicts yet", flush=True) | |
| pusher = Pusher() | |
| review_server.Handler.store = store = review_server.Store(BATCH, debounce=2.0, after_build=pusher) | |
| store.dirty.set() # build the workbook once at start (with any restored verdicts) | |
| def stop(*_): | |
| if store.excel["pending"]: | |
| time.sleep(3) # let a rebuild in flight finish | |
| pusher.push() | |
| sys.exit(0) | |
| signal.signal(signal.SIGTERM, stop) | |
| srv = ThreadingHTTPServer(("0.0.0.0", 7860), partial(review_server.Handler, directory=str(BATCH))) | |
| print(f"serving {BATCH} ({len(store.reviews)} saved reviews) on :7860", flush=True) | |
| srv.serve_forever() | |
| if __name__ == "__main__": | |
| main() | |