sitecustomize.py

import os


def _install_hf_folder_shim() -> None:
    try:
        import huggingface_hub
    except Exception:
        return

    if hasattr(huggingface_hub, "HfFolder"):
        return

    class HfFolder:
        @staticmethod
        def get_token() -> str | None:
            return os.getenv("HUGGINGFACE_HUB_TOKEN") or os.getenv("HF_TOKEN")

        @staticmethod
        def save_token(token: str) -> None:
            os.environ["HUGGINGFACE_HUB_TOKEN"] = token

        @staticmethod
        def delete_token() -> None:
            os.environ.pop("HUGGINGFACE_HUB_TOKEN", None)

    huggingface_hub.HfFolder = HfFolder


_install_hf_folder_shim()