fix alignment when using uneven batches and left pad
@@ -9,6 +9,7 @@ import math
import torch
from fairseq import utils
+from fairseq.data import LanguagePairDataset
from fairseq.models import FairseqIncrementalDecoder
@@ -140,6 +141,7 @@ class SequenceGenerator(object):
# helper function for allocating buffers on the fly
buffers = {}
+
def buffer(name, type_of=tokens): # noqa
if name not in buffers:
buffers[name] = type_of.new()
@@ -216,11 +218,17 @@ class SequenceGenerator(object):
sents_seen.add((sent, unfin_idx))
def get_hypo():
- _, alignment = attn_clone[i].max(dim=0)
+
+ # remove padding tokens from attn scores
+ nonpad_idxs = src_tokens[sent].ne(self.pad)
+ hypo_attn = attn_clone[i][nonpad_idxs]
+
+ _, alignment = hypo_attn.max(dim=0)
+
return {
'tokens': tokens_clone[i],
'score': score,
- 'attention': attn_clone[i], # src_len x tgt_len
+ 'attention': hypo_attn, # src_len x tgt_len
'alignment': alignment,
'positional_scores': pos_scores[i],
}
