[BUGFIX] BART and mBART support 2D attention mask from tokenizer (#1637)

* fix: bart and mbart support 2d attention mask * fix: bart and mbart support 2d attention mask
PaddlePaddle · Jan 26, 2022 · 32d01fa · 32d01fa
1 parent 82b1cc4
commit 32d01fa
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/paddlenlp/transformers/bart/modeling.py b/paddlenlp/transformers/bart/modeling.py
@@ -214,7 +214,12 @@ def forward(self, input_ids=None, attention_mask=None, **kwargs):
             attention_mask = paddle.cast(
                 input_ids == self.pad_token_id,
                 dtype=paddle.get_default_dtype()).unsqueeze([1, 2]) * -1e4
-            attention_mask.stop_gradient = True
+        # For 2D attention_mask from tokenizer
+        elif attention_mask.ndim == 2:
+            attention_mask = paddle.unsqueeze(
+                attention_mask, axis=[1, 2]).astype(paddle.get_default_dtype())
+            attention_mask = (1.0 - attention_mask) * -1e4
+        attention_mask.stop_gradient = True
 
         encoder_output = self.encoder(encoder_input, src_mask=attention_mask)
         return encoder_output

diff --git a/paddlenlp/transformers/mbart/modeling.py b/paddlenlp/transformers/mbart/modeling.py
@@ -286,7 +286,12 @@ def forward(self, input_ids=None, attention_mask=None, **kwargs):
             attention_mask = paddle.cast(
                 input_ids == self.pad_token_id,
                 dtype=paddle.get_default_dtype()).unsqueeze([1, 2]) * -1e4
-            attention_mask.stop_gradient = True
+        # For 2D attention_mask from tokenizer
+        elif attention_mask.ndim == 2:
+            attention_mask = paddle.unsqueeze(
+                attention_mask, axis=[1, 2]).astype(paddle.get_default_dtype())
+            attention_mask = (1.0 - attention_mask) * -1e4
+        attention_mask.stop_gradient = True
 
         encoder_output = self.encoder(encoder_input, src_mask=attention_mask)
         return encoder_output