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), ('1K<n<10K', 1_000, 10_000), ('10K<n<100K', 10_000, 100_000), ('100K<n<1M', 100_000, 1_000_000), ('1M<n<10M', 1_000_000, 10_000_000), ('10M<n<100M', 10_000_000, 100_000_000), ('100M<n<1B', 100_000_000, 1_000_000_000), ('1B<n<10B', 1_000_000_000, 10_000_000_000), ('10B<n<100B', 10_000_000_000, 100_000_000_000), ('100B<n<1T', 100_000_000_000, 1_000_000_000_000), ('n>1T', 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'))