From b85c59f9978e090881caad844c164fe4a2b03c7a Mon Sep 17 00:00:00 2001 From: Julien Chaumond Date: Thu, 30 Jan 2020 16:45:52 -0500 Subject: [PATCH] config.architectures --- src/transformers/configuration_utils.py | 1 + src/transformers/modeling_utils.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index 325c4591c2..dc0b675bfa 100644 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -82,6 +82,7 @@ class PretrainedConfig(object): self.num_return_sequences = kwargs.pop("num_return_sequences", 1) # Fine-tuning task arguments + self.architectures = kwargs.pop("architectures", None) self.finetuning_task = kwargs.pop("finetuning_task", None) self.num_labels = kwargs.pop("num_labels", 2) self.id2label = kwargs.pop("id2label", {i: "LABEL_{}".format(i) for i in range(self.num_labels)}) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 7ea606b256..2e1215cb45 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -284,6 +284,9 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin): # Only save the model itself if we are using distributed training model_to_save = self.module if hasattr(self, "module") else self + # Attach architecture to the config + model_to_save.config.architectures = [model_to_save.__class__.__name__] + # Save configuration file model_to_save.config.save_pretrained(save_directory)