Coverage for transformer_lens/utilities/hf_utils.py: 75%
216 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""hf_utils.
3This module contains utility functions related to HuggingFace
4"""
6from __future__ import annotations
8import errno
9import inspect
10import json
11import logging
12import os
13import random
14import shutil
15import stat
16import time
17from typing import Any, Callable, Dict, TypeVar
19import torch
20from datasets.arrow_dataset import Dataset
21from datasets.iterable_dataset import IterableDataset
22from datasets.load import load_dataset
23from huggingface_hub import hf_hub_download
24from huggingface_hub.constants import HF_HUB_CACHE
26CACHE_DIR = HF_HUB_CACHE
27logger = logging.getLogger(__name__)
29T = TypeVar("T")
31_HF_RETRY_MAX_ATTEMPTS = 3
32_HF_RETRY_BASE_DELAY_SECONDS = 10.0
33_HF_RETRY_MAX_DELAY_SECONDS = 120.0
36def _is_hf_rate_limit_error(exc: BaseException) -> bool:
37 """Duck-typed check for HTTP 429 — covers HfHubHTTPError, requests.HTTPError, and subclasses."""
38 response = getattr(exc, "response", None)
39 return response is not None and getattr(response, "status_code", None) == 429
42def _retry_after_seconds(exc: BaseException) -> float | None:
43 """Parse the Retry-After header from a 429 response, if present and numeric."""
44 response = getattr(exc, "response", None)
45 if response is None: 45 ↛ 46line 45 didn't jump to line 46 because the condition on line 45 was never true
46 return None
47 headers = getattr(response, "headers", None) or {}
48 raw = headers.get("Retry-After") if hasattr(headers, "get") else None
49 if raw is None:
50 return None
51 try:
52 return float(raw)
53 except (TypeError, ValueError):
54 return None
57_TL_RETRY_WRAPPED_ATTR = "_tl_hf_retry_wrapped"
60def enable_hf_retry() -> None:
61 """Globally wrap transformers ``Auto*.from_pretrained`` with retry-on-429.
63 Opt-in via ``TRANSFORMERLENS_HF_RETRY=1`` or by calling this function.
64 Idempotent. See :func:`call_hf_with_retry`.
65 """
66 from transformers import (
67 AutoConfig,
68 AutoFeatureExtractor,
69 AutoModel,
70 AutoProcessor,
71 AutoTokenizer,
72 )
74 for cls in (AutoConfig, AutoModel, AutoTokenizer, AutoProcessor, AutoFeatureExtractor):
75 original = cls.from_pretrained
76 if getattr(original, _TL_RETRY_WRAPPED_ATTR, False):
77 continue
78 underlying = original.__func__ if hasattr(original, "__func__") else original
80 def _wrapped(klass, *args: Any, _orig: Any = underlying, **kwargs: Any) -> Any:
81 return call_hf_with_retry(_orig, klass, *args, **kwargs)
83 setattr(_wrapped, _TL_RETRY_WRAPPED_ATTR, True)
84 cls.from_pretrained = classmethod(_wrapped) # type: ignore[method-assign,assignment]
87def call_hf_with_retry(
88 func: Callable[..., T],
89 *args: Any,
90 max_attempts: int = _HF_RETRY_MAX_ATTEMPTS,
91 base_delay: float = _HF_RETRY_BASE_DELAY_SECONDS,
92 **kwargs: Any,
93) -> T:
94 """Retry ``func(*args, **kwargs)`` on HTTP 429, honoring ``Retry-After``.
96 Exponential backoff with ±20% jitter, capped at ``_HF_RETRY_MAX_DELAY_SECONDS``.
97 Non-429 exceptions propagate immediately.
98 """
99 for attempt in range(max_attempts): 99 ↛ 116line 99 didn't jump to line 116 because the loop on line 99 didn't complete
100 try:
101 return func(*args, **kwargs)
102 except Exception as exc:
103 if not _is_hf_rate_limit_error(exc) or attempt == max_attempts - 1:
104 raise
105 wait = _retry_after_seconds(exc)
106 if wait is None:
107 wait = min(base_delay * (2**attempt), _HF_RETRY_MAX_DELAY_SECONDS)
108 wait *= 0.8 + 0.4 * random.random()
109 logger.warning(
110 "HuggingFace Hub rate-limited (HTTP 429); retrying in %.1fs (attempt %d/%d)",
111 wait,
112 attempt + 1,
113 max_attempts,
114 )
115 time.sleep(wait)
116 raise RuntimeError("call_hf_with_retry exited loop without returning or raising")
119def get_hf_token() -> str | None:
120 """Get HuggingFace token from environment. Returns None if not set."""
121 return os.environ.get("HF_TOKEN", "") or None
124def get_rotary_pct_from_config(config: Any) -> float:
125 """Get the rotary percentage from a config object.
127 In transformers v5, rotary_pct was moved to rope_parameters['partial_rotary_factor'].
128 This function handles both the old and new config formats.
130 Args:
131 config: Config object (HuggingFace or custom)
133 Returns:
134 float: The rotary percentage (0.0 to 1.0)
135 """
136 if config is None: 136 ↛ 137line 136 didn't jump to line 137 because the condition on line 136 was never true
137 return 1.0
139 # het view: hasattr does NOT suppress the per-layer-registered raise.
140 from transformer_lens.utilities.heterogeneous_config import het_safe_view
142 config = het_safe_view(config)
144 # Try the old attribute first (transformers v4)
145 if hasattr(config, "rotary_pct"):
146 return getattr(config, "rotary_pct", 1.0)
148 # Try the new rope_parameters format (transformers v5)
149 if hasattr(config, "rope_parameters"):
150 rope_params = getattr(config, "rope_parameters", None)
151 if isinstance(rope_params, dict) and "partial_rotary_factor" in rope_params:
152 return rope_params["partial_rotary_factor"]
154 # Default to 1.0 (full rotary) if not found
155 return 1.0
158def select_compatible_kwargs(kwargs_dict: Dict[str, Any], callable: Callable) -> Dict[str, Any]:
159 """Return a dict with the elements kwargs_dict that are parameters of callable"""
160 return {k: v for k, v in kwargs_dict.items() if k in inspect.getfullargspec(callable).args}
163def download_file_from_hf(
164 repo_name,
165 file_name,
166 subfolder=".",
167 cache_dir=CACHE_DIR,
168 force_is_torch=False,
169 **kwargs,
170):
171 """
172 Helper function to download files from the HuggingFace Hub, from subfolder/file_name in repo_name, saving locally to cache_dir and returning the loaded file (if a json or Torch object) and the file path otherwise.
174 If it's a Torch file without the ".pth" extension, set force_is_torch=True to load it as a Torch object.
175 """
176 file_path = call_hf_with_retry(
177 hf_hub_download,
178 repo_id=repo_name,
179 filename=file_name,
180 subfolder=subfolder,
181 cache_dir=cache_dir,
182 **select_compatible_kwargs(kwargs, hf_hub_download),
183 )
185 if file_path.endswith(".pth") or force_is_torch:
186 return torch.load(file_path, map_location="cpu", weights_only=False)
187 elif file_path.endswith(".json"): 187 ↛ 190line 187 didn't jump to line 190 because the condition on line 187 was always true
188 return json.load(open(file_path, "r"))
189 else:
190 print("File type not supported:", file_path.split(".")[-1])
191 return file_path
194def clear_huggingface_cache():
195 """Delete the HuggingFace cache directory and all its contents.
197 Safe to call under parallel test execution (handles concurrent-delete races).
198 """
200 print("Deleting Hugging Face cache directory and all its contents.")
202 if not os.path.exists(CACHE_DIR): 202 ↛ 203line 202 didn't jump to line 203 because the condition on line 202 was never true
203 return
205 try:
206 # Use a custom error handler that only ignores specific race condition errors
207 def handle_remove_readonly(func, path, exc_info):
208 """Error handler for Windows readonly files and race conditions."""
210 excvalue = exc_info[1]
211 # Ignore "directory not empty" errors (race condition - another process deleted contents)
212 if isinstance(excvalue, OSError) and excvalue.errno == errno.ENOTEMPTY:
213 return
214 # Ignore "no such file or directory" errors (race condition - already deleted)
215 if isinstance(excvalue, FileNotFoundError):
216 return
217 if isinstance(excvalue, OSError) and excvalue.errno == errno.ENOENT:
218 return
219 # For readonly files on Windows, try to make writable and retry
220 if os.path.exists(path) and not os.access(path, os.W_OK):
221 try:
222 os.chmod(path, stat.S_IWUSR)
223 func(path)
224 except (OSError, FileNotFoundError):
225 # File disappeared or became inaccessible - race condition, ignore
226 return
227 else:
228 raise
230 shutil.rmtree(CACHE_DIR, onerror=handle_remove_readonly)
231 except FileNotFoundError:
232 # Directory was deleted by another process - that's fine
233 pass
234 except OSError as e:
235 # Only ignore "directory not empty" and "no such file" errors (race conditions)
236 if e.errno not in (errno.ENOTEMPTY, errno.ENOENT):
237 print(f"Warning: Could not fully clear cache: {e}")
240def keep_single_column(dataset: Dataset | IterableDataset, col_name: str):
241 """
242 Acts on a HuggingFace dataset to delete all columns apart from a single column name - useful when we want to tokenize and mix together different strings
243 """
244 for key in dataset.features:
245 if key != col_name: 245 ↛ 246line 245 didn't jump to line 246 because the condition on line 245 was never true
246 dataset = dataset.remove_columns(key)
247 return dataset
250def get_dataset(dataset_name: str, **kwargs) -> Dataset:
251 """
252 Returns a small HuggingFace dataset, for easy testing and exploration. Accesses several convenience datasets with 10,000 elements (dealing with the enormous 100GB - 2TB datasets is a lot of effort!). Note that it returns a dataset (ie a dictionary containing all the data), *not* a DataLoader (iterator over the data + some fancy features). But you can easily convert it to a DataLoader.
254 Each dataset has a 'text' field, which contains the relevant info, some also have several meta data fields
256 Kwargs will be passed to the huggingface dataset loading function, e.g. "data_dir"
258 Possible inputs:
259 * openwebtext (approx the GPT-2 training data https://huggingface.co/datasets/openwebtext)
260 * pile (The Pile, a big mess of tons of diverse data https://pile.eleuther.ai/)
261 * c4 (Colossal, Cleaned, Common Crawl - basically openwebtext but bigger https://huggingface.co/datasets/c4)
262 * code (Codeparrot Clean, a Python code dataset https://huggingface.co/datasets/codeparrot/codeparrot-clean )
263 * c4_code (c4 + code - the 20K data points from c4-10k and code-10k. This is the mix of datasets used to train my interpretability-friendly models, though note that they are *not* in the correct ratio! There's 10K texts for each, but about 22M tokens of code and 5M tokens of C4)
264 * wiki (Wikipedia, generated from the 20220301.en split of https://huggingface.co/datasets/wikipedia )
265 """
266 dataset_aliases = {
267 "openwebtext": "stas/openwebtext-10k",
268 "owt": "stas/openwebtext-10k",
269 "pile": "NeelNanda/pile-10k",
270 "c4": "NeelNanda/c4-10k",
271 "code": "NeelNanda/code-10k",
272 "python": "NeelNanda/code-10k",
273 "c4_code": "NeelNanda/c4-code-20k",
274 "c4-code": "NeelNanda/c4-code-20k",
275 "wiki": "NeelNanda/wiki-10k",
276 }
277 if dataset_name in dataset_aliases:
278 dataset = load_dataset(dataset_aliases[dataset_name], split="train", **kwargs)
279 else:
280 raise ValueError(f"Dataset {dataset_name} not supported")
281 return dataset
284def autoconfig_with_remote_post_init_compat(
285 model_id: str, auto_config: Any = None, **kwargs: Any
286) -> Any:
287 """``AutoConfig.from_pretrained`` tolerating 4.x-era remote-code configs.
289 Two upstream quirks, each retried once:
291 * transformers>=5 delivers non-base fields via ``__post_init__(**extras)``;
292 a 4.x-era argless override (OpenELM) crashes on the first kwarg. On exactly
293 that TypeError, wrap the class's ``__post_init__``.
294 * a repo that pins ``flash_attention_2`` in its config cannot be asked for
295 ``output_attentions``. The bridge needs attention outputs and loads the model
296 eager regardless, so retry asking for the implementation the loader will use.
297 """
298 if auto_config is None: 298 ↛ 301line 298 didn't jump to line 301 because the condition on line 298 was never true
299 # Callers with a patchable module-level AutoConfig (boot) pass it in so
300 # test seams keep seeing the call; others get the real one.
301 from transformers import AutoConfig as auto_config # noqa: N813
303 try:
304 return auto_config.from_pretrained(model_id, **kwargs)
305 except TypeError as err:
306 if "__post_init__() got an unexpected keyword argument" not in str(err):
307 raise
308 config_class = _resolve_remote_config_class(model_id, **kwargs)
309 if config_class is None:
310 raise
311 make_post_init_kwarg_tolerant(config_class)
312 return auto_config.from_pretrained(model_id, **kwargs)
313 except _strict_dataclass_error() as err:
314 if "output_attentions" not in str(err) or "attn_implementation" in kwargs:
315 raise
316 return auto_config.from_pretrained(model_id, attn_implementation="eager", **kwargs)
319def _strict_dataclass_error() -> type[BaseException]:
320 """``StrictDataclassError``, or a placeholder that is never raised."""
321 try:
322 from huggingface_hub.errors import StrictDataclassError
324 return StrictDataclassError
325 except ImportError: # pragma: no cover - depends on huggingface_hub version
327 class _Unraisable(Exception):
328 pass
330 return _Unraisable
333def autotokenizer_with_special_token_compat(
334 model_id: str, auto_tokenizer: Any = None, **kwargs: Any
335) -> Any:
336 """``AutoTokenizer.from_pretrained`` tolerating rejected special-token declarations.
338 Some repos declare extra special tokens that their ``tokenizer.json`` already
339 contains -- the catherinearnett/B-GPT family lists ~1200 ``[XXXXXn]`` tokens --
340 and re-adding them raises ``TypeError: argument 'special_tokens': Expected
341 Union[...]``. The serialized fast tokenizer is intact and encodes identically,
342 so build straight from it and carry the config's token roles across.
343 """
344 if auto_tokenizer is None: 344 ↛ 345line 344 didn't jump to line 345 because the condition on line 344 was never true
345 from transformers import AutoTokenizer as auto_tokenizer # noqa: N813
347 try:
348 return auto_tokenizer.from_pretrained(model_id, **kwargs)
349 except TypeError as err:
350 if "special_tokens" not in str(err):
351 raise
352 tokenizer = _fast_tokenizer_from_file(model_id, **kwargs)
353 if tokenizer is None:
354 raise
355 logger.warning(
356 "%s declares special tokens its backend rejects; built from tokenizer.json instead",
357 model_id,
358 )
359 return tokenizer
362def _fast_tokenizer_from_file(model_id: str, **kwargs: Any) -> Any:
363 """Build a fast tokenizer from the repo's ``tokenizer.json``, or None."""
364 from transformers import PreTrainedTokenizerFast
366 passthrough = {key: kwargs[key] for key in ("token", "revision") if key in kwargs}
367 try:
368 tokenizer_file = hf_hub_download(model_id, "tokenizer.json", **passthrough)
369 config_file = hf_hub_download(model_id, "tokenizer_config.json", **passthrough)
370 except Exception:
371 return None # no serialized fast tokenizer: nothing to rebuild from
373 with open(config_file) as handle:
374 config = json.load(handle)
376 # A role may be a bare string or a full AddedToken dict.
377 roles = {}
378 for role in ("bos_token", "eos_token", "unk_token", "pad_token", "cls_token", "sep_token"):
379 value = config.get(role)
380 if isinstance(value, dict):
381 value = value.get("content")
382 if isinstance(value, str):
383 roles[role] = value
384 if config.get("model_max_length"): 384 ↛ 387line 384 didn't jump to line 387 because the condition on line 384 was always true
385 roles["model_max_length"] = config["model_max_length"]
386 # AutoTokenizer sets this from the repo id; downstream code reads it.
387 roles["name_or_path"] = model_id
388 for key in ("add_bos_token", "add_eos_token"):
389 if key in kwargs: 389 ↛ 390line 389 didn't jump to line 390 because the condition on line 389 was never true
390 roles[key] = kwargs[key]
392 return PreTrainedTokenizerFast(tokenizer_file=tokenizer_file, **roles)
395_TL_NUMERIC_TOWER_ATTR = "_tl_numeric_tower_tolerant"
398def enable_hf_numeric_tower() -> None:
399 """Let huggingface_hub's strict dataclasses accept an int in a float field.
401 The hub validates with a bare ``isinstance``, so a config that writes an
402 integral value for a float field -- ``"routed_scaling_factor": 1`` on the
403 DeepseekV3 family, and any ``rope_theta: 10000`` -- fails to load, though
404 PEP 484 and Python's numeric tower both accept int wherever float is asked
405 for. transformers calls ``AutoConfig`` from inside ``AutoTokenizer`` and the
406 model loaders, so nothing short of the validator itself covers every path.
408 Idempotent, and a no-op if the hub's internals move.
409 """
410 try:
411 from huggingface_hub import dataclasses as hub_dataclasses
412 except ImportError: # pragma: no cover - huggingface_hub is a hard dependency
413 return
415 original = getattr(hub_dataclasses, "_validate_simple_type", None)
416 if original is None or getattr(original, _TL_NUMERIC_TOWER_ATTR, False):
417 return
419 def tolerant(name: str, value: Any, expected_type: type) -> None:
420 # `type(value) is int` on purpose: bool is an int subclass and stays wrong.
421 if expected_type is float and type(value) is int:
422 return
423 original(name, value, expected_type)
425 setattr(tolerant, _TL_NUMERIC_TOWER_ATTR, True)
426 tolerant.__wrapped__ = original # type: ignore[attr-defined]
427 hub_dataclasses._validate_simple_type = tolerant
430def _resolve_remote_config_class(model_id: str, **kwargs: Any) -> Any:
431 """The repo's auto_map AutoConfig class, or None if there is none."""
432 from transformers.configuration_utils import PretrainedConfig
433 from transformers.dynamic_module_utils import get_class_from_dynamic_module
435 passthrough = {key: kwargs[key] for key in ("token", "revision") if key in kwargs}
436 config_dict, _ = PretrainedConfig.get_config_dict(model_id, **passthrough)
437 class_ref = (config_dict.get("auto_map") or {}).get("AutoConfig")
438 if not class_ref:
439 return None
440 return get_class_from_dynamic_module(class_ref, model_id, **passthrough)
443def make_post_init_kwarg_tolerant(config_class: Any) -> None:
444 """Wrap an argless remote ``__post_init__`` to accept 5.x extras.
446 Base-first ordering is load-bearing: the base handler setattrs the extras
447 (the class's own fields arrive that way under 5.x) and the original body
448 derives from them. Idempotent via marker.
449 """
450 from transformers.configuration_utils import PretrainedConfig
452 original = config_class.__post_init__
453 if getattr(original, "_tl_kwarg_tolerant", False):
454 return
456 def tolerant_post_init(self: Any, **extras: Any) -> None:
457 PretrainedConfig.__post_init__(self, **extras)
458 original(self)
460 tolerant_post_init._tl_kwarg_tolerant = True # type: ignore[attr-defined]
461 config_class.__post_init__ = tolerant_post_init