fix alignment when using uneven batches and left pad (9f1b37dd) · Commits · Simon Will / fairseq

fairseq/sequence_generator.py

+10 −2

Changes for fairseq/sequence_generator.py: 10 added lines, 2 removed lines.
Original line number	Diff line number	Diff line
		@@ -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],
		}