| |
| |
| """The model's files from the Hugging Face Hub, checked against this package's own list. |
| |
| from anuvani_diar.hub import bundle_dir, golden_audio_dir |
| bundle_dir() # model/ of CIIL/Anuvani-Diarization, downloaded once into the Hugging Face cache |
| golden_audio_dir() # golden/audio/: the 12 recordings anuvani-diar-check diarises |
| |
| Every file's sha256 must equal the one in `assets/bundle_files.json` (or `assets/golden12.json` for the recordings), |
| which ship inside the package, so the weights a Diarizer loads are the ones this package was tested with whatever |
| revision the Hub serves. Downloads come from the release tag this package was made for (`REVISION`), or from the |
| main branch if that tag is missing; `ANUVANI_DIAR_REVISION` (or `revision=`) names another branch, tag or commit. No |
| Hugging Face account is needed; huggingface_hub's own settings apply (`HF_HUB_CACHE`, `HF_HUB_OFFLINE`, |
| `HF_ENDPOINT`). With `model/` comes the repository's `config.json` (the release's description): the Hub counts |
| requests for it as the model's downloads; nothing else is sent. |
| """ |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| from typing import Dict, Optional |
|
|
| REPO = "CIIL/Anuvani-Diarization" |
| |
| REVISION = "v1.0.0" |
| _HERE = Path(__file__).resolve().parent |
| _CHECKED: Dict[tuple, str] = {} |
|
|
|
|
| def assets_dir() -> Path: |
| """The package's assets: `$ANUVANI_DIAR_ASSETS` (set by the package to its own copy), else the source tree's.""" |
| env = os.environ.get("ANUVANI_DIAR_ASSETS") |
| if env: |
| return Path(env) |
| return _HERE / "assets" if (_HERE / "assets").is_dir() else _HERE.parents[1] / "assets" |
|
|
|
|
| def sha256(path: Path) -> str: |
| """The file's sha256, computed once per process for each (file, size, modification time).""" |
| st = path.stat() |
| key = (str(path.resolve()), st.st_size, st.st_mtime_ns) |
| if key not in _CHECKED: |
| h = hashlib.sha256() |
| with open(path, "rb") as f: |
| for block in iter(lambda: f.read(1 << 24), b""): |
| h.update(block) |
| _CHECKED[key] = h.hexdigest() |
| return _CHECKED[key] |
|
|
|
|
| def _download(patterns, revision: Optional[str], token) -> Path: |
| try: |
| from huggingface_hub import snapshot_download |
| except ImportError as e: |
| raise ImportError("the model's files come from the Hugging Face Hub: pip install huggingface_hub, or pass " |
| "bundle= (or set ANUVANI_DIAR_BUNDLE) to a downloaded copy") from e |
| chosen = revision or os.environ.get("ANUVANI_DIAR_REVISION") |
| try: |
| try: |
| return Path(snapshot_download(REPO, revision=chosen or REVISION, allow_patterns=patterns, token=token)) |
| |
| except Exception as e: |
| if chosen or type(e).__name__ != "RevisionNotFoundError": |
| raise |
| return Path(snapshot_download(REPO, revision=None, allow_patterns=patterns, token=token)) |
| except Exception as e: |
| raise RuntimeError(f"could not download {patterns[0].split('/')[0]}/ from {REPO} ({type(e).__name__}: {e}). " |
| "Check the network connection; without one, download the repository elsewhere " |
| f"(hf download {REPO}) and set ANUVANI_DIAR_BUNDLE to its model/ folder.") from e |
|
|
|
|
| def _check(root: Path, want: Dict[str, str], what: str) -> None: |
| bad = [rel for rel, h in want.items() if not (root / rel).is_file() or sha256(root / rel) != h] |
| if bad: |
| raise RuntimeError(f"{what}: {', '.join(bad)} under {root} do not match this package's sha256 list " |
| "(another revision of the repository?). Pin the revision this package was made for with " |
| "ANUVANI_DIAR_REVISION, or update the package.") |
|
|
|
|
| def bundle_dir(revision: Optional[str] = None, token=None) -> Path: |
| """`model/` of the Hub repository (MANIFEST, decoder, both networks' safetensors), every file sha256-checked.""" |
| spec = json.loads((assets_dir() / "bundle_files.json").read_text()) |
| sub, want = spec["subdir"], spec["files"] |
| |
| root = _download([f"{sub}/{rel}" for rel in sorted(want)] + ["config.json"], revision, token) / sub |
| _check(root, want, f"{REPO}/{sub}") |
| return root |
|
|
|
|
| def golden_audio_dir(revision: Optional[str] = None, token=None) -> Path: |
| """`golden/audio/` of the Hub repository: the release's 12 install-check recordings, sha256-checked.""" |
| files = json.loads((assets_dir() / "golden12.json").read_text())["files"] |
| want = {f"{sid}.flac": v["flac_sha256"] for sid, v in files.items()} |
| root = _download([f"golden/audio/{f}" for f in sorted(want)], revision, token) / "golden" / "audio" |
| _check(root, want, f"{REPO}/golden/audio") |
| return root |
|
|
