Spaces:
Runtime error
Runtime error
| """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()) | |