File size: 1,635 Bytes
d5147ba | 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 | """HuggingFace cache location for evaluation (EVAL_HF_ROOT, default eval/hf_cache).
EVAL_OFFLINE=1 reads only what is already cached there; the default downloads on first use.
"""
from __future__ import annotations
import os
from pathlib import Path
from datasets import DownloadMode
from datasets import load_dataset as _load_dataset
_EVAL_DIR = Path(__file__).resolve().parent
DEFAULT_HF_ROOT = Path(os.environ.get("EVAL_HF_ROOT", str(_EVAL_DIR / "hf_cache")))
def is_offline_eval() -> bool:
return os.environ.get("EVAL_OFFLINE", "0") != "0"
def enable_offline_mode(offline: bool | None = None) -> None:
"""Skip HuggingFace Hub network checks; use only local cache."""
if offline is None:
offline = is_offline_eval()
if offline:
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["HF_DATASETS_OFFLINE"] = "1"
def setup_hf_datasets_cache(hf_root: str | Path | None = None) -> dict[str, str]:
"""Point HF hub/datasets caches at a shared directory (not ~/.cache)."""
root = Path(hf_root or DEFAULT_HF_ROOT)
paths = {
"HF_HOME": str(root / "hf_home"),
"HF_DATASETS_CACHE": str(root / "hf_datasets"),
"HUGGINGFACE_HUB_CACHE": str(root / "hf_hub"),
}
for key, value in paths.items():
os.environ[key] = value
Path(value).mkdir(parents=True, exist_ok=True)
enable_offline_mode()
return paths
def load_dataset(*args, **kwargs):
kwargs.setdefault("trust_remote_code", True)
if is_offline_eval():
kwargs.setdefault("download_mode", DownloadMode.REUSE_DATASET_IF_EXISTS)
return _load_dataset(*args, **kwargs)
|