Spaces:
Paused
Paused
| """ | |
| Download inference artifacts from HuggingFace Hub at startup if not available locally. | |
| """ | |
| import os | |
| import shutil | |
| from pathlib import Path | |
| DATA_PROCESSED = Path("data/processed") | |
| MODELS_DIR = Path("models") | |
| DEFAULT_REPO_ID = os.environ.get( | |
| "CROSSPAPER_ARTIFACTS_REPO", "zkmine/crosspaper-artifacts" | |
| ) | |
| # Local paths that must exist before the app can serve recommendations | |
| REQUIRED_PATHS = [ | |
| DATA_PROCESSED / "base.index", | |
| DATA_PROCESSED / "finetuned.index", | |
| DATA_PROCESSED / "base_embeddings.npy", | |
| DATA_PROCESSED / "finetuned_embeddings.npy", | |
| DATA_PROCESSED / "paper_metadata.pkl", | |
| MODELS_DIR / "base" / "model.safetensors", | |
| MODELS_DIR / "fine_tuned" / "model.safetensors", | |
| ] | |
| # Mapping from the layout inside the Hub repo to the local layout | |
| HUB_LAYOUT = [ | |
| ("data/base.index", DATA_PROCESSED / "base.index"), | |
| ("data/finetuned.index", DATA_PROCESSED / "finetuned.index"), | |
| ("data/base_embeddings.npy", DATA_PROCESSED / "base_embeddings.npy"), | |
| ("data/finetuned_embeddings.npy", DATA_PROCESSED / "finetuned_embeddings.npy"), | |
| ("data/paper_metadata.pkl", DATA_PROCESSED / "paper_metadata.pkl"), | |
| ] | |
| HUB_MODEL_DIRS = [ | |
| ("models/base", MODELS_DIR / "base"), | |
| ("models/fine_tuned", MODELS_DIR / "fine_tuned"), | |
| ] | |
| def artifacts_present(): | |
| """Check if all required artifacts exist locally.""" | |
| return all(path.exists() for path in REQUIRED_PATHS) | |
| def ensure_artifacts(repo_id=DEFAULT_REPO_ID): | |
| """Download artifacts from Hub if not present, otherwise skip.""" | |
| if artifacts_present(): | |
| print("Artifacts found locally, skipping download.") | |
| return | |
| # Imported lazily so local runs do not require huggingface_hub | |
| from huggingface_hub import snapshot_download | |
| print(f"Artifacts not found locally. Downloading from {repo_id}...") | |
| print("This runs once per deployment and takes a few minutes.") | |
| snapshot_path = Path(snapshot_download( | |
| repo_id=repo_id, | |
| repo_type="dataset", | |
| )) | |
| DATA_PROCESSED.mkdir(parents=True, exist_ok=True) | |
| MODELS_DIR.mkdir(parents=True, exist_ok=True) | |
| for hub_relative, local_path in HUB_LAYOUT: | |
| source = snapshot_path / hub_relative | |
| if not source.exists(): | |
| continue | |
| local_path.parent.mkdir(parents=True, exist_ok=True) | |
| if not local_path.exists(): | |
| shutil.copy2(source, local_path) | |
| print(f" {local_path}") | |
| for hub_relative, local_dir in HUB_MODEL_DIRS: | |
| source = snapshot_path / hub_relative | |
| if not source.exists(): | |
| continue | |
| if not (local_dir / "model.safetensors").exists(): | |
| shutil.copytree(source, local_dir, dirs_exist_ok=True) | |
| print(f" {local_dir}/") | |
| missing = [str(p) for p in REQUIRED_PATHS if not p.exists()] | |
| if missing: | |
| raise RuntimeError( | |
| "Artifact download finished but these paths are still missing:\n " | |
| + "\n ".join(missing) | |
| + f"\nCheck the layout of {repo_id} on the Hub." | |
| ) | |
| print("All artifacts ready.") | |