From e99a2314cdea112de34db62fb8a7722eb62849fc Mon Sep 17 00:00:00 2001 From: Matt Date: Thu, 11 Nov 2021 14:21:50 +0000 Subject: [PATCH] Experimenting with adding proper get_config() and from_config() methods (#14361) * Experimenting with adding proper get_config() and from_config() methods * Adding a test for get/from config * Fix test for get/from config --- src/transformers/modeling_tf_utils.py | 7 +++++++ tests/test_modeling_tf_common.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/src/transformers/modeling_tf_utils.py b/src/transformers/modeling_tf_utils.py index d46b905c4f..d7fcb6e56e 100644 --- a/src/transformers/modeling_tf_utils.py +++ b/src/transformers/modeling_tf_utils.py @@ -692,6 +692,13 @@ class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin, TFGenerationMixin, Pu self.config = config self.name_or_path = config.name_or_path + def get_config(self): + return self.config + + @classmethod + def from_config(cls, config, **kwargs): + return cls._from_config(config, **kwargs) + @classmethod def _from_config(cls, config, **kwargs): """ diff --git a/tests/test_modeling_tf_common.py b/tests/test_modeling_tf_common.py index 2709b4ff76..adb68221d3 100644 --- a/tests/test_modeling_tf_common.py +++ b/tests/test_modeling_tf_common.py @@ -160,6 +160,20 @@ class TFModelTesterMixin: self.assert_outputs_same(after_outputs, outputs) + def test_save_load_config(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + for model_class in self.all_model_classes: + model = model_class(config) + outputs = model(self._prepare_for_class(inputs_dict, model_class)) + + new_model = model_class.from_config(model.get_config()) + _ = new_model(self._prepare_for_class(inputs_dict, model_class)) # Build model + new_model.set_weights(model.get_weights()) + after_outputs = new_model(self._prepare_for_class(inputs_dict, model_class)) + + self.assert_outputs_same(after_outputs, outputs) + @tooslow def test_graph_mode(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()