SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

added cache_dir option in from_pretrained

Browse Source
This commit is contained in:
thomwolf
2018-11-26 10:21:56 +01:00
parent 029bdc0d50
commit 63ae5d2134
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
pytorch_pretrained_bert/modeling.py
View File

@@ -443,7 +443,7 @@ class PreTrainedBertModel(nn.Module):
module.bias.data.zero_() module.bias.data.zero_()
@classmethod @classmethod
def from_pretrained(cls, pretrained_model_name, *inputs, **kwargs): def from_pretrained(cls, pretrained_model_name, cache_dir=None, *inputs, **kwargs):
""" """
Instantiate a PreTrainedBertModel from a pre-trained model file. Instantiate a PreTrainedBertModel from a pre-trained model file.
Download and cache the pre-trained model file if needed. Download and cache the pre-trained model file if needed.
@@ -468,7 +468,7 @@ class PreTrainedBertModel(nn.Module):
archive_file = pretrained_model_name archive_file = pretrained_model_name
# redirect to the cache, if necessary # redirect to the cache, if necessary
try: try:
resolved_archive_file = cached_path(archive_file) resolved_archive_file = cached_path(archive_file, cache_dir=cache_dir)
except FileNotFoundError: except FileNotFoundError:
logger.error( logger.error(
"Model name '{}' was not found in model name list ({}). " "Model name '{}' was not found in model name list ({}). "