SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Fix some TF slow tests (#9728)

Browse Source 
* Fix saved model tests + fix a graph issue in longformer

* Apply style
This commit is contained in:
Julien Plu
2021-01-22 14:50:46 +01:00
committed by GitHub
parent 08b22722c7
commit d7c31abf38
2 changed files with 14 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/transformers/models/longformer/modeling_tf_longformer.py
View File

@@ -2438,10 +2438,16 @@ class TFLongformerForSequenceClassification(TFLongformerPreTrainedModel, TFSeque
logger.info("Initializing global attention on CLS token...") logger.info("Initializing global attention on CLS token...")
# global attention on cls token # global attention on cls token
inputs["global_attention_mask"] = tf.zeros_like(inputs["input_ids"]) inputs["global_attention_mask"] = tf.zeros_like(inputs["input_ids"])
updates = tf.ones(shape_list(inputs["input_ids"])[0], dtype=tf.int32)
indices = tf.pad(
tensor=tf.expand_dims(tf.range(shape_list(inputs["input_ids"])[0]), axis=1),
paddings=[[0, 0], [0, 1]],
constant_values=0,
)
inputs["global_attention_mask"] = tf.tensor_scatter_nd_update( inputs["global_attention_mask"] = tf.tensor_scatter_nd_update(
inputs["global_attention_mask"], inputs["global_attention_mask"],
[[i, 0] for i in range(shape_list(inputs["input_ids"])[0])], indices,
[1 for _ in range(shape_list(inputs["input_ids"])[0])], updates,
) )
outputs = self.longformer( outputs = self.longformer(

10
tests/test_modeling_tf_common.py
View File

@@ -184,7 +184,7 @@ class TFModelTesterMixin:
with tempfile.TemporaryDirectory() as tmpdirname: with tempfile.TemporaryDirectory() as tmpdirname:
model.save_pretrained(tmpdirname, saved_model=True) model.save_pretrained(tmpdirname, saved_model=True)
saved_model_dir = os.path.join(tmpdirname, "saved_model") saved_model_dir = os.path.join(tmpdirname, "saved_model", "1")
self.assertTrue(os.path.exists(saved_model_dir)) self.assertTrue(os.path.exists(saved_model_dir))
@slow @slow
@@ -204,7 +204,7 @@ class TFModelTesterMixin:
with tempfile.TemporaryDirectory() as tmpdirname: with tempfile.TemporaryDirectory() as tmpdirname:
model.save_pretrained(tmpdirname, saved_model=True) model.save_pretrained(tmpdirname, saved_model=True)
saved_model_dir = os.path.join(tmpdirname, "saved_model") saved_model_dir = os.path.join(tmpdirname, "saved_model", "1")
self.assertTrue(os.path.exists(saved_model_dir)) self.assertTrue(os.path.exists(saved_model_dir))
@slow @slow
@@ -223,7 +223,8 @@ class TFModelTesterMixin:
with tempfile.TemporaryDirectory() as tmpdirname: with tempfile.TemporaryDirectory() as tmpdirname:
model.save_pretrained(tmpdirname, saved_model=True) model.save_pretrained(tmpdirname, saved_model=True)
model = tf.keras.models.load_model(os.path.join(tmpdirname, "saved_model", "1")) saved_model_dir = os.path.join(tmpdirname, "saved_model", "1")
model = tf.keras.models.load_model(saved_model_dir)
outputs = model(class_inputs_dict) outputs = model(class_inputs_dict)
if self.is_encoder_decoder: if self.is_encoder_decoder:
@@ -262,7 +263,8 @@ class TFModelTesterMixin:
with tempfile.TemporaryDirectory() as tmpdirname: with tempfile.TemporaryDirectory() as tmpdirname:
model.save_pretrained(tmpdirname, saved_model=True) model.save_pretrained(tmpdirname, saved_model=True)
model = tf.keras.models.load_model(os.path.join(tmpdirname, "saved_model", "1")) saved_model_dir = os.path.join(tmpdirname, "saved_model", "1")
model = tf.keras.models.load_model(saved_model_dir)
outputs = model(class_inputs_dict) outputs = model(class_inputs_dict)
if self.is_encoder_decoder: if self.is_encoder_decoder: