FIX [Gemma / CI] Make sure our runners have access to the model (#29242)

* pu hf token in gemma tests

* update suggestion

* add to flax

* revert

* fix

* fixup

* forward contrib credits from discussion

---------

Co-authored-by: ArthurZucker <ArthurZucker@users.noreply.github.com>
This commit is contained in:
Younes Belkada
2024-02-28 06:25:23 +01:00
committed by GitHub
parent bd5b986306
commit ad00c482c7
3 changed files with 20 additions and 4 deletions

View File

@@ -31,12 +31,14 @@ import time
import unittest
from collections import defaultdict
from collections.abc import Mapping
from functools import wraps
from io import StringIO
from pathlib import Path
from typing import Callable, Dict, Iterable, Iterator, List, Optional, Union
from unittest import mock
from unittest.mock import patch
import huggingface_hub
import urllib3
from transformers import logging as transformers_logging
@@ -460,6 +462,20 @@ def require_torch_sdpa(test_case):
return unittest.skipUnless(is_torch_sdpa_available(), "test requires PyTorch SDPA")(test_case)
def require_read_token(fn):
"""
A decorator that loads the HF token for tests that require to load gated models.
"""
token = os.getenv("HF_HUB_READ_TOKEN", None)
@wraps(fn)
def _inner(*args, **kwargs):
with patch(huggingface_hub.utils._headers, "get_token", return_value=token):
return fn(*args, **kwargs)
return _inner
def require_peft(test_case):
"""
Decorator marking a test that requires PEFT.