fix mssking bias bug in Attention._attn

Author: xiaoda99

generate.py

@@ -65,6 +65,7 @@ def append_batch(X, next_idx):
     parser.add_argument('--e', type=float, default=1e-8)
     parser.add_argument('--n_valid', type=int, default=374)
     parser.add_argument('--gen_len', type=int, default=20)
+    parser.add_argument('--topk', type=int, default=10)
 
     args = parser.parse_args()
     print(args)
@@ -108,7 +109,11 @@ def append_batch(X, next_idx):
 
         for _ in range(args.gen_len):
             lm_probs = lm_model(XMB)
-            next_idx = torch.multinomial(lm_probs[:, -1, :], 1)
+            if args.topk == 0:
+                next_idx = torch.multinomial(lm_probs[:, -1, :], 1)
+            else:
+                values, indices = lm_probs[:, -1, :].topk(args.topk)
+                next_idx = indices.gather(-1, torch.multinomial(values, 1))
             next_token = text_encoder.decoder[next_idx.item()].replace('</w>', '')
             print(next_token, end=' ')
             XMB = append_batch(XMB, next_idx)

model_pytorch.py

@@ -86,7 +86,7 @@ def _attn(self, q, k, v):
             w = w / math.sqrt(v.size(-1))
         # w = w * self.b + -1e9 * (1 - self.b)  # TF implem method: mask_attn_weights
         # XD: self.b may be larger than w, so we need to crop it
-        b = self.b[:, :, w.size(-2), w.size(-1)]
+        b = self.b[:, :, :w.size(-2), :w.size(-1)]
         w = w * b + -1e9 * (1 - b)
 
         w = nn.Softmax(dim=-1)(w)