Spaces:
Paused
Paused
File size: 3,080 Bytes
a7f5b86 5514429 a7f5b86 5514429 a7f5b86 5514429 a7f5b86 | 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 | """
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.")
|