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 8c33aacd9ff890..c2daf1da6db30c 100644 --- a/src/transformers/modeling_tf_utils.py +++ b/src/transformers/modeling_tf_utils.py @@ -1147,6 +1147,7 @@ def __init__(self, config, *inputs, **kwargs): 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 537d48a57e48e5..9c284a78aee56c 100644 --- a/tests/models/auto/test_modeling_tf_auto.py +++ b/tests/models/auto/test_modeling_tf_auto.py @@ -211,7 +211,7 @@ def test_from_pretrained_with_tuple_values(self): 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 @@ def test_new_model_registration(self): 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 a88435acba3e63..d636097dc28622 100644 --- a/tests/models/gpt2/test_modeling_tf_gpt2.py +++ b/tests/models/gpt2/test_modeling_tf_gpt2.py @@ -445,7 +445,7 @@ def test_onnx_runtime_optimize(self): 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 e7ac93a5adbe6a..9d1a3299349a94 100644 --- a/tests/models/whisper/test_modeling_tf_whisper.py +++ b/tests/models/whisper/test_modeling_tf_whisper.py @@ -312,7 +312,7 @@ def test_encoder_sinusoidal_embed_positions(self): 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 ebd5dfda6bc019..7ec2198dd1523f 100644 --- a/tests/utils/test_modeling_tf_core.py +++ b/tests/utils/test_modeling_tf_core.py @@ -217,7 +217,7 @@ def test_saved_model_creation_extended(self): 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()):