Spaces:
Runtime error
Runtime error
File size: 4,891 Bytes
31e8a9f | 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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """Model artifact registry and HuggingFace Hub download (PRD §17.3).
Production policy (user decision, see DECISIONS.md): the worker **hard-fails**
if any artifact is missing. There is no heuristic fallback.
CLI (from ``backend/``)::
python -m app.ml.models.loader # download missing artifacts, then verify
python -m app.ml.models.loader --check # verify only, no network
"""
from __future__ import annotations
import argparse
import logging
import sys
from pathlib import Path
from typing import Final
logger = logging.getLogger(__name__)
#: Every file the production pipeline needs, with the notebook that produces it.
REQUIRED_ARTIFACTS: Final[dict[str, str]] = {
"experience_model.onnx": "notebooks/04_finetune_experience.ipynb (INT8 ONNX)",
"experience_tokenizer.json": "notebooks/04_finetune_experience.ipynb",
"absa_model.onnx": "notebooks/04_finetune_experience.ipynb (ABSA export section)",
"absa_tokenizer.json": "notebooks/04_finetune_experience.ipynb (ABSA export section)",
"absa_config.json": "notebooks/04_finetune_experience.ipynb (ABSA export section)",
"aspect_detector.joblib": "training/scripts/train_all.py --stages aspects",
"domain_tfidf.joblib": "notebooks/05_train_fusion.ipynb (domain model section)",
"fusion_ensemble.json": "notebooks/05_train_fusion.ipynb",
"conformal_qhat.json": "notebooks/06_calibration.ipynb",
}
class ModelArtifactsMissingError(RuntimeError):
"""Raised when required model artifacts are absent after any download attempt."""
def __init__(self, model_dir: Path, missing: list[str]) -> None:
self.model_dir = model_dir
self.missing = missing
details = "\n".join(
f" - {name} <- produced by {REQUIRED_ARTIFACTS[name]}" for name in missing
)
super().__init__(
f"{len(missing)} required model artifact(s) missing in {model_dir}:\n{details}\n"
"Run notebooks 04 (Colab) and 05-06 (local), then upload the files to "
"HF_MODEL_REPO or copy them into MODEL_DIR."
)
def missing_artifacts(model_dir: Path) -> list[str]:
"""Return the names of required artifacts not present (or empty) in ``model_dir``."""
return [
name
for name in REQUIRED_ARTIFACTS
if not (model_dir / name).is_file() or (model_dir / name).stat().st_size == 0
]
def artifact_status(model_dir: Path) -> dict[str, bool]:
"""Map each required artifact to whether it is present. Used by ``/health``."""
missing = set(missing_artifacts(model_dir))
return {name: name not in missing for name in REQUIRED_ARTIFACTS}
def ensure_models(model_dir: Path, hf_repo: str = "", hf_token: str = "") -> dict[str, Path]:
"""Download missing artifacts from the Hub (if configured) and verify all exist.
Args:
model_dir: Local artifact directory (created if absent).
hf_repo: ``user/repo`` on HuggingFace Hub; empty disables downloading.
hf_token: Optional token for private repos.
Returns:
Mapping of artifact name to its local path.
Raises:
ModelArtifactsMissingError: If anything is still missing afterwards.
"""
model_dir.mkdir(parents=True, exist_ok=True)
missing = missing_artifacts(model_dir)
if missing and hf_repo:
from huggingface_hub import hf_hub_download # imported lazily: network-only path
for name in missing:
logger.info("Downloading %s from %s ...", name, hf_repo)
try:
hf_hub_download(
repo_id=hf_repo,
filename=name,
local_dir=model_dir,
token=hf_token or None,
)
except Exception as exc: # report every failure, then fail once below
logger.error("Failed to download %s: %s", name, exc)
missing = missing_artifacts(model_dir)
if missing:
raise ModelArtifactsMissingError(model_dir, missing)
logger.info("All %d model artifacts ready in %s", len(REQUIRED_ARTIFACTS), model_dir)
return {name: model_dir / name for name in REQUIRED_ARTIFACTS}
def _main(argv: list[str] | None = None) -> int:
from app.config import settings
parser = argparse.ArgumentParser(description="Download / verify VERA model artifacts.")
parser.add_argument("--check", action="store_true", help="verify only; never download")
args = parser.parse_args(argv)
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
try:
ensure_models(
settings.MODEL_DIR,
hf_repo="" if args.check else settings.HF_MODEL_REPO,
hf_token=settings.HF_TOKEN,
)
except ModelArtifactsMissingError as exc:
logger.error("%s", exc)
return 1
return 0
if __name__ == "__main__":
sys.exit(_main())
|