social-data-attribution-w2/src/data_attribution/config.py

"""Project configuration helpers."""

from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
DEFAULT_HF_CACHE_ROOT = PROJECT_ROOT / ".hf_cache"
ALLOWED_MODEL_IDS = (
    "allenai/Olmo-3-1025-7B",
    "allenai/Olmo-3-1125-32B",
)
DEFAULT_MODEL_ID = ALLOWED_MODEL_IDS[0]
DEPRECATED_TRANSFORMERS_CACHE_VARS = (
    "PYTORCH_PRETRAINED_BERT_CACHE",
    "PYTORCH_TRANSFORMERS_CACHE",
    "TRANSFORMERS_CACHE",
)


@dataclass(frozen=True)
class HFCachePaths:
    root: Path
    hub: Path
    datasets: Path
    transformers: Path


def resolve_hf_cache_root(cache_root: Path | None = None) -> Path:
    if cache_root is not None:
        return Path(cache_root)
    env_root = os.environ.get("HF_CACHE_ROOT")
    if env_root:
        return Path(env_root)
    return DEFAULT_HF_CACHE_ROOT


def resolve_hf_datasets_cache(cache_dir: Path | None = None) -> Path:
    if cache_dir is not None:
        return Path(cache_dir)
    env_cache = os.environ.get("HF_DATASETS_CACHE")
    if env_cache:
        return Path(env_cache)
    return Path.home() / ".cache" / "huggingface" / "datasets"


def configure_hf_cache(
    cache_root: Path | None = None, *, create_dirs: bool = False
) -> HFCachePaths:
    root = resolve_hf_cache_root(cache_root).expanduser()
    paths = HFCachePaths(
        root=root,
        hub=root / "hub",
        datasets=root / "datasets",
        transformers=root / "transformers",
    )

    os.environ["HF_HOME"] = str(paths.root)
    os.environ["HUGGINGFACE_HUB_CACHE"] = str(paths.hub)
    os.environ["HF_DATASETS_CACHE"] = str(paths.datasets)
    for env_var in DEPRECATED_TRANSFORMERS_CACHE_VARS:
        os.environ.pop(env_var, None)

    if create_dirs:
        for path in (paths.root, paths.hub, paths.datasets, paths.transformers):
            path.mkdir(parents=True, exist_ok=True)

    return paths


__all__ = [
    "DEFAULT_HF_CACHE_ROOT",
    "DEFAULT_MODEL_ID",
    "ALLOWED_MODEL_IDS",
    "DEPRECATED_TRANSFORMERS_CACHE_VARS",
    "PROJECT_ROOT",
    "HFCachePaths",
    "configure_hf_cache",
    "resolve_hf_cache_root",
    "resolve_hf_datasets_cache",
]