aiice.core.huggingface

  1from collections import defaultdict
  2from concurrent.futures import ThreadPoolExecutor, as_completed
  3from datetime import date, timedelta
  4from functools import lru_cache
  5from io import BytesIO
  6
  7import requests
  8from huggingface_hub import HfApi
  9from huggingface_hub.constants import DEFAULT_REQUEST_TIMEOUT
 10from huggingface_hub.errors import RemoteEntryNotFoundError
 11from huggingface_hub.file_download import http_get
 12from huggingface_hub.utils import build_hf_headers
 13
 14from aiice.constants import (
 15    BYTES_IN_MB,
 16    DATASET_SHAPE,
 17    DEFAULT_BACKOFF,
 18    DEFAULT_RETRIES,
 19    HF_BASE_URL,
 20    HF_DATASET_REPO,
 21    HF_PACKAGE_NAME,
 22    HF_REPO_TYPE,
 23    KEY_DATASET_END,
 24    KEY_DATASET_START,
 25    KEY_FILES,
 26    KEY_PER_YEAR,
 27    KEY_SHAPE,
 28    KEY_SIZE_BYTES,
 29    KEY_SIZE_MB,
 30    MAX_DATASET_END,
 31    MIN_DATASET_START,
 32    YEAR_STATS_CACHE_SIZE,
 33)
 34from aiice.core.utils import (
 35    convert_step_to_delta,
 36    get_filename_template,
 37    retry_on_network_errors,
 38)
 39
 40
 41class HfDatasetClient:
 42    """
 43    Client for accessing the AIICE Hugging Face dataset.
 44    """
 45
 46    def __init__(self):
 47        self._api_base_url = HF_BASE_URL
 48        self._api = HfApi(endpoint=self._api_base_url, library_name=HF_PACKAGE_NAME)
 49        self._api_headers = build_hf_headers(library_name=HF_PACKAGE_NAME)
 50
 51        self._dataset_repo = HF_DATASET_REPO
 52        self._dataset_repo_type = HF_REPO_TYPE
 53
 54        self._min_dataset_start, self._max_dataset_end = (
 55            MIN_DATASET_START,
 56            MAX_DATASET_END,
 57        )
 58        self._shape = DATASET_SHAPE
 59
 60    @property
 61    def dataset_start(self) -> date:
 62        """
 63        Earliest available date in the dataset.
 64        """
 65        return self._min_dataset_start
 66
 67    @property
 68    def dataset_end(self) -> date:
 69        """
 70        Latest available date in the dataset.
 71        """
 72        return self._max_dataset_end
 73
 74    @property
 75    def shape(self) -> tuple[int, ...]:
 76        """
 77        Shape of a single dataset sample.
 78        """
 79        return self._shape
 80
 81    @retry_on_network_errors(retries=DEFAULT_RETRIES, backoff=DEFAULT_BACKOFF)
 82    def info(self, per_year: bool = False, threads: int = 24) -> dict[str, any]:
 83        """
 84        Collect dataset size statistics.
 85
 86        Args:
 87            per_year (`bool`, optional): If True, include per-year file and size statistics. Defaults to False.
 88            threads (`int`, optional): Number of threads used for parallel HTTP requests. Defaults to 24.
 89        """
 90        total_files, total_size = 0, 0
 91        per_year_result = defaultdict(
 92            lambda: {
 93                KEY_FILES: 0,
 94                KEY_SIZE_BYTES: 0,
 95                KEY_SIZE_MB: 0.0,
 96            }
 97        )
 98
 99        with ThreadPoolExecutor(max_workers=threads) as executor:
100            futures = [
101                executor.submit(self._fetch_year_stats, year)
102                for year in range(
103                    self.dataset_start.year,
104                    self.dataset_end.year + 1,
105                )
106            ]
107
108            for future in as_completed(futures):
109                year, files, size = future.result()
110
111                per_year_result[year][KEY_FILES] = files
112                per_year_result[year][KEY_SIZE_BYTES] = size
113                per_year_result[year][KEY_SIZE_MB] = round(size / BYTES_IN_MB, 2)
114
115                total_files += files
116                total_size += size
117
118        result: dict[str, any] = {
119            KEY_DATASET_START: self.dataset_start,
120            KEY_DATASET_END: self.dataset_end,
121            KEY_SHAPE: self.shape,
122            f"total_{KEY_FILES}": total_files,
123            f"total_{KEY_SIZE_BYTES}": total_size,
124            f"total_{KEY_SIZE_MB}": round(total_size / BYTES_IN_MB, 2),
125        }
126
127        if per_year:
128            result[KEY_PER_YEAR] = dict(per_year_result)
129
130        return result
131
132    def get_filenames(
133        self,
134        start: date | None = None,
135        end: date | None = None,
136        step: int | str | None = None,
137    ) -> list[str]:
138        """
139        Generate dataset filenames for a date range.
140
141        Args:
142            start (`date`, optional): Start date (inclusive). Defaults to dataset start.
143            end (`date`, optional): End date (inclusive). Defaults to dataset end.
144            step (`int` or `str`, optional): Step between files. If `int` - number of days.
145                If `str` - format like `"1d"`, `"1w"`, `"1m"`, `"1y"`.
146                For month or years steps (`"1m"`, `"2m"`, etc.), the date always lands on the last day
147                of the month (e.g., Jan 31 + 1 month = Feb 28/29, then Mar 31).
148                Defaults to 1 day.
149        """
150        start = start or self.dataset_start
151        end = end or self.dataset_end
152
153        if start < self.dataset_start:
154            raise ValueError(f"date start value should be > {self.dataset_start}")
155
156        if end > self.dataset_end:
157            raise ValueError(f"date end value should be < {self.dataset_end}")
158
159        if start > end:
160            raise ValueError("start date must be <= date end")
161
162        filenames: list[str] = []
163        current = start
164        delta = convert_step_to_delta(step=step)
165
166        while current <= end:
167            filenames.append(get_filename_template(current))
168            current += delta
169
170        return filenames
171
172    @retry_on_network_errors(retries=DEFAULT_RETRIES, backoff=DEFAULT_BACKOFF)
173    def read_file(self, filename: str) -> bytes | None:
174        """
175        Load a dataset file from Hugging Face into memory.
176
177        Args:
178            filename (`str`): Relative path to the dataset file.
179        """
180        url = f"{self._api_base_url}/datasets/{self._dataset_repo}/resolve/main/{filename}"
181        buffer = BytesIO()
182        try:
183            http_get(
184                url=url,
185                temp_file=buffer,
186                displayed_filename=filename,
187                headers=self._api_headers,
188            )
189            return buffer.getvalue()
190
191        # ignore if file isn't found
192        except RemoteEntryNotFoundError:
193            return None
194
195        except Exception as e:
196            raise RuntimeError(f"Failed to get file {filename}") from e
197
198    @retry_on_network_errors(retries=DEFAULT_RETRIES, backoff=DEFAULT_BACKOFF)
199    def download_file(self, filename: str, local_dir: str) -> str | None:
200        """
201        Download a dataset file to a local directory.
202
203        Args:
204            filename (`str`): Dataset file path.
205            local_dir (`str`): Target directory for download.
206        """
207        try:
208            return self._api.hf_hub_download(
209                repo_id=self._dataset_repo,
210                repo_type=self._dataset_repo_type,
211                filename=filename,
212                local_dir=local_dir,
213            )
214
215        # ignore if file isn't found
216        except RemoteEntryNotFoundError:
217            return None
218
219        except Exception as e:
220            raise RuntimeError(f"Failed to download file {filename}") from e
221
222    @lru_cache(maxsize=YEAR_STATS_CACHE_SIZE)
223    def _fetch_year_stats(self, year: int) -> tuple[int, int, int]:
224        url = f"{self._api_base_url}/api/datasets/{self._dataset_repo}/tree/main/global_series/{year}"
225
226        resp = requests.get(
227            url, timeout=DEFAULT_REQUEST_TIMEOUT, headers=self._api_headers
228        )
229        resp.raise_for_status()
230
231        files, size = 0, 0
232        for item in resp.json():
233            if item.get("type") != "file":
234                continue
235
236            files += 1
237            size += item.get("size", 0)
238
239        return year, files, size