From 0e31e06a75b0022171056e51c8b5d53078ac5170 Mon Sep 17 00:00:00 2001 From: thomwolf Date: Fri, 24 Jan 2020 17:49:02 -0500 Subject: [PATCH] Add AutoModelForPreTraining --- src/transformers/__init__.py | 2 + src/transformers/modeling_auto.py | 169 +++++++++++++++++++++++++++ src/transformers/modeling_tf_auto.py | 167 +++++++++++++++++++++++++- tests/test_modeling_auto.py | 17 +++ tests/test_modeling_tf_auto.py | 19 +++ 5 files changed, 373 insertions(+), 1 deletion(-) diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index e305c8b15b..d2ae943c27 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -133,6 +133,7 @@ if is_torch_available(): from .modeling_utils import PreTrainedModel, prune_layer, Conv1D from .modeling_auto import ( AutoModel, + AutoModelForPreTraining, AutoModelForSequenceClassification, AutoModelForQuestionAnswering, AutoModelWithLMHead, @@ -267,6 +268,7 @@ if is_tf_available(): from .modeling_tf_utils import TFPreTrainedModel, TFSharedEmbeddings, TFSequenceSummary, shape_list from .modeling_tf_auto import ( TFAutoModel, + TFAutoModelForPreTraining, TFAutoModelForSequenceClassification, TFAutoModelForQuestionAnswering, TFAutoModelWithLMHead, diff --git a/src/transformers/modeling_auto.py b/src/transformers/modeling_auto.py index 35aaa6e61a..8b8eb98588 100644 --- a/src/transformers/modeling_auto.py +++ b/src/transformers/modeling_auto.py @@ -49,6 +49,7 @@ from .modeling_bert import ( BertForSequenceClassification, BertForTokenClassification, BertModel, + BertForPreTraining, ) from .modeling_camembert import ( CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_MAP, @@ -143,6 +144,24 @@ MODEL_MAPPING = OrderedDict( ] ) +MODEL_FOR_PRETRAINING_MAPPING = OrderedDict( + [ + (T5Config, T5WithLMHeadModel), + (DistilBertConfig, DistilBertForMaskedLM), + (AlbertConfig, AlbertForMaskedLM), + (CamembertConfig, CamembertForMaskedLM), + (XLMRobertaConfig, XLMRobertaForMaskedLM), + (RobertaConfig, RobertaForMaskedLM), + (BertConfig, BertForPreTraining), + (OpenAIGPTConfig, OpenAIGPTLMHeadModel), + (GPT2Config, GPT2LMHeadModel), + (TransfoXLConfig, TransfoXLLMHeadModel), + (XLNetConfig, XLNetLMHeadModel), + (XLMConfig, XLMWithLMHeadModel), + (CTRLConfig, CTRLLMHeadModel), + ] +) + MODEL_WITH_LM_HEAD_MAPPING = OrderedDict( [ (T5Config, T5WithLMHeadModel), @@ -348,6 +367,156 @@ class AutoModel(object): ) +class AutoModelForPreTraining(object): + r""" + :class:`~transformers.AutoModelForPreTraining` is a generic model class + that will be instantiated as one of the model classes of the library -with the architecture used for pretraining this model– when created with the `AutoModelForPreTraining.from_pretrained(pretrained_model_name_or_path)` + class method. + + This class cannot be instantiated using `__init__()` (throws an error). + """ + + def __init__(self): + raise EnvironmentError( + "AutoModelForPreTraining is designed to be instantiated " + "using the `AutoModelForPreTraining.from_pretrained(pretrained_model_name_or_path)` or " + "`AutoModelForPreTraining.from_config(config)` methods." + ) + + @classmethod + def from_config(cls, config): + r""" Instantiates one of the base model classes of the library + from a configuration. + + Args: + config (:class:`~transformers.PretrainedConfig`): + The model class to instantiate is selected based on the configuration class: + + - isInstance of `distilbert` configuration class: :class:`~transformers.DistilBertModelForMaskedLM` (DistilBERT model) + - isInstance of `roberta` configuration class: :class:`~transformers.RobertaModelForMaskedLM` (RoBERTa model) + - isInstance of `bert` configuration class: :class:`~transformers.BertForPreTraining` (Bert model) + - isInstance of `openai-gpt` configuration class: :class:`~transformers.OpenAIGPTLMHeadModel` (OpenAI GPT model) + - isInstance of `gpt2` configuration class: :class:`~transformers.GPT2ModelLMHeadModel` (OpenAI GPT-2 model) + - isInstance of `ctrl` configuration class: :class:`~transformers.CTRLModelLMHeadModel` (Salesforce CTRL model) + - isInstance of `transfo-xl` configuration class: :class:`~transformers.TransfoXLLMHeadModel` (Transformer-XL model) + - isInstance of `xlnet` configuration class: :class:`~transformers.XLNetLMHeadModel` (XLNet model) + - isInstance of `xlm` configuration class: :class:`~transformers.XLMWithLMHeadModel` (XLM model) + + Examples:: + + config = BertConfig.from_pretrained('bert-base-uncased') # Download configuration from S3 and cache. + model = AutoModelForPreTraining.from_config(config) # E.g. model was saved using `save_pretrained('./test/saved_model/')` + """ + for config_class, model_class in MODEL_FOR_PRETRAINING_MAPPING.items(): + if isinstance(config, config_class): + return model_class(config) + raise ValueError( + "Unrecognized configuration class {} for this kind of AutoModel: {}.\n" + "Model type should be one of {}.".format( + config.__class__, cls.__name__, ", ".join(c.__name__ for c in MODEL_FOR_PRETRAINING_MAPPING.keys()) + ) + ) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): + r""" Instantiates one of the model classes of the library -with the architecture used for pretraining this model– from a pre-trained model configuration. + + The `from_pretrained()` method takes care of returning the correct model class instance + based on the `model_type` property of the config object, or when it's missing, + falling back to using pattern matching on the `pretrained_model_name_or_path` string. + + The model class to instantiate is selected as the first pattern matching + in the `pretrained_model_name_or_path` string (in the following order): + - contains `t5`: :class:`~transformers.T5ModelWithLMHead` (T5 model) + - contains `distilbert`: :class:`~transformers.DistilBertForMaskedLM` (DistilBERT model) + - contains `albert`: :class:`~transformers.AlbertForMaskedLM` (ALBERT model) + - contains `camembert`: :class:`~transformers.CamembertForMaskedLM` (CamemBERT model) + - contains `xlm-roberta`: :class:`~transformers.XLMRobertaForMaskedLM` (XLM-RoBERTa model) + - contains `roberta`: :class:`~transformers.RobertaForMaskedLM` (RoBERTa model) + - contains `bert`: :class:`~transformers.BertForPreTraining` (Bert model) + - contains `openai-gpt`: :class:`~transformers.OpenAIGPTLMHeadModel` (OpenAI GPT model) + - contains `gpt2`: :class:`~transformers.GPT2LMHeadModel` (OpenAI GPT-2 model) + - contains `transfo-xl`: :class:`~transformers.TransfoXLLMHeadModel` (Transformer-XL model) + - contains `xlnet`: :class:`~transformers.XLNetLMHeadModel` (XLNet model) + - contains `xlm`: :class:`~transformers.XLMWithLMHeadModel` (XLM model) + - contains `ctrl`: :class:`~transformers.CTRLLMHeadModel` (Salesforce CTRL model) + + The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated) + To train the model, you should first set it back in training mode with `model.train()` + + Args: + pretrained_model_name_or_path: + Either: + + - a string with the `shortcut name` of a pre-trained model to load from cache or download, e.g.: ``bert-base-uncased``. + - a string with the `identifier name` of a pre-trained model that was user-uploaded to our S3, e.g.: ``dbmdz/bert-base-german-cased``. + - a path to a `directory` containing model weights saved using :func:`~transformers.PreTrainedModel.save_pretrained`, e.g.: ``./my_model_directory/``. + - a path or url to a `tensorflow index checkpoint file` (e.g. `./tf_model/model.ckpt.index`). In this case, ``from_tf`` should be set to True and a configuration object should be provided as ``config`` argument. This loading path is slower than converting the TensorFlow checkpoint in a PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards. + model_args: (`optional`) Sequence of positional arguments: + All remaning positional arguments will be passed to the underlying model's ``__init__`` method + config: (`optional`) instance of a class derived from :class:`~transformers.PretrainedConfig`: + Configuration for the model to use instead of an automatically loaded configuation. Configuration can be automatically loaded when: + + - the model is a model provided by the library (loaded with the ``shortcut-name`` string of a pretrained model), or + - the model was saved using :func:`~transformers.PreTrainedModel.save_pretrained` and is reloaded by suppling the save directory. + - the model is loaded by suppling a local directory as ``pretrained_model_name_or_path`` and a configuration JSON file named `config.json` is found in the directory. + + state_dict: (`optional`) dict: + an optional state dictionnary for the model to use instead of a state dictionary loaded from saved weights file. + This option can be used if you want to create a model from a pretrained configuration but load your own weights. + In this case though, you should check if using :func:`~transformers.PreTrainedModel.save_pretrained` and :func:`~transformers.PreTrainedModel.from_pretrained` is not a simpler option. + cache_dir: (`optional`) string: + Path to a directory in which a downloaded pre-trained model + configuration should be cached if the standard cache should not be used. + force_download: (`optional`) boolean, default False: + Force to (re-)download the model weights and configuration files and override the cached versions if they exists. + resume_download: (`optional`) boolean, default False: + Do not delete incompletely received file. Attempt to resume the download if such a file exists. + proxies: (`optional`) dict, default None: + A dictionary of proxy servers to use by protocol or endpoint, e.g.: {'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}. + The proxies are used on each request. + output_loading_info: (`optional`) boolean: + Set to ``True`` to also return a dictionnary containing missing keys, unexpected keys and error messages. + kwargs: (`optional`) Remaining dictionary of keyword arguments: + Can be used to update the configuration object (after it being loaded) and initiate the model. + (e.g. ``output_attention=True``). Behave differently depending on whether a `config` is provided or + automatically loaded: + + - If a configuration is provided with ``config``, ``**kwargs`` will be directly passed to the + underlying model's ``__init__`` method (we assume all relevant updates to the configuration have + already been done) + - If a configuration is not provided, ``kwargs`` will be first passed to the configuration class + initialization function (:func:`~transformers.PretrainedConfig.from_pretrained`). Each key of + ``kwargs`` that corresponds to a configuration attribute will be used to override said attribute + with the supplied ``kwargs`` value. Remaining keys that do not correspond to any configuration + attribute will be passed to the underlying model's ``__init__`` function. + + Examples:: + + model = AutoModelForPreTraining.from_pretrained('bert-base-uncased') # Download model and configuration from S3 and cache. + model = AutoModelForPreTraining.from_pretrained('./test/bert_model/') # E.g. model was saved using `save_pretrained('./test/saved_model/')` + model = AutoModelForPreTraining.from_pretrained('bert-base-uncased', output_attention=True) # Update configuration during loading + assert model.config.output_attention == True + # Loading from a TF checkpoint file instead of a PyTorch model (slower) + config = AutoConfig.from_json_file('./tf_model/bert_tf_model_config.json') + model = AutoModelForPreTraining.from_pretrained('./tf_model/bert_tf_checkpoint.ckpt.index', from_tf=True, config=config) + + """ + config = kwargs.pop("config", None) + if not isinstance(config, PretrainedConfig): + config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) + + for config_class, model_class in MODEL_FOR_PRETRAINING_MAPPING.items(): + if isinstance(config, config_class): + return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs) + raise ValueError( + "Unrecognized configuration class {} for this kind of AutoModel: {}.\n" + "Model type should be one of {}.".format( + config.__class__, cls.__name__, ", ".join(c.__name__ for c in MODEL_FOR_PRETRAINING_MAPPING.keys()) + ) + ) + + class AutoModelWithLMHead(object): r""" :class:`~transformers.AutoModelWithLMHead` is a generic model class diff --git a/src/transformers/modeling_tf_auto.py b/src/transformers/modeling_tf_auto.py index a20ffdcc4d..c5b4d29e60 100644 --- a/src/transformers/modeling_tf_auto.py +++ b/src/transformers/modeling_tf_auto.py @@ -46,6 +46,7 @@ from .modeling_tf_bert import ( TFBertForSequenceClassification, TFBertForTokenClassification, TFBertModel, + TFBertForPreTraining, ) from .modeling_tf_ctrl import TF_CTRL_PRETRAINED_MODEL_ARCHIVE_MAP, TFCTRLLMHeadModel, TFCTRLModel from .modeling_tf_distilbert import ( @@ -125,6 +126,22 @@ TF_MODEL_MAPPING = OrderedDict( ] ) +TF_MODEL_FOR_PRETRAINING_MAPPING = OrderedDict( + [ + (T5Config, TFT5WithLMHeadModel), + (DistilBertConfig, TFDistilBertForMaskedLM), + (AlbertConfig, TFAlbertForMaskedLM), + (RobertaConfig, TFRobertaForMaskedLM), + (BertConfig, TFBertForPreTraining), + (OpenAIGPTConfig, TFOpenAIGPTLMHeadModel), + (GPT2Config, TFGPT2LMHeadModel), + (TransfoXLConfig, TFTransfoXLLMHeadModel), + (XLNetConfig, TFXLNetLMHeadModel), + (XLMConfig, TFXLMWithLMHeadModel), + (CTRLConfig, TFCTRLLMHeadModel), + ] +) + TF_MODEL_WITH_LM_HEAD_MAPPING = OrderedDict( [ (T5Config, TFT5WithLMHeadModel), @@ -329,6 +346,154 @@ class TFAutoModel(object): ) +class TFAutoModelForPreTraining(object): + r""" + :class:`~transformers.TFAutoModelForPreTraining` is a generic model class + that will be instantiated as one of the model classes of the library -with the architecture used for pretraining this model– when created with the `TFAutoModelForPreTraining.from_pretrained(pretrained_model_name_or_path)` + class method. + + This class cannot be instantiated using `__init__()` (throws an error). + """ + + def __init__(self): + raise EnvironmentError( + "TFAutoModelForPreTraining is designed to be instantiated " + "using the `TFAutoModelForPreTraining.from_pretrained(pretrained_model_name_or_path)` or " + "`TFAutoModelForPreTraining.from_config(config)` methods." + ) + + @classmethod + def from_config(cls, config): + r""" Instantiates one of the base model classes of the library + from a configuration. + + Args: + config (:class:`~transformers.PretrainedConfig`): + The model class to instantiate is selected based on the configuration class: + + - isInstance of `distilbert` configuration class: :class:`~transformers.TFDistilBertModelForMaskedLM` (DistilBERT model) + - isInstance of `roberta` configuration class: :class:`~transformers.TFRobertaModelForMaskedLM` (RoBERTa model) + - isInstance of `bert` configuration class: :class:`~transformers.TFBertForPreTraining` (Bert model) + - isInstance of `openai-gpt` configuration class: :class:`~transformers.TFOpenAIGPTLMHeadModel` (OpenAI GPT model) + - isInstance of `gpt2` configuration class: :class:`~transformers.TFGPT2ModelLMHeadModel` (OpenAI GPT-2 model) + - isInstance of `ctrl` configuration class: :class:`~transformers.TFCTRLModelLMHeadModel` (Salesforce CTRL model) + - isInstance of `transfo-xl` configuration class: :class:`~transformers.TFTransfoXLLMHeadModel` (Transformer-XL model) + - isInstance of `xlnet` configuration class: :class:`~transformers.TFXLNetLMHeadModel` (XLNet model) + - isInstance of `xlm` configuration class: :class:`~transformers.TFXLMWithLMHeadModel` (XLM model) + + Examples:: + + config = BertConfig.from_pretrained('bert-base-uncased') # Download configuration from S3 and cache. + model = TFAutoModelForPreTraining.from_config(config) # E.g. model was saved using `save_pretrained('./test/saved_model/')` + """ + for config_class, model_class in TF_MODEL_FOR_PRETRAINING_MAPPING.items(): + if isinstance(config, config_class): + return model_class(config) + raise ValueError( + "Unrecognized configuration class {} for this kind of AutoModel: {}.\n" + "Model type should be one of {}.".format( + config.__class__, cls.__name__, ", ".join(c.__name__ for c in TF_MODEL_FOR_PRETRAINING_MAPPING.keys()) + ) + ) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): + r""" Instantiates one of the model classes of the library -with the architecture used for pretraining this model– from a pre-trained model configuration. + + The `from_pretrained()` method takes care of returning the correct model class instance + based on the `model_type` property of the config object, or when it's missing, + falling back to using pattern matching on the `pretrained_model_name_or_path` string. + + The model class to instantiate is selected as the first pattern matching + in the `pretrained_model_name_or_path` string (in the following order): + - contains `t5`: :class:`~transformers.TFT5ModelWithLMHead` (T5 model) + - contains `distilbert`: :class:`~transformers.TFDistilBertForMaskedLM` (DistilBERT model) + - contains `albert`: :class:`~transformers.TFAlbertForMaskedLM` (ALBERT model) + - contains `roberta`: :class:`~transformers.TFRobertaForMaskedLM` (RoBERTa model) + - contains `bert`: :class:`~transformers.TFBertForPreTraining` (Bert model) + - contains `openai-gpt`: :class:`~transformers.TFOpenAIGPTLMHeadModel` (OpenAI GPT model) + - contains `gpt2`: :class:`~transformers.TFGPT2LMHeadModel` (OpenAI GPT-2 model) + - contains `transfo-xl`: :class:`~transformers.TFTransfoXLLMHeadModel` (Transformer-XL model) + - contains `xlnet`: :class:`~transformers.TFXLNetLMHeadModel` (XLNet model) + - contains `xlm`: :class:`~transformers.TFXLMWithLMHeadModel` (XLM model) + - contains `ctrl`: :class:`~transformers.TFCTRLLMHeadModel` (Salesforce CTRL model) + + The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated) + To train the model, you should first set it back in training mode with `model.train()` + + Args: + pretrained_model_name_or_path: + Either: + + - a string with the `shortcut name` of a pre-trained model to load from cache or download, e.g.: ``bert-base-uncased``. + - a string with the `identifier name` of a pre-trained model that was user-uploaded to our S3, e.g.: ``dbmdz/bert-base-german-cased``. + - a path to a `directory` containing model weights saved using :func:`~transformers.PreTrainedModel.save_pretrained`, e.g.: ``./my_model_directory/``. + - a path or url to a `tensorflow index checkpoint file` (e.g. `./tf_model/model.ckpt.index`). In this case, ``from_tf`` should be set to True and a configuration object should be provided as ``config`` argument. This loading path is slower than converting the TensorFlow checkpoint in a PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards. + model_args: (`optional`) Sequence of positional arguments: + All remaning positional arguments will be passed to the underlying model's ``__init__`` method + config: (`optional`) instance of a class derived from :class:`~transformers.PretrainedConfig`: + Configuration for the model to use instead of an automatically loaded configuation. Configuration can be automatically loaded when: + + - the model is a model provided by the library (loaded with the ``shortcut-name`` string of a pretrained model), or + - the model was saved using :func:`~transformers.PreTrainedModel.save_pretrained` and is reloaded by suppling the save directory. + - the model is loaded by suppling a local directory as ``pretrained_model_name_or_path`` and a configuration JSON file named `config.json` is found in the directory. + + state_dict: (`optional`) dict: + an optional state dictionnary for the model to use instead of a state dictionary loaded from saved weights file. + This option can be used if you want to create a model from a pretrained configuration but load your own weights. + In this case though, you should check if using :func:`~transformers.PreTrainedModel.save_pretrained` and :func:`~transformers.PreTrainedModel.from_pretrained` is not a simpler option. + cache_dir: (`optional`) string: + Path to a directory in which a downloaded pre-trained model + configuration should be cached if the standard cache should not be used. + force_download: (`optional`) boolean, default False: + Force to (re-)download the model weights and configuration files and override the cached versions if they exists. + resume_download: (`optional`) boolean, default False: + Do not delete incompletely received file. Attempt to resume the download if such a file exists. + proxies: (`optional`) dict, default None: + A dictionary of proxy servers to use by protocol or endpoint, e.g.: {'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}. + The proxies are used on each request. + output_loading_info: (`optional`) boolean: + Set to ``True`` to also return a dictionnary containing missing keys, unexpected keys and error messages. + kwargs: (`optional`) Remaining dictionary of keyword arguments: + Can be used to update the configuration object (after it being loaded) and initiate the model. + (e.g. ``output_attention=True``). Behave differently depending on whether a `config` is provided or + automatically loaded: + + - If a configuration is provided with ``config``, ``**kwargs`` will be directly passed to the + underlying model's ``__init__`` method (we assume all relevant updates to the configuration have + already been done) + - If a configuration is not provided, ``kwargs`` will be first passed to the configuration class + initialization function (:func:`~transformers.PretrainedConfig.from_pretrained`). Each key of + ``kwargs`` that corresponds to a configuration attribute will be used to override said attribute + with the supplied ``kwargs`` value. Remaining keys that do not correspond to any configuration + attribute will be passed to the underlying model's ``__init__`` function. + + Examples:: + + model = TFAutoModelForPreTraining.from_pretrained('bert-base-uncased') # Download model and configuration from S3 and cache. + model = TFAutoModelForPreTraining.from_pretrained('./test/bert_model/') # E.g. model was saved using `save_pretrained('./test/saved_model/')` + model = TFAutoModelForPreTraining.from_pretrained('bert-base-uncased', output_attention=True) # Update configuration during loading + assert model.config.output_attention == True + # Loading from a TF checkpoint file instead of a PyTorch model (slower) + config = AutoConfig.from_json_file('./tf_model/bert_tf_model_config.json') + model = TFAutoModelForPreTraining.from_pretrained('./tf_model/bert_tf_checkpoint.ckpt.index', from_tf=True, config=config) + + """ + config = kwargs.pop("config", None) + if not isinstance(config, PretrainedConfig): + config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) + + for config_class, model_class in TF_MODEL_FOR_PRETRAINING_MAPPING.items(): + if isinstance(config, config_class): + return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs) + raise ValueError( + "Unrecognized configuration class {} for this kind of AutoModel: {}.\n" + "Model type should be one of {}.".format( + config.__class__, cls.__name__, ", ".join(c.__name__ for c in TF_MODEL_FOR_PRETRAINING_MAPPING.keys()) + ) + ) + + class TFAutoModelWithLMHead(object): r""" :class:`~transformers.TFAutoModelWithLMHead` is a generic model class @@ -383,7 +548,7 @@ class TFAutoModelWithLMHead(object): Examples:: config = BertConfig.from_pretrained('bert-base-uncased') # Download configuration from S3 and cache. - model = AutoModelWithLMHead.from_config(config) # E.g. model was saved using `save_pretrained('./test/saved_model/')` + model = TFAutoModelWithLMHead.from_config(config) # E.g. model was saved using `save_pretrained('./test/saved_model/')` """ for config_class, model_class in TF_MODEL_WITH_LM_HEAD_MAPPING.items(): if isinstance(config, config_class): diff --git a/tests/test_modeling_auto.py b/tests/test_modeling_auto.py index 435fc04cf0..9f027800b0 100644 --- a/tests/test_modeling_auto.py +++ b/tests/test_modeling_auto.py @@ -28,6 +28,8 @@ if is_torch_available(): BertConfig, AutoModel, BertModel, + AutoModelForPreTraining, + BertForPreTraining, AutoModelWithLMHead, BertForMaskedLM, RobertaForMaskedLM, @@ -56,6 +58,21 @@ class AutoModelTest(unittest.TestCase): for value in loading_info.values(): self.assertEqual(len(value), 0) + @slow + def test_model_for_pretraining_from_pretrained(self): + logging.basicConfig(level=logging.INFO) + for model_name in list(BERT_PRETRAINED_MODEL_ARCHIVE_MAP.keys())[:1]: + config = AutoConfig.from_pretrained(model_name) + self.assertIsNotNone(config) + self.assertIsInstance(config, BertConfig) + + model = AutoModelForPreTraining.from_pretrained(model_name) + model, loading_info = AutoModelForPreTraining.from_pretrained(model_name, output_loading_info=True) + self.assertIsNotNone(model) + self.assertIsInstance(model, BertForPreTraining) + for value in loading_info.values(): + self.assertEqual(len(value), 0) + @slow def test_lmhead_model_from_pretrained(self): logging.basicConfig(level=logging.INFO) diff --git a/tests/test_modeling_tf_auto.py b/tests/test_modeling_tf_auto.py index e86499572f..6994f6eaa9 100644 --- a/tests/test_modeling_tf_auto.py +++ b/tests/test_modeling_tf_auto.py @@ -28,6 +28,8 @@ if is_tf_available(): BertConfig, TFAutoModel, TFBertModel, + TFAutoModelForPreTraining, + TFBertForPreTraining, TFAutoModelWithLMHead, TFBertForMaskedLM, TFRobertaForMaskedLM, @@ -57,6 +59,23 @@ class TFAutoModelTest(unittest.TestCase): self.assertIsNotNone(model) self.assertIsInstance(model, TFBertModel) + @slow + def test_model_for_pretraining_from_pretrained(self): + import h5py + + self.assertTrue(h5py.version.hdf5_version.startswith("1.10")) + + logging.basicConfig(level=logging.INFO) + # for model_name in list(TF_BERT_PRETRAINED_MODEL_ARCHIVE_MAP.keys())[:1]: + for model_name in ["bert-base-uncased"]: + config = AutoConfig.from_pretrained(model_name) + self.assertIsNotNone(config) + self.assertIsInstance(config, BertConfig) + + model = TFAutoModelForPreTraining.from_pretrained(model_name) + self.assertIsNotNone(model) + self.assertIsInstance(model, TFBertForPreTraining) + @slow def test_lmhead_model_from_pretrained(self): logging.basicConfig(level=logging.INFO)