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(