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...")
# global attention on cls token
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"],
[[i, 0] for i in range(shape_list(inputs["input_ids"])[0])],
[1 for _ in range(shape_list(inputs["input_ids"])[0])],
indices,
updates,
)
outputs = self.longformer(

10
tests/test_modeling_tf_common.py
View File

@@ -184,7 +184,7 @@ class TFModelTesterMixin:
with tempfile.TemporaryDirectory() as tmpdirname:
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))
@slow
@@ -204,7 +204,7 @@ class TFModelTesterMixin:
with tempfile.TemporaryDirectory() as tmpdirname:
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))
@slow
@@ -223,7 +223,8 @@ class TFModelTesterMixin:
with tempfile.TemporaryDirectory() as tmpdirname:
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)
if self.is_encoder_decoder:
@@ -262,7 +263,8 @@ class TFModelTesterMixin:
with tempfile.TemporaryDirectory() as tmpdirname:
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)
if self.is_encoder_decoder: