From 06a6cb6f360e88866afdac5c0c4e295ab7da2c9b Mon Sep 17 00:00:00 2001
From: Tom Hosking <tomhosking@gmail.com>
Date: Thu, 5 Mar 2020 13:45:41 +0000
Subject: [PATCH] Refactor BartModel so that input checks are handled within
 BartEncoder and BartDecoder

---
 src/.DS_Store                     | Bin 0 -> 6148 bytes
 src/transformers/modeling_bart.py |  18 +++++++++++++-----
 2 files changed, 13 insertions(+), 5 deletions(-)
 create mode 100644 src/.DS_Store

diff --git a/src/.DS_Store b/src/.DS_Store
new file mode 100644
index 0000000000000000000000000000000000000000..cb8a7b505b4a75af5f9d6540383431344c729f6c
GIT binary patch
literal 6148
zcmeHKyH3ME5S)b+k)TLPdB4CPoTBgr`2mo4G~h@MP`l#0_%vo8Lb4nR3K}#k?atla
z&RL$q>jhxjxBF{g17J;e#D|Ba`MLYT&MIR>I`25*5eJNT!Na2Zbi%opcqQi!!;C-T
z_`V;9=fm)n_n*E%St%d|q<|EV0#e{t3V83OtuGQ4rGONW0^bVw_o2}pd*PHApALo?
z0f;N6!?=!Fg4jGk?1fVzGc-#oF{xG~h9#Z(R&~8_N=!Peh7YSJTTLhyPv`wD%3(cG
zQ3^<bxdP|8-FW}Mq5m-d&q>-z0V(jW6tLB9zuWSas<$p)&U<a6-_pJ2lkUcKP#B^e
k6Qdn-<L&q=in6Zxn&-W6N(?&lK_}{Gz;%&HfxlMZ3lPB;rT_o{

literal 0
HcmV?d00001

diff --git a/src/transformers/modeling_bart.py b/src/transformers/modeling_bart.py
index 21c51f971e..f1a6abdb19 100644
--- a/src/transformers/modeling_bart.py
+++ b/src/transformers/modeling_bart.py
@@ -271,6 +271,12 @@ class BartEncoder(nn.Module):
                 - **all_attentions** (List[Tensor]): Attention weights for each layer.
                 During training might not be of length n_layers because of layer dropout.
         """
+        # check attention mask and invert
+        if attention_mask is not None:
+            assert attention_mask.dim() == 2
+
+            attention_mask = (1.0 - attention_mask.long()) * -10000.0
+            assert attention_mask.max() <= 0
         inputs_embeds = self.embed_tokens(input_ids)
         embed_pos = self.embed_positions(input_ids)
         x = inputs_embeds + embed_pos
@@ -448,6 +454,13 @@ class BartDecoder(nn.Module):
                 - hidden states
                 - attentions
         """
+        # check attention mask and invert
+        if encoder_padding_mask is not None:
+            assert encoder_padding_mask.dim() == 2
+
+            encoder_padding_mask = (1.0 - encoder_padding_mask.long()) * -10000.0
+            assert encoder_padding_mask.max() <= 0
+
         # embed positions
         positions = self.embed_positions(input_ids, generation_mode=self.generation_mode)
 
@@ -823,11 +836,6 @@ class BartModel(PretrainedBartModel):
         decoder_attention_mask=None,
         decoder_cached_states=None,
     ):
-        if attention_mask is not None:
-            assert attention_mask.dim() == 2
-
-            attention_mask = (1.0 - attention_mask.long()) * -10000.0
-            assert attention_mask.max() <= 0
 
         # make masks if user doesn't supply
         if not self.decoder.generation_mode: