| """ |
| Download model and data artifacts from HuggingFace Hub at startup. |
| |
| Run once before starting the API server. Files are cached locally so |
| subsequent restarts skip the download entirely. |
| |
| Usage: |
| python scripts/download_artifacts.py |
| |
| Environment variables: |
| HF_REPO_ID — HuggingFace model repo, e.g. "yashb18/deepreads-artifacts" |
| HF_TOKEN — Optional HF token (only needed if repo is private) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import os |
| import sys |
| from pathlib import Path |
|
|
| |
| |
| |
| |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
|
|
| ARTIFACT_MAP: list[tuple[str, Path]] = [ |
| |
| ("models/two_tower.pt", ROOT / "models" / "two_tower.pt"), |
| ("models/item_embeddings.npy", ROOT / "models" / "item_embeddings.npy"), |
| ("models/item_index.pkl", ROOT / "models" / "item_index.pkl"), |
| ("models/naive_baseline.pkl", ROOT / "models" / "naive_baseline.pkl"), |
| ("models/naive_category.pkl", ROOT / "models" / "naive_category.pkl"), |
| ("models/lgbm_model.pkl", ROOT / "models" / "lgbm_model.pkl"), |
| |
| ("data/processed/item_features.parquet", ROOT / "data" / "processed" / "item_features.parquet"), |
| ("data/processed/user_features.parquet", ROOT / "data" / "processed" / "user_features.parquet"), |
| ("data/processed/train.parquet", ROOT / "data" / "processed" / "train.parquet"), |
| ("data/processed/feature_artifacts.pkl", ROOT / "data" / "processed" / "feature_artifacts.pkl"), |
| ("data/processed/item_encoder.pkl", ROOT / "data" / "processed" / "item_encoder.pkl"), |
| ("data/processed/user_encoder.pkl", ROOT / "data" / "processed" / "user_encoder.pkl"), |
| ] |
|
|
|
|
| def main() -> None: |
| repo_id = os.environ.get("HF_REPO_ID", "") |
| if not repo_id: |
| print("ERROR: HF_REPO_ID environment variable is not set.", file=sys.stderr) |
| print(" Set it to your HuggingFace model repo, e.g.:", file=sys.stderr) |
| print(" export HF_REPO_ID=yourname/deepreads-artifacts", file=sys.stderr) |
| sys.exit(1) |
|
|
| token = os.environ.get("HF_TOKEN") or None |
|
|
| try: |
| from huggingface_hub import hf_hub_download |
| except ImportError: |
| print("ERROR: huggingface_hub is not installed. Run: pip install huggingface-hub", file=sys.stderr) |
| sys.exit(1) |
|
|
| missing = [(repo_fn, local) for repo_fn, local in ARTIFACT_MAP if not local.exists()] |
|
|
| if not missing: |
| print("All artifacts already present — skipping download.") |
| return |
|
|
| print(f"Downloading {len(missing)} artifact(s) from {repo_id} …") |
|
|
| for repo_filename, local_path in missing: |
| local_path.parent.mkdir(parents=True, exist_ok=True) |
| print(f" ↓ {repo_filename}") |
| try: |
| downloaded = hf_hub_download( |
| repo_id=repo_id, |
| filename=repo_filename, |
| repo_type="model", |
| token=token, |
| local_dir=ROOT, |
| local_dir_use_symlinks=False, |
| ) |
| |
| |
| _ = downloaded |
| except Exception as exc: |
| print(f"ERROR: Failed to download {repo_filename}: {exc}", file=sys.stderr) |
| sys.exit(1) |
|
|
| print("Download complete.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
