import math import os from functools import partial from huggingface_hub import configure_http_backend, HfApi, HfFileSystem from .session import get_requests_session _NUM_TAGS = [ ('n<1K', 0, 1_000), ('1K1T', 1_000_000_000_000, math.inf), ] def number_to_tag(v): for tag, min_, max_ in _NUM_TAGS: if min_ <= v < max_: return tag raise ValueError(f'No tags found for {v!r}') configure_http_backend(partial(get_requests_session, timeout=120)) def get_hf_client() -> HfApi: return HfApi(token=os.environ.get('HF_TOKEN')) def get_hf_fs() -> HfFileSystem: return HfFileSystem(token=os.environ.get('HF_TOKEN'))