Diverse Beam Search (8c0ca1a0) · Commits · Simon Will / fairseq

fairseq/options.py

+4 −0

Original line number	Diff line number	Diff line
		@@ -305,6 +305,10 @@ def add_generation_args(parser):
		help='sample from top K likely next words instead of all words')
		group.add_argument('--sampling-temperature', default=1, type=float, metavar='N',
		help='temperature for random sampling')
		group.add_argument('--diverse-beam-groups', default=1, type=int, metavar='N',
		help='number of groups for Diverse Beam Search')
		group.add_argument('--diverse-beam-strength', default=0.5, type=float, metavar='N',
		help='strength of diversity penalty for Diverse Beam Search')
		group.add_argument('--print-alignment', action='store_true',
		help='if set, uses attention feedback to compute and print alignment to source tokens')
		group.add_argument('--model-overrides', default="{}", type=str, metavar='DICT',

fairseq/search.py

+63 −0

Original line number	Diff line number	Diff line
		@@ -80,6 +80,69 @@ class BeamSearch(Search):
		return self.scores_buf, self.indices_buf, self.beams_buf


		class DiverseBeamSearch(Search):
		"""Diverse Beam Search.

		See "Diverse Beam Search: Decoding Diverse Solutions from Neural Sequence
		Models" for details.

		We only implement the Hamming Diversity penalty here, which performed best
		in the original paper.
		"""

		def __init__(self, tgt_dict, num_groups, diversity_strength):
		super().__init__(tgt_dict)
		self.num_groups = num_groups
		self.diversity_strength = -diversity_strength
		self.diversity_buf = None
		self.beam = BeamSearch(tgt_dict)

		def step(self, step, lprobs, scores):
		super()._init_buffers(lprobs)
		bsz, beam_size, vocab_size = lprobs.size()
		if beam_size % self.num_groups != 0:
		raise ValueError(
		'DiverseBeamSearch requires --beam to be divisible by the number of groups'
		)
		group_size = beam_size // self.num_groups

		# initialize diversity penalty
		if self.diversity_buf is None:
		self.diversity_buf = lprobs.new()
		torch.zeros(lprobs[:, 0, :].size(), out=self.diversity_buf)

		scores_G, indices_G, beams_G = [], [], []
		for g in range(self.num_groups):
		lprobs_g = lprobs[:, g::self.num_groups, :]
		scores_g = scores[:, g::self.num_groups, :] if step > 0 else None

		# apply diversity penalty
		if g > 0:
		lprobs_g = torch.add(lprobs_g, self.diversity_strength, self.diversity_buf.unsqueeze(1))
		else:
		lprobs_g = lprobs_g.contiguous()

		scores_buf, indices_buf, beams_buf = self.beam.step(step, lprobs_g, scores_g)
		beams_buf.mul_(self.num_groups).add_(g)

		scores_G.append(scores_buf.clone())
		indices_G.append(indices_buf.clone())
		beams_G.append(beams_buf.clone())

		# update diversity penalty
		self.diversity_buf.scatter_add_(
		1,
		indices_buf,
		self.diversity_buf.new_ones(indices_buf.size())
		)

		# interleave results from different groups
		self.scores_buf = torch.stack(scores_G, dim=2, out=self.scores_buf).view(bsz, -1)
		self.indices_buf = torch.stack(indices_G, dim=2, out=self.indices_buf).view(bsz, -1)
		self.beams_buf = torch.stack(beams_G, dim=2, out=self.beams_buf).view(bsz, -1)
		return self.scores_buf, self.indices_buf, self.beams_buf


		class Sampling(Search):

		def __init__(self, tgt_dict, sampling_topk=-1, sampling_temperature=1.):

fairseq/sequence_generator.py

+4 −0

Original line number	Diff line number	Diff line
		@@ -18,6 +18,7 @@ class SequenceGenerator(object):
		self, models, tgt_dict, beam_size=1, minlen=1, maxlen=None, stop_early=True,
		normalize_scores=True, len_penalty=1, unk_penalty=0, retain_dropout=False,
		sampling=False, sampling_topk=-1, sampling_temperature=1,
		diverse_beam_groups=-1, diverse_beam_strength=0.5,
		):
		"""Generates translations of a given source sentence.
		Args:
		@@ -48,6 +49,8 @@ class SequenceGenerator(object):

		if sampling:
		self.search = search.Sampling(tgt_dict, sampling_topk, sampling_temperature)
		elif diverse_beam_groups > 0:
		self.search = search.DiverseBeamSearch(tgt_dict, diverse_beam_groups, diverse_beam_strength)
		else:
		self.search = search.BeamSearch(tgt_dict)

		@@ -402,6 +405,7 @@ class SequenceGenerator(object):
		active_mask, k=beam_size, dim=1, largest=False,
		out=(_ignore, active_hypos)
		)

		active_bbsz_idx = buffer('active_bbsz_idx')
		torch.gather(
		cand_bbsz_idx, dim=1, index=active_hypos,

generate.py

+3 −3

Original line number	Diff line number	Diff line
		@@ -71,11 +71,11 @@ def main(args):
		translator = SequenceScorer(models, task.target_dictionary)
		else:
		translator = SequenceGenerator(
		models, task.target_dictionary, beam_size=args.beam,
		models, task.target_dictionary, beam_size=args.beam, minlen=args.min_len,
		stop_early=(not args.no_early_stop), normalize_scores=(not args.unnormalized),
		len_penalty=args.lenpen, unk_penalty=args.unkpen,
		sampling=args.sampling, sampling_topk=args.sampling_topk, minlen=args.min_len,
		sampling_temperature=args.sampling_temperature,
		sampling=args.sampling, sampling_topk=args.sampling_topk, sampling_temperature=args.sampling_temperature,
		diverse_beam_groups=args.diverse_beam_groups, diverse_beam_strength=args.diverse_beam_strength,
		)

		if use_cuda:

interactive.py

+5 −4

Original line number	Diff line number	Diff line
		@@ -90,10 +90,11 @@ def main(args):

		# Initialize generator
		translator = SequenceGenerator(
		models, tgt_dict, beam_size=args.beam, stop_early=(not args.no_early_stop),
		normalize_scores=(not args.unnormalized), len_penalty=args.lenpen,
		unk_penalty=args.unkpen, sampling=args.sampling, sampling_topk=args.sampling_topk,
		minlen=args.min_len, sampling_temperature=args.sampling_temperature
		models, tgt_dict, beam_size=args.beam, minlen=args.min_len,
		stop_early=(not args.no_early_stop), normalize_scores=(not args.unnormalized),
		len_penalty=args.lenpen, unk_penalty=args.unkpen,
		sampling=args.sampling, sampling_topk=args.sampling_topk, sampling_temperature=args.sampling_temperature,
		diverse_beam_groups=args.diverse_beam_groups, diverse_beam_strength=args.diverse_beam_strength,
		)

		if use_cuda: