From 71d47f0ad498b7649f11d3a9cca3cd3585e4341f Mon Sep 17 00:00:00 2001 From: Matt Date: Mon, 18 Dec 2023 15:26:03 +0000 Subject: [PATCH] More TF fixes (#28081) * More build_in_name_scope() * Make sure we set the save spec now we don't do it with dummies anymore * make fixup --- src/transformers/modeling_tf_utils.py | 1 + tests/models/auto/test_modeling_tf_auto.py | 4 ++-- tests/models/gpt2/test_modeling_tf_gpt2.py | 2 +- tests/models/whisper/test_modeling_tf_whisper.py | 2 +- tests/utils/test_modeling_tf_core.py | 2 +- 5 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/transformers/modeling_tf_utils.py b/src/transformers/modeling_tf_utils.py index 8c33aacd9f..c2daf1da6d 100644 --- a/src/transformers/modeling_tf_utils.py +++ b/src/transformers/modeling_tf_utils.py @@ -1147,6 +1147,7 @@ class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin, TFGenerationMixin, Pu self.config = config self.name_or_path = config.name_or_path self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None + self._set_save_spec(self.input_signature) def get_config(self): return self.config.to_dict() diff --git a/tests/models/auto/test_modeling_tf_auto.py b/tests/models/auto/test_modeling_tf_auto.py index 537d48a57e..9c284a78ae 100644 --- a/tests/models/auto/test_modeling_tf_auto.py +++ b/tests/models/auto/test_modeling_tf_auto.py @@ -211,7 +211,7 @@ class TFAutoModelTest(unittest.TestCase): config = copy.deepcopy(model.config) config.architectures = ["FunnelBaseModel"] model = TFAutoModel.from_config(config) - model.build() + model.build_in_name_scope() self.assertIsInstance(model, TFFunnelBaseModel) @@ -249,7 +249,7 @@ class TFAutoModelTest(unittest.TestCase): config = NewModelConfig(**tiny_config.to_dict()) model = auto_class.from_config(config) - model.build() + model.build_in_name_scope() self.assertIsInstance(model, TFNewModel) diff --git a/tests/models/gpt2/test_modeling_tf_gpt2.py b/tests/models/gpt2/test_modeling_tf_gpt2.py index a88435acba..d636097dc2 100644 --- a/tests/models/gpt2/test_modeling_tf_gpt2.py +++ b/tests/models/gpt2/test_modeling_tf_gpt2.py @@ -445,7 +445,7 @@ class TFGPT2ModelTest(TFModelTesterMixin, TFCoreModelTesterMixin, PipelineTester continue model = model_class(config) - model.build() + model.build_in_name_scope() onnx_model_proto, _ = tf2onnx.convert.from_keras(model, opset=self.onnx_min_opset) diff --git a/tests/models/whisper/test_modeling_tf_whisper.py b/tests/models/whisper/test_modeling_tf_whisper.py index e7ac93a5ad..9d1a329934 100644 --- a/tests/models/whisper/test_modeling_tf_whisper.py +++ b/tests/models/whisper/test_modeling_tf_whisper.py @@ -312,7 +312,7 @@ class TFWhisperModelTest(TFModelTesterMixin, PipelineTesterMixin, unittest.TestC config = self.model_tester.get_config() for model_class in self.all_model_classes: model = model_class(config) - model.build() + model.build_in_name_scope() embeds = model.get_encoder().embed_positions.get_weights()[0] sinusoids = sinusoidal_embedding_init(embeds.shape).numpy() diff --git a/tests/utils/test_modeling_tf_core.py b/tests/utils/test_modeling_tf_core.py index ebd5dfda6b..7ec2198dd1 100644 --- a/tests/utils/test_modeling_tf_core.py +++ b/tests/utils/test_modeling_tf_core.py @@ -217,7 +217,7 @@ class TFCoreModelTesterMixin: for model_class in self.all_model_classes[:2]: class_inputs_dict = self._prepare_for_class(inputs_dict, model_class) model = model_class(config) - model.build() + model.build_in_name_scope() num_out = len(model(class_inputs_dict)) for key in list(class_inputs_dict.keys()):