Experimenting with adding proper get_config() and from_config() methods (huggingface#14361)

Rocketknight1 · LysandreJik · commit e99a2314cdea · 2021-11-16T17:13:21.000-05:00
* Experimenting with adding proper get_config() and from_config() methods

* Adding a test for get/from config

* Fix test for get/from config
diff --git a/src/transformers/modeling_tf_utils.py b/src/transformers/modeling_tf_utils.py
@@ -692,6 +692,13 @@ def __init__(self, config, *inputs, **kwargs):
         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
@@ -160,6 +160,20 @@ def test_save_load(self):
 
                 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()
