def decode_beam(self,
beam_size,
batch_size,
encoder_states,
block_ngram=0,
expand_beam=1):
dev = self.opts['device']
beams = [
Beam(beam_size,
device='cuda',
block_ngram=block_ngram,
expand_beam=expand_beam) for _ in range(batch_size)
]
decoder_input = self.sos_buffer.expand(batch_size * beam_size,
1).to(dev)
inds = torch.arange(batch_size).to(dev).unsqueeze(1).repeat(
1, beam_size).view(-1)
encoder_states = self.reorder_encoder_states(
encoder_states, inds) # not reordering but expanding
incr_state = encoder_states[1]
for ts in range(self.longest_label):
if all((b.done() for b in beams)):
break
score, incr_state, attn_w_log = self.decoder(
decoder_input, incr_state, encoder_states)
score = score[:, -1:, :]
score = score.view(batch_size, beam_size, -1)
score = F.log_softmax(score, dim=-1)
for i, b in enumerate(beams):
if not b.done():
b.advance(score[i])
incr_state_inds = torch.cat([
beam_size * i + b.get_backtrack_from_current_step()
for i, b in enumerate(beams)
])
incr_state = self.reorder_decoder_incremental_state(
incr_state, incr_state_inds)
selection = torch.cat([
b.get_output_from_current_step() for b in beams
]).unsqueeze(-1)
decoder_input = selection
for b in beams:
b.check_finished()
beam_preds_scores = [list(b.get_top_hyp()) for b in beams]
for pair in beam_preds_scores:
pair[0] = Beam.get_pretty_hypothesis(pair[0])
return beam_preds_scores, beams
