file_utils.py 17.5 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
2
3
4
5
"""
Utilities for working with the local dataset cache.
This file is adapted from the AllenNLP library at https://github.com/allenai/allennlp
Copyright by the AllenNLP authors.
"""
Aymeric Augustin's avatar
Aymeric Augustin committed
6

Aymeric Augustin's avatar
Aymeric Augustin committed
7
import fnmatch
thomwolf's avatar
thomwolf committed
8
import json
thomwolf's avatar
thomwolf committed
9
import logging
thomwolf's avatar
thomwolf committed
10
import os
11
import shutil
Aymeric Augustin's avatar
Aymeric Augustin committed
12
import sys
13
import tarfile
thomwolf's avatar
thomwolf committed
14
import tempfile
Aymeric Augustin's avatar
Aymeric Augustin committed
15
from contextlib import contextmanager
16
from functools import partial, wraps
thomwolf's avatar
thomwolf committed
17
from hashlib import sha256
18
from typing import Optional
Aymeric Augustin's avatar
Aymeric Augustin committed
19
from urllib.parse import urlparse
20
from zipfile import ZipFile, is_zipfile
thomwolf's avatar
thomwolf committed
21
22

import boto3
Aymeric Augustin's avatar
Aymeric Augustin committed
23
import requests
24
from botocore.config import Config
thomwolf's avatar
thomwolf committed
25
from botocore.exceptions import ClientError
Aymeric Augustin's avatar
Aymeric Augustin committed
26
from filelock import FileLock
27
from tqdm.auto import tqdm
Aymeric Augustin's avatar
Aymeric Augustin committed
28

29
from . import __version__
thomwolf's avatar
thomwolf committed
30

Lysandre's avatar
Lysandre committed
31

thomwolf's avatar
thomwolf committed
32
33
logger = logging.getLogger(__name__)  # pylint: disable=invalid-name

thomwolf's avatar
thomwolf committed
34
try:
35
36
37
    USE_TF = os.environ.get("USE_TF", "AUTO").upper()
    USE_TORCH = os.environ.get("USE_TORCH", "AUTO").upper()
    if USE_TORCH in ("1", "ON", "YES", "AUTO") and USE_TF not in ("1", "ON", "YES"):
38
        import torch
39

40
41
        _torch_available = True  # pylint: disable=invalid-name
        logger.info("PyTorch version {} available.".format(torch.__version__))
42
    else:
43
        logger.info("Disabling PyTorch because USE_TF is set")
44
        _torch_available = False
thomwolf's avatar
thomwolf committed
45
46
47
except ImportError:
    _torch_available = False  # pylint: disable=invalid-name

Lysandre's avatar
Lysandre committed
48
try:
49
50
51
52
    USE_TF = os.environ.get("USE_TF", "AUTO").upper()
    USE_TORCH = os.environ.get("USE_TORCH", "AUTO").upper()

    if USE_TF in ("1", "ON", "YES", "AUTO") and USE_TORCH not in ("1", "ON", "YES"):
thomwolf's avatar
thomwolf committed
53
        import tensorflow as tf
54
55

        assert hasattr(tf, "__version__") and int(tf.__version__[0]) >= 2
thomwolf's avatar
thomwolf committed
56
57
58
        _tf_available = True  # pylint: disable=invalid-name
        logger.info("TensorFlow version {} available.".format(tf.__version__))
    else:
59
        logger.info("Disabling Tensorflow because USE_TORCH is set")
thomwolf's avatar
thomwolf committed
60
        _tf_available = False
Lysandre's avatar
Lysandre committed
61
62
except (ImportError, AssertionError):
    _tf_available = False  # pylint: disable=invalid-name
thomwolf's avatar
thomwolf committed
63

64
65
try:
    from torch.hub import _get_torch_home
66

67
68
69
    torch_cache_home = _get_torch_home()
except ImportError:
    torch_cache_home = os.path.expanduser(
70
71
72
        os.getenv("TORCH_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "torch"))
    )
default_cache_path = os.path.join(torch_cache_home, "transformers")
73

thomwolf's avatar
thomwolf committed
74
75
try:
    from pathlib import Path
76

77
    PYTORCH_PRETRAINED_BERT_CACHE = Path(
78
79
        os.getenv("PYTORCH_TRANSFORMERS_CACHE", os.getenv("PYTORCH_PRETRAINED_BERT_CACHE", default_cache_path))
    )
80
except (AttributeError, ImportError):
81
82
83
    PYTORCH_PRETRAINED_BERT_CACHE = os.getenv(
        "PYTORCH_TRANSFORMERS_CACHE", os.getenv("PYTORCH_PRETRAINED_BERT_CACHE", default_cache_path)
    )
84
85

PYTORCH_TRANSFORMERS_CACHE = PYTORCH_PRETRAINED_BERT_CACHE  # Kept for backward compatibility
86
TRANSFORMERS_CACHE = PYTORCH_PRETRAINED_BERT_CACHE  # Kept for backward compatibility
thomwolf's avatar
thomwolf committed
87

88
WEIGHTS_NAME = "pytorch_model.bin"
89
90
TF2_WEIGHTS_NAME = "tf_model.h5"
TF_WEIGHTS_NAME = "model.ckpt"
91
CONFIG_NAME = "config.json"
92
MODEL_CARD_NAME = "modelcard.json"
Thomas Wolf's avatar
Thomas Wolf committed
93

Lysandre's avatar
Lysandre committed
94
95

MULTIPLE_CHOICE_DUMMY_INPUTS = [[[0], [1]], [[0], [1]]]
96
97
98
DUMMY_INPUTS = [[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]]
DUMMY_MASK = [[1, 1, 1, 1, 1], [1, 1, 1, 0, 0], [0, 0, 0, 1, 1]]

99
S3_BUCKET_PREFIX = "https://s3.amazonaws.com/models.huggingface.co/bert"
100
CLOUDFRONT_DISTRIB_PREFIX = "https://d2ws9o8vfrpkyk.cloudfront.net"
101

Thomas Wolf's avatar
Thomas Wolf committed
102

thomwolf's avatar
thomwolf committed
103
104
105
def is_torch_available():
    return _torch_available

106

thomwolf's avatar
thomwolf committed
107
108
109
def is_tf_available():
    return _tf_available

110

Aymeric Augustin's avatar
Aymeric Augustin committed
111
112
def add_start_docstrings(*docstr):
    def docstring_decorator(fn):
113
114
115
116
117
118
119
120
121
122
        fn.__doc__ = "".join(docstr) + (fn.__doc__ if fn.__doc__ is not None else "")
        return fn

    return docstring_decorator


def add_start_docstrings_to_callable(*docstr):
    def docstring_decorator(fn):
        class_name = ":class:`~transformers.{}`".format(fn.__qualname__.split(".")[0])
        intro = "   The {} forward method, overrides the :func:`__call__` special method.".format(class_name)
Lysandre's avatar
Lysandre committed
123
124
        note = r"""

125
126
127
128
    .. note::
        Although the recipe for forward pass needs to be defined within
        this function, one should call the :class:`Module` instance afterwards
        instead of this since the former takes care of running the
129
        pre and post processing steps while the latter silently ignores them.
130
131
        """
        fn.__doc__ = intro + note + "".join(docstr) + (fn.__doc__ if fn.__doc__ is not None else "")
Aymeric Augustin's avatar
Aymeric Augustin committed
132
        return fn
133

Aymeric Augustin's avatar
Aymeric Augustin committed
134
    return docstring_decorator
135

136

Aymeric Augustin's avatar
Aymeric Augustin committed
137
138
139
140
def add_end_docstrings(*docstr):
    def docstring_decorator(fn):
        fn.__doc__ = fn.__doc__ + "".join(docstr)
        return fn
141

Aymeric Augustin's avatar
Aymeric Augustin committed
142
    return docstring_decorator
thomwolf's avatar
thomwolf committed
143

144
145
146

def is_remote_url(url_or_filename):
    parsed = urlparse(url_or_filename)
147
148
    return parsed.scheme in ("http", "https", "s3")

149

150
def hf_bucket_url(identifier, postfix=None, cdn=False) -> str:
151
    endpoint = CLOUDFRONT_DISTRIB_PREFIX if cdn else S3_BUCKET_PREFIX
152
    if postfix is None:
153
        return "/".join((endpoint, identifier))
154
    else:
155
        return "/".join((endpoint, identifier, postfix))
156
157


thomwolf's avatar
thomwolf committed
158
def url_to_filename(url, etag=None):
thomwolf's avatar
thomwolf committed
159
160
161
162
    """
    Convert `url` into a hashed filename in a repeatable way.
    If `etag` is specified, append its hash to the url's, delimited
    by a period.
163
    If the url ends with .h5 (Keras HDF5 weights) adds '.h5' to the name
thomwolf's avatar
thomwolf committed
164
165
    so that TF 2.0 can identify it as a HDF5 file
    (see https://github.com/tensorflow/tensorflow/blob/00fad90125b18b80fe054de1055770cfb8fe4ba3/tensorflow/python/keras/engine/network.py#L1380)
thomwolf's avatar
thomwolf committed
166
    """
167
    url_bytes = url.encode("utf-8")
thomwolf's avatar
thomwolf committed
168
169
170
171
    url_hash = sha256(url_bytes)
    filename = url_hash.hexdigest()

    if etag:
172
        etag_bytes = etag.encode("utf-8")
thomwolf's avatar
thomwolf committed
173
        etag_hash = sha256(etag_bytes)
174
        filename += "." + etag_hash.hexdigest()
thomwolf's avatar
thomwolf committed
175

176
177
    if url.endswith(".h5"):
        filename += ".h5"
thomwolf's avatar
thomwolf committed
178

thomwolf's avatar
thomwolf committed
179
180
181
    return filename


thomwolf's avatar
thomwolf committed
182
def filename_to_url(filename, cache_dir=None):
thomwolf's avatar
thomwolf committed
183
184
    """
    Return the url and etag (which may be ``None``) stored for `filename`.
thomwolf's avatar
thomwolf committed
185
    Raise ``EnvironmentError`` if `filename` or its stored metadata do not exist.
thomwolf's avatar
thomwolf committed
186
187
    """
    if cache_dir is None:
188
        cache_dir = TRANSFORMERS_CACHE
189
    if isinstance(cache_dir, Path):
190
        cache_dir = str(cache_dir)
thomwolf's avatar
thomwolf committed
191
192
193

    cache_path = os.path.join(cache_dir, filename)
    if not os.path.exists(cache_path):
thomwolf's avatar
thomwolf committed
194
        raise EnvironmentError("file {} not found".format(cache_path))
thomwolf's avatar
thomwolf committed
195

196
    meta_path = cache_path + ".json"
thomwolf's avatar
thomwolf committed
197
    if not os.path.exists(meta_path):
thomwolf's avatar
thomwolf committed
198
        raise EnvironmentError("file {} not found".format(meta_path))
thomwolf's avatar
thomwolf committed
199

thomwolf's avatar
thomwolf committed
200
    with open(meta_path, encoding="utf-8") as meta_file:
thomwolf's avatar
thomwolf committed
201
        metadata = json.load(meta_file)
202
203
    url = metadata["url"]
    etag = metadata["etag"]
thomwolf's avatar
thomwolf committed
204
205
206
207

    return url, etag


208
def cached_path(
209
210
211
212
213
214
215
216
    url_or_filename,
    cache_dir=None,
    force_download=False,
    proxies=None,
    resume_download=False,
    user_agent=None,
    extract_compressed_file=False,
    force_extract=False,
217
    local_files_only=False,
218
) -> Optional[str]:
thomwolf's avatar
thomwolf committed
219
220
221
222
223
    """
    Given something that might be a URL (or might be a local path),
    determine which. If it's a URL, download the file and cache it, and
    return the path to the cached file. If it's already a local path,
    make sure the file exists and then return the path.
224
225
226
    Args:
        cache_dir: specify a cache directory to save the file to (overwrite the default cache dir).
        force_download: if True, re-dowload the file even if it's already cached in the cache dir.
227
        resume_download: if True, resume the download if incompletly recieved file is found.
228
        user_agent: Optional string or dict that will be appended to the user-agent on remote requests.
229
230
231
232
        extract_compressed_file: if True and the path point to a zip or tar file, extract the compressed
            file in a folder along the archive.
        force_extract: if True when extract_compressed_file is True and the archive was already extracted,
            re-extract the archive and overide the folder where it was extracted.
233
234
235
236

    Return:
        None in case of non-recoverable file (non-existent or inaccessible url + no cache on disk).
        Local path (string) otherwise
thomwolf's avatar
thomwolf committed
237
238
    """
    if cache_dir is None:
239
        cache_dir = TRANSFORMERS_CACHE
240
    if isinstance(url_or_filename, Path):
241
        url_or_filename = str(url_or_filename)
242
    if isinstance(cache_dir, Path):
243
        cache_dir = str(cache_dir)
thomwolf's avatar
thomwolf committed
244

245
    if is_remote_url(url_or_filename):
thomwolf's avatar
thomwolf committed
246
        # URL, so get it from the cache (downloading if necessary)
247
        output_path = get_from_cache(
248
249
250
251
252
253
            url_or_filename,
            cache_dir=cache_dir,
            force_download=force_download,
            proxies=proxies,
            resume_download=resume_download,
            user_agent=user_agent,
254
            local_files_only=local_files_only,
255
        )
thomwolf's avatar
thomwolf committed
256
257
    elif os.path.exists(url_or_filename):
        # File, and it exists.
258
        output_path = url_or_filename
259
    elif urlparse(url_or_filename).scheme == "":
thomwolf's avatar
thomwolf committed
260
        # File, but it doesn't exist.
thomwolf's avatar
thomwolf committed
261
        raise EnvironmentError("file {} not found".format(url_or_filename))
thomwolf's avatar
thomwolf committed
262
263
264
265
    else:
        # Something unknown
        raise ValueError("unable to parse {} as a URL or as a local path".format(url_or_filename))

266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
    if extract_compressed_file:
        if not is_zipfile(output_path) and not tarfile.is_tarfile(output_path):
            return output_path

        # Path where we extract compressed archives
        # We avoid '.' in dir name and add "-extracted" at the end: "./model.zip" => "./model-zip-extracted/"
        output_dir, output_file = os.path.split(output_path)
        output_extract_dir_name = output_file.replace(".", "-") + "-extracted"
        output_path_extracted = os.path.join(output_dir, output_extract_dir_name)

        if os.path.isdir(output_path_extracted) and os.listdir(output_path_extracted) and not force_extract:
            return output_path_extracted

        # Prevent parallel extractions
        lock_path = output_path + ".lock"
        with FileLock(lock_path):
            shutil.rmtree(output_path_extracted, ignore_errors=True)
            os.makedirs(output_path_extracted)
            if is_zipfile(output_path):
                with ZipFile(output_path, "r") as zip_file:
                    zip_file.extractall(output_path_extracted)
                    zip_file.close()
            elif tarfile.is_tarfile(output_path):
                tar_file = tarfile.open(output_path)
                tar_file.extractall(output_path_extracted)
                tar_file.close()
thomwolf's avatar
cleanup  
thomwolf committed
292
293
            else:
                raise EnvironmentError("Archive format of {} could not be identified".format(output_path))
294
295
296
297
298

        return output_path_extracted

    return output_path

thomwolf's avatar
thomwolf committed
299

thomwolf's avatar
thomwolf committed
300
def split_s3_path(url):
thomwolf's avatar
thomwolf committed
301
302
303
304
305
306
307
308
309
310
311
312
    """Split a full s3 path into the bucket name and path."""
    parsed = urlparse(url)
    if not parsed.netloc or not parsed.path:
        raise ValueError("bad s3 path {}".format(url))
    bucket_name = parsed.netloc
    s3_path = parsed.path
    # Remove '/' at beginning of path.
    if s3_path.startswith("/"):
        s3_path = s3_path[1:]
    return bucket_name, s3_path


thomwolf's avatar
thomwolf committed
313
def s3_request(func):
thomwolf's avatar
thomwolf committed
314
315
316
317
318
319
    """
    Wrapper function for s3 requests in order to create more helpful error
    messages.
    """

    @wraps(func)
thomwolf's avatar
thomwolf committed
320
    def wrapper(url, *args, **kwargs):
thomwolf's avatar
thomwolf committed
321
322
323
324
        try:
            return func(url, *args, **kwargs)
        except ClientError as exc:
            if int(exc.response["Error"]["Code"]) == 404:
thomwolf's avatar
thomwolf committed
325
                raise EnvironmentError("file {} not found".format(url))
thomwolf's avatar
thomwolf committed
326
327
328
329
330
331
332
            else:
                raise

    return wrapper


@s3_request
333
def s3_etag(url, proxies=None):
thomwolf's avatar
thomwolf committed
334
    """Check ETag on S3 object."""
335
    s3_resource = boto3.resource("s3", config=Config(proxies=proxies))
thomwolf's avatar
thomwolf committed
336
337
338
339
340
341
    bucket_name, s3_path = split_s3_path(url)
    s3_object = s3_resource.Object(bucket_name, s3_path)
    return s3_object.e_tag


@s3_request
342
def s3_get(url, temp_file, proxies=None):
thomwolf's avatar
thomwolf committed
343
    """Pull a file directly from S3."""
344
    s3_resource = boto3.resource("s3", config=Config(proxies=proxies))
thomwolf's avatar
thomwolf committed
345
346
347
348
    bucket_name, s3_path = split_s3_path(url)
    s3_resource.Bucket(bucket_name).download_fileobj(s3_path, temp_file)


349
350
def http_get(url, temp_file, proxies=None, resume_size=0, user_agent=None):
    ua = "transformers/{}; python/{}".format(__version__, sys.version.split()[0])
351
352
353
354
    if is_torch_available():
        ua += "; torch/{}".format(torch.__version__)
    if is_tf_available():
        ua += "; tensorflow/{}".format(tf.__version__)
355
    if isinstance(user_agent, dict):
356
        ua += "; " + "; ".join("{}/{}".format(k, v) for k, v in user_agent.items())
Aymeric Augustin's avatar
Aymeric Augustin committed
357
    elif isinstance(user_agent, str):
358
359
        ua += "; " + user_agent
    headers = {"user-agent": ua}
360
    if resume_size > 0:
361
        headers["Range"] = "bytes=%d-" % (resume_size,)
362
363
364
    response = requests.get(url, stream=True, proxies=proxies, headers=headers)
    if response.status_code == 416:  # Range not satisfiable
        return
365
    content_length = response.headers.get("Content-Length")
366
    total = resume_size + int(content_length) if content_length is not None else None
367
368
369
370
371
372
    progress = tqdm(
        unit="B",
        unit_scale=True,
        total=total,
        initial=resume_size,
        desc="Downloading",
thomwolf's avatar
thomwolf committed
373
        disable=bool(logger.getEffectiveLevel() == logging.NOTSET),
374
    )
375
    for chunk in response.iter_content(chunk_size=1024):
376
        if chunk:  # filter out keep-alive new chunks
thomwolf's avatar
thomwolf committed
377
378
379
380
381
            progress.update(len(chunk))
            temp_file.write(chunk)
    progress.close()


382
def get_from_cache(
383
384
385
386
387
388
389
390
    url,
    cache_dir=None,
    force_download=False,
    proxies=None,
    etag_timeout=10,
    resume_download=False,
    user_agent=None,
    local_files_only=False,
391
) -> Optional[str]:
thomwolf's avatar
thomwolf committed
392
    """
393
    Given a URL, look for the corresponding file in the local cache.
thomwolf's avatar
thomwolf committed
394
    If it's not there, download it. Then return the path to the cached file.
395
396
397
398

    Return:
        None in case of non-recoverable file (non-existent or inaccessible url + no cache on disk).
        Local path (string) otherwise
thomwolf's avatar
thomwolf committed
399
400
    """
    if cache_dir is None:
401
        cache_dir = TRANSFORMERS_CACHE
402
    if isinstance(cache_dir, Path):
403
        cache_dir = str(cache_dir)
thomwolf's avatar
thomwolf committed
404

405
    os.makedirs(cache_dir, exist_ok=True)
thomwolf's avatar
thomwolf committed
406

407
408
409
410
411
412
413
414
415
416
417
418
419
    etag = None
    if not local_files_only:
        # Get eTag to add to filename, if it exists.
        if url.startswith("s3://"):
            etag = s3_etag(url, proxies=proxies)
        else:
            try:
                response = requests.head(url, allow_redirects=True, proxies=proxies, timeout=etag_timeout)
                if response.status_code == 200:
                    etag = response.headers.get("ETag")
            except (EnvironmentError, requests.exceptions.Timeout):
                # etag is already None
                pass
thomwolf's avatar
thomwolf committed
420
421
422
423
424
425

    filename = url_to_filename(url, etag)

    # get cache path to put the file
    cache_path = os.path.join(cache_dir, filename)

426
    # etag is None = we don't have a connection, or url doesn't exist, or is otherwise inaccessible.
427
    # try to get the last downloaded one
428
429
430
431
432
433
434
435
436
437
438
439
    if etag is None:
        if os.path.exists(cache_path):
            return cache_path
        else:
            matching_files = [
                file
                for file in fnmatch.filter(os.listdir(cache_dir), filename + ".*")
                if not file.endswith(".json") and not file.endswith(".lock")
            ]
            if len(matching_files) > 0:
                return os.path.join(cache_dir, matching_files[-1])
            else:
440
441
442
443
444
445
446
447
448
                # If files cannot be found and local_files_only=True,
                # the models might've been found if local_files_only=False
                # Notify the user about that
                if local_files_only:
                    raise ValueError(
                        "Cannot find the requested files in the cached path and outgoing traffic has been"
                        " disabled. To enable model look-ups and downloads online, set 'local_files_only'"
                        " to False."
                    )
449
450
451
452
453
                return None

    # From now on, etag is not None.
    if os.path.exists(cache_path) and not force_download:
        return cache_path
454

455
    # Prevent parallel downloads of the same file with a lock.
456
    lock_path = cache_path + ".lock"
457
458
459
    with FileLock(lock_path):

        if resume_download:
460
461
            incomplete_path = cache_path + ".incomplete"

462
463
            @contextmanager
            def _resumable_file_manager():
464
                with open(incomplete_path, "a+b") as f:
465
                    yield f
466

467
468
469
470
471
            temp_file_manager = _resumable_file_manager
            if os.path.exists(incomplete_path):
                resume_size = os.stat(incomplete_path).st_size
            else:
                resume_size = 0
472
        else:
473
            temp_file_manager = partial(tempfile.NamedTemporaryFile, dir=cache_dir, delete=False)
474
            resume_size = 0
475

476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
        # Download to temporary file, then copy to cache dir once finished.
        # Otherwise you get corrupt cache entries if the download gets interrupted.
        with temp_file_manager() as temp_file:
            logger.info("%s not found in cache or force_download set to True, downloading to %s", url, temp_file.name)

            # GET file object
            if url.startswith("s3://"):
                if resume_download:
                    logger.warn('Warning: resumable downloads are not implemented for "s3://" urls')
                s3_get(url, temp_file, proxies=proxies)
            else:
                http_get(url, temp_file, proxies=proxies, resume_size=resume_size, user_agent=user_agent)

        logger.info("storing %s in cache at %s", url, cache_path)
        os.rename(temp_file.name, cache_path)

        logger.info("creating metadata file for %s", cache_path)
        meta = {"url": url, "etag": etag}
        meta_path = cache_path + ".json"
        with open(meta_path, "w") as meta_file:
            json.dump(meta, meta_file)
thomwolf's avatar
thomwolf committed
497
498

    return cache_path