Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """ | |
| Fetch the (private) eval data into a local dir before Streamlit starts serving. | |
| The Space repo is PUBLIC, so the eval data is deliberately NOT committed to it. | |
| Instead the data lives in a private HF dataset and is pulled at container | |
| startup using the `HF_TOKEN` Space secret. The download is targeted — only the | |
| two file types the dashboard reads, for the one subfolder — so it is a small | |
| (~tens of MB) fetch that completes in seconds, well before the first visitor. | |
| Because this runs in the Docker CMD *before* `streamlit run`, no visitor ever | |
| sees a partial download: by the time the app serves, DATA_DIR is fully on disk | |
| and `_resolve_eval_dir()` reads it as a local path (no network on the request | |
| path). On a cold restart this re-fetches, but the targeted snapshot is fast. | |
| Env / secrets: | |
| HF_TOKEN — required to read the private dataset | |
| HF_REPO_ID — dataset repo (default below) | |
| HF_DEFAULT_BRANCH — branch/revision to pull | |
| HF_DEFAULT_SUBFOLDER — subfolder within the branch to pull | |
| DATA_DIR — where to write the data (default /app/data) | |
| """ | |
| import os | |
| import shutil | |
| import sys | |
| def main() -> None: | |
| repo_id = os.environ.get("HF_REPO_ID", "srmdtranslations/model_gu_en_evaluations") | |
| branch = os.environ.get("HF_DEFAULT_BRANCH", "main") | |
| subfolder = os.environ.get("HF_DEFAULT_SUBFOLDER", "") | |
| token = os.environ.get("HF_TOKEN") or None | |
| data_dir = os.environ.get("DATA_DIR", "/app/data") | |
| if not token: | |
| print("[prewarm] HF_TOKEN not set — skipping fetch (app will try at runtime).", | |
| flush=True) | |
| return | |
| if not subfolder: | |
| print("[prewarm] HF_DEFAULT_SUBFOLDER not set — skipping fetch.", flush=True) | |
| return | |
| try: | |
| from huggingface_hub import snapshot_download | |
| print(f"[prewarm] fetching {repo_id}@{branch}/{subfolder} → {data_dir} …", | |
| flush=True) | |
| # Only the two file types the dashboard reads — keeps the fetch small/fast. | |
| local_root = snapshot_download( | |
| repo_id=repo_id, | |
| repo_type="dataset", | |
| revision=branch, | |
| token=token, | |
| allow_patterns=[ | |
| f"{subfolder}/**/*.holistic.jsonl", | |
| f"{subfolder}/**/*.intent_entity.jsonl", | |
| ], | |
| ) | |
| src = os.path.join(local_root, subfolder) | |
| if not os.path.isdir(src): | |
| print(f"[prewarm] subfolder missing after download: {src}", flush=True) | |
| return | |
| # Copy the matching files into DATA_DIR (resolve symlinks so the dir is | |
| # self-contained and readable without the HF cache layout). | |
| if os.path.exists(data_dir): | |
| shutil.rmtree(data_dir) | |
| n = 0 | |
| for root, _, files in os.walk(src): | |
| for fname in files: | |
| if not (fname.endswith(".holistic.jsonl") | |
| or fname.endswith(".intent_entity.jsonl")): | |
| continue | |
| rel = os.path.relpath(os.path.join(root, fname), src) | |
| out = os.path.join(data_dir, rel) | |
| os.makedirs(os.path.dirname(out), exist_ok=True) | |
| shutil.copy(os.path.realpath(os.path.join(root, fname)), out) | |
| n += 1 | |
| print(f"[prewarm] staged {n} files into {data_dir}.", flush=True) | |
| except Exception as exc: # noqa: BLE001 — prewarm must never block startup | |
| print(f"[prewarm] skipped ({exc})", flush=True) | |
| if __name__ == "__main__": | |
| main() | |