Given a URL, look for the corresponding dataset in the local cache. If it's not there, download it. Then return the path to the cached file.
(url, cache_dir=None)
| 173 | |
| 174 | |
| 175 | def get_from_cache(url, cache_dir=None): |
| 176 | """ |
| 177 | Given a URL, look for the corresponding dataset in the local cache. |
| 178 | If it's not there, download it. Then return the path to the cached file. |
| 179 | """ |
| 180 | if cache_dir is None: |
| 181 | cache_dir = PYTORCH_PRETRAINED_BERT_CACHE |
| 182 | if sys.version_info[0] == 3 and isinstance(cache_dir, Path): |
| 183 | cache_dir = str(cache_dir) |
| 184 | |
| 185 | if not os.path.exists(cache_dir): |
| 186 | os.makedirs(cache_dir) |
| 187 | |
| 188 | # Get eTag to add to filename, if it exists. |
| 189 | if url.startswith("s3://"): |
| 190 | etag = s3_etag(url) |
| 191 | else: |
| 192 | response = requests.head(url, allow_redirects=True) |
| 193 | if response.status_code != 200: |
| 194 | raise IOError("HEAD request failed for url {} with status code {}" |
| 195 | .format(url, response.status_code)) |
| 196 | etag = response.headers.get("ETag") |
| 197 | |
| 198 | filename = url_to_filename(url, etag) |
| 199 | |
| 200 | # get cache path to put the file |
| 201 | cache_path = os.path.join(cache_dir, filename) |
| 202 | |
| 203 | if not os.path.exists(cache_path): |
| 204 | # Download to temporary file, then copy to cache dir once finished. |
| 205 | # Otherwise you get corrupt cache entries if the download gets interrupted. |
| 206 | with tempfile.NamedTemporaryFile() as temp_file: |
| 207 | logger.info("%s not found in cache, downloading to %s", url, temp_file.name) |
| 208 | |
| 209 | # GET file object |
| 210 | if url.startswith("s3://"): |
| 211 | s3_get(url, temp_file) |
| 212 | else: |
| 213 | http_get(url, temp_file) |
| 214 | |
| 215 | # we are copying the file before closing it, so flush to avoid truncation |
| 216 | temp_file.flush() |
| 217 | # shutil.copyfileobj() starts at the current position, so go to the start |
| 218 | temp_file.seek(0) |
| 219 | |
| 220 | logger.info("copying %s to cache at %s", temp_file.name, cache_path) |
| 221 | with open(cache_path, 'wb') as cache_file: |
| 222 | shutil.copyfileobj(temp_file, cache_file) |
| 223 | |
| 224 | logger.info("creating metadata file for %s", cache_path) |
| 225 | meta = {'url': url, 'etag': etag} |
| 226 | meta_path = cache_path + '.json' |
| 227 | with open(meta_path, 'w', encoding="utf-8") as meta_file: |
| 228 | json.dump(meta, meta_file) |
| 229 | |
| 230 | logger.info("removing temp file %s", temp_file.name) |
| 231 | |
| 232 | return cache_path |
no test coverage detected