medrax/config.py

import os
from pathlib import Path
from typing import Optional


def _resolve_default_cache_path() -> Path:
    """Pick the best-available cache directory."""
    candidates = (
        os.environ.get("MEDRAX_MODEL_DIR"),
        os.environ.get("HF_HOME"),
        os.environ.get("TRANSFORMERS_CACHE"),
    )
    for candidate in candidates:
        if candidate:
            return Path(candidate).expanduser()
    return Path.home() / ".cache" / "medrax_hf"


def _ensure_subdirs(base: Path) -> None:
    """Create the common Hugging Face cache subdirectories."""
    for subdir in ("datasets", "hub"):
        (base / subdir).mkdir(parents=True, exist_ok=True)


_DEFAULT_CACHE_PATH = _resolve_default_cache_path()
_DEFAULT_CACHE_PATH.mkdir(parents=True, exist_ok=True)
_ensure_subdirs(_DEFAULT_CACHE_PATH)

# Expose a string form for convenience
DEFAULT_CACHE_DIR = str(_DEFAULT_CACHE_PATH)


def _ensure_hf_env_vars(base: Path) -> None:
    """Populate the typical HF cache env vars if the user did not set them."""
    os.environ.setdefault("MEDRAX_MODEL_DIR", str(base))
    os.environ.setdefault("HF_HOME", str(base))
    os.environ.setdefault("TRANSFORMERS_CACHE", str(base))
    os.environ.setdefault("HF_DATASETS_CACHE", str(base / "datasets"))
    os.environ.setdefault("HF_HUB_CACHE", str(base / "hub"))


_ensure_hf_env_vars(_DEFAULT_CACHE_PATH)


def resolve_cache_dir(cache_dir: Optional[str] = None) -> str:
    """Expand, create, and return a usable cache directory."""
    if cache_dir:
        path = Path(cache_dir).expanduser()
        path.mkdir(parents=True, exist_ok=True)
        _ensure_subdirs(path)
        return str(path)
    return DEFAULT_CACHE_DIR