def beam_decode(self, inpute, inputc, src_pad_mask=None, context_mask=None, beam_size=8, max_len=512, length_penalty=0.0, return_all=False, clip_beam=False, fill_pad=False):
bsize, seql = inpute.size()[:2]
beam_size2 = beam_size * beam_size
bsizeb2 = bsize * beam_size2
real_bsize = bsize * beam_size
sos_emb = self.get_sos_emb(inpute)
isize = sos_emb.size(-1)
sqrt_isize = sqrt(isize)
if length_penalty > 0.0:
lpv = sos_emb.new_ones(real_bsize, 1)
lpv_base = 6.0 ** length_penalty
out = sos_emb * sqrt_isize
if self.pemb is not None:
out = out + self.pemb.get_pos(0)
if self.drop is not None:
out = self.drop(out)
states = {}
for _tmp, net in enumerate(self.nets):
out, _state = net(inpute, None, inputc, src_pad_mask, context_mask, None, out)
states[_tmp] = _state
if self.out_normer is not None:
out = self.out_normer(out)
out = self.lsm(self.classifier(out))
scores, wds = out.topk(beam_size, dim=-1)
scores = scores.squeeze(1)
sum_scores = scores
wds = wds.view(real_bsize, 1)
trans = wds
done_trans = wds.view(bsize, beam_size).eq(2)
inpute = inpute.repeat(1, beam_size, 1).view(real_bsize, seql, isize)
_src_pad_mask = None if src_pad_mask is None else src_pad_mask.repeat(1, beam_size, 1).view(real_bsize, 1, seql)
_cbsize, _cseql = inputc[0].size()[:2]
_creal_bsize = _cbsize * beam_size
_context_mask = [None if cu is None else cu.repeat(1, beam_size, 1).view(_creal_bsize, 1, _cseql) for cu in context_mask]
_inputc = [inputu.repeat(1, beam_size, 1).view(_creal_bsize, _cseql, isize) for inputu in inputc]
for key, value in states.items():
states[key] = repeat_bsize_for_beam_tensor(value, beam_size)
for step in range(1, max_len):
out = self.wemb(wds) * sqrt_isize
if self.pemb is not None:
out = out + self.pemb.get_pos(step)
if self.drop is not None:
out = self.drop(out)
for _tmp, net in enumerate(self.nets):
out, _state = net(inpute, states[_tmp], _inputc, _src_pad_mask, None, _context_mask, out)
states[_tmp] = _state
if self.out_normer is not None:
out = self.out_normer(out)
out = self.lsm(self.classifier(out)).view(bsize, beam_size, -1)
_scores, _wds = out.topk(beam_size, dim=-1)
_scores = (_scores.masked_fill(done_trans.unsqueeze(2).expand(bsize, beam_size, beam_size), 0.0) + sum_scores.unsqueeze(2).expand(bsize, beam_size, beam_size))
if length_penalty > 0.0:
lpv = lpv.masked_fill(~done_trans.view(real_bsize, 1), ((step + 6.0) ** length_penalty) / lpv_base)
if clip_beam and (length_penalty > 0.0):
scores, _inds = (_scores.view(real_bsize, beam_size) / lpv.expand(real_bsize, beam_size)).view(bsize, beam_size2).topk(beam_size, dim=-1)
_tinds = (_inds + torch.arange(0, bsizeb2, beam_size2, dtype=_inds.dtype, device=_inds.device).unsqueeze(1).expand_as(_inds)).view(real_bsize)
sum_scores = _scores.view(bsizeb2).index_select(0, _tinds).view(bsize, beam_size)
else:
scores, _inds = _scores.view(bsize, beam_size2).topk(beam_size, dim=-1)
_tinds = (_inds + torch.arange(0, bsizeb2, beam_size2, dtype=_inds.dtype, device=_inds.device).unsqueeze(1).expand_as(_inds)).view(real_bsize)
sum_scores = scores
wds = _wds.view(bsizeb2).index_select(0, _tinds).view(real_bsize, 1)
_inds = (_inds // beam_size + torch.arange(0, real_bsize, beam_size, dtype=_inds.dtype, device=_inds.device).unsqueeze(1).expand_as(_inds)).view(real_bsize)
trans = torch.cat((trans.index_select(0, _inds), wds.masked_fill(done_trans.view(real_bsize, 1), 0) if fill_pad else wds), 1)
done_trans = (done_trans.view(real_bsize).index_select(0, _inds) | wds.eq(2).squeeze(1)).view(bsize, beam_size)
_done = False
if length_penalty > 0.0:
lpv = lpv.index_select(0, _inds)
elif (not return_all) and done_trans.select(1, 0).int().sum().item() == bsize:
_done = True
if _done or (done_trans.int().sum().item() == real_bsize):
break
for key, value in states.items():
states[key] = value.index_select(0, _inds)
if (not clip_beam) and (length_penalty > 0.0):
scores = scores / lpv.view(bsize, beam_size)
scores, _inds = scores.topk(beam_size, dim=-1)
_inds = (_inds + torch.arange(0, real_bsize, beam_size, dtype=_inds.dtype, device=_inds.device).unsqueeze(1).expand_as(_inds)).view(real_bsize)
trans = trans.view(real_bsize, -1).index_select(0, _inds).view(bsize, beam_size, -1)
if return_all:
return trans, scores
else:
return trans.view(bsize, beam_size, -1).select(1, 0)
def beam_decode(self,
inpute,
src_pad_mask=None,
beam_size=8,
max_len=512,
length_penalty=0.0,
return_all=False,
clip_beam=False,
fill_pad=False):
bsize, seql, isize = inpute[0].size()
beam_size2 = beam_size * beam_size
bsizeb2 = bsize * beam_size2
real_bsize = bsize * beam_size
sqrt_isize = sqrt(isize)
if length_penalty > 0.0:
# lpv: length penalty vector for each beam (bsize * beam_size, 1)
lpv = inpute[0].new_ones(real_bsize, 1)
lpv_base = 6.0**length_penalty
states = {}
outs = []
for _inum, (model, inputu) in enumerate(zip(self.nets, inpute)):
out = model.get_sos_emb(inputu) * sqrt_isize
if model.pemb is not None:
out = out + model.pemb.get_pos(0)
if model.drop is not None:
out = model.drop(out)
states[_inum] = {}
for _tmp, net in enumerate(model.nets):
out, _state = net(inputu, None, src_pad_mask, out, 1)
states[_inum][_tmp] = _state
if model.out_normer is not None:
out = model.out_normer(out)
# outs: [(bsize, 1, nwd)]
outs.append(model.classifier(out).softmax(dim=-1))
out = torch.stack(outs).mean(0).log()
# scores: (bsize, 1, beam_size) => (bsize, beam_size)
# wds: (bsize * beam_size, 1)
# trans: (bsize * beam_size, 1)
scores, wds = out.topk(beam_size, dim=-1)
scores = scores.squeeze(1)
sum_scores = scores
wds = wds.view(real_bsize, 1)
trans = wds
# done_trans: (bsize, beam_size)
done_trans = wds.view(bsize, beam_size).eq(2)
# inpute: (bsize, seql, isize) => (bsize * beam_size, seql, isize)
inpute = [
inputu.repeat(1, beam_size, 1).view(real_bsize, seql, isize)
for inputu in inpute
]
# _src_pad_mask: (bsize, 1, seql) => (bsize * beam_size, 1, seql)
_src_pad_mask = None if src_pad_mask is None else src_pad_mask.repeat(
1, beam_size, 1).view(real_bsize, 1, seql)
# states[i][j]: (bsize, 1, isize) => (bsize * beam_size, 1, isize)
for key, value in states.items():
for _key, _value in value.items():
value[_key] = repeat_bsize_for_beam_tensor(_value, beam_size)
for step in range(2, max_len + 1):
outs = []
for _inum, (model, inputu) in enumerate(zip(self.nets, inpute)):
out = model.wemb(wds) * sqrt_isize
if model.pemb is not None:
out = out + model.pemb.get_pos(step - 1)
if model.drop is not None:
out = model.drop(out)
for _tmp, net in enumerate(model.nets):
out, _state = net(inputu, states[_inum][_tmp],
_src_pad_mask, out, step)
states[_inum][_tmp] = _state
if model.out_normer is not None:
out = model.out_normer(out)
# outs: [(bsize, beam_size, nwd)...]
outs.append(
model.classifier(out).softmax(dim=-1).view(
bsize, beam_size, -1))
out = torch.stack(outs).mean(0).log()
# find the top k ** 2 candidates and calculate route scores for them
# _scores: (bsize, beam_size, beam_size)
# done_trans: (bsize, beam_size)
# scores: (bsize, beam_size)
# _wds: (bsize, beam_size, beam_size)
# mask_from_done_trans: (bsize, beam_size) => (bsize, beam_size * beam_size)
# added_scores: (bsize, 1, beam_size) => (bsize, beam_size, beam_size)
_scores, _wds = out.topk(beam_size, dim=-1)
_scores = (_scores.masked_fill(
done_trans.unsqueeze(2).expand(bsize, beam_size, beam_size),
0.0) + scores.unsqueeze(2).expand(bsize, beam_size, beam_size))
if length_penalty > 0.0:
lpv = lpv.masked_fill(1 - done_trans.view(real_bsize, 1),
((step + 5.0)**length_penalty) /
lpv_base)
# clip from k ** 2 candidate and remain the top-k for each path
# scores: (bsize, beam_size * beam_size) => (bsize, beam_size)
# _inds: indexes for the top-k candidate (bsize, beam_size)
if clip_beam and (length_penalty > 0.0):
scores, _inds = (_scores.view(real_bsize, beam_size) /
lpv.expand(real_bsize, beam_size)).view(
bsize, beam_size2).topk(beam_size, dim=-1)
_tinds = (_inds + torch.arange(
0,
bsizeb2,
beam_size2,
dtype=_inds.dtype,
device=_inds.device).unsqueeze(1).expand_as(_inds)
).view(real_bsize)
sum_scores = _scores.view(bsizeb2).index_select(
0, _tinds).view(bsize, beam_size)
else:
scores, _inds = _scores.view(bsize, beam_size2).topk(beam_size,
dim=-1)
_tinds = (_inds + torch.arange(
0,
bsizeb2,
beam_size2,
dtype=_inds.dtype,
device=_inds.device).unsqueeze(1).expand_as(_inds)
).view(real_bsize)
sum_scores = scores
# select the top-k candidate with higher route score and update translation record
# wds: (bsize, beam_size, beam_size) => (bsize * beam_size, 1)
wds = _wds.view(bsizeb2).index_select(0,
_tinds).view(real_bsize, 1)
# reduces indexes in _inds from (beam_size ** 2) to beam_size
# thus the fore path of the top-k candidate is pointed out
# _inds: indexes for the top-k candidate (bsize, beam_size)
_inds = (
_inds // beam_size +
torch.arange(0,
real_bsize,
beam_size,
dtype=_inds.dtype,
device=_inds.device).unsqueeze(1).expand_as(_inds)
).view(real_bsize)
# select the corresponding translation history for the top-k candidate and update translation records
# trans: (bsize * beam_size, nquery) => (bsize * beam_size, nquery + 1)
trans = torch.cat(
(trans.index_select(0, _inds),
wds.masked_fill(done_trans.view(real_bsize, 1), pad_id)
if fill_pad else wds), 1)
done_trans = (done_trans.view(real_bsize).index_select(0, _inds)
| wds.eq(2).squeeze(1)).view(bsize, beam_size)
# check early stop for beam search
# done_trans: (bsize, beam_size)
# scores: (bsize, beam_size)
_done = False
if length_penalty > 0.0:
lpv = lpv.index_select(0, _inds)
elif (not return_all) and all_done(done_trans.select(1, 0), bsize):
_done = True
# check beam states(done or not)
if _done or all_done(done_trans, real_bsize):
break
# update the corresponding hidden states
# states[i][j]: (bsize * beam_size, nquery, isize)
# _inds: (bsize, beam_size) => (bsize * beam_size)
for key, value in states.items():
for _key, _value in value.items():
value[_key] = _value.index_select(0, _inds)
# if length penalty is only applied in the last step, apply length penalty
if (not clip_beam) and (length_penalty > 0.0):
scores = scores / lpv.view(bsize, beam_size)
scores, _inds = scores.topk(beam_size, dim=-1)
_inds = (
_inds +
torch.arange(0,
real_bsize,
beam_size,
dtype=_inds.dtype,
device=_inds.device).unsqueeze(1).expand_as(_inds)
).view(real_bsize)
trans = trans.view(real_bsize, -1).index_select(0, _inds).view(
bsize, beam_size, -1)
if return_all:
return trans, scores
else:
return trans.view(bsize, beam_size, -1).select(1, 0)
def beam_decode(self,
inpute,
src_pad_mask=None,
beam_size=8,
max_len=512,
length_penalty=0.0,
return_all=False,
clip_beam=False):
bsize, seql = inpute.size()[:2]
beam_size2 = beam_size * beam_size
bsizeb2 = bsize * beam_size2
real_bsize = bsize * beam_size
sos_emb = self.get_sos_emb(inpute)
isize = sos_emb.size(-1)
sqrt_isize = sqrt(isize)
if length_penalty > 0.0:
# lpv: length penalty vector for each beam (bsize * beam_size, 1)
lpv = sos_emb.new_ones(real_bsize, 1)
lpv_base = 6.0**length_penalty
out = sos_emb * sqrt_isize + self.pemb.get_pos(0)
if self.drop is not None:
out = self.drop(out)
states = {}
attns = []
for _tmp, net in enumerate(self.nets):
out, _attn, _state = net(inpute, None, src_pad_mask, None, out,
True)
states[_tmp] = _state
attns.append(_attn)
if self.out_normer is not None:
out = self.out_normer(out)
attns = torch.cat(attns, dim=1).permute(0, 2, 3, 1)
_asize = attns.size()
out = torch.cat([
out,
attns.contiguous().view(-1, _asize[-1]).mv(
self.tattn_w.softmax(dim=0) if self.tattn_drop is None else
self.tattn_drop(self.tattn_w).softmax(dim=0)).view(
_asize[:-1]).bmm(inpute)
],
dim=-1)
# out: (bsize, 1, nwd)
out = self.lsm(self.classifier(out))
# scores: (bsize, 1, beam_size) => (bsize, beam_size)
# wds: (bsize * beam_size, 1)
# trans: (bsize * beam_size, 1)
scores, wds = out.topk(beam_size, dim=-1)
scores = scores.squeeze(1)
sum_scores = scores
wds = wds.view(real_bsize, 1)
trans = wds
# done_trans: (bsize, beam_size)
done_trans = wds.view(bsize, beam_size).eq(2)
# inpute: (bsize, seql, isize) => (bsize * beam_size, seql, isize)
inpute = inpute.repeat(1, beam_size, 1).view(real_bsize, seql, isize)
# _src_pad_mask: (bsize, 1, seql) => (bsize * beam_size, 1, seql)
_src_pad_mask = None if src_pad_mask is None else src_pad_mask.repeat(
1, beam_size, 1).view(real_bsize, 1, seql)
# states[i]: (bsize, 1, isize) => (bsize * beam_size, 1, isize)
for key, value in states.items():
states[key] = repeat_bsize_for_beam_tensor(value, beam_size)
for step in range(1, max_len):
out = self.wemb(wds) * sqrt_isize + self.pemb.get_pos(step)
if self.drop is not None:
out = self.drop(out)
attns = []
for _tmp, net in enumerate(self.nets):
out, _attn, _state = net(inpute, states[_tmp], _src_pad_mask,
None, out, True)
states[_tmp] = _state
attns.append(_attn)
if self.out_normer is not None:
out = self.out_normer(out)
attns = torch.cat(attns, dim=1).permute(0, 2, 3, 1)
_asize = attns.size()
out = torch.cat([
out,
attns.contiguous().view(-1, _asize[-1]).mv(
self.tattn_w.softmax(dim=0) if self.tattn_drop is None else
self.tattn_drop(self.tattn_w).softmax(dim=0)).view(
_asize[:-1]).bmm(inpute)
],
dim=-1)
# out: (bsize, beam_size, nwd)
out = self.lsm(self.classifier(out)).view(bsize, beam_size, -1)
# find the top k ** 2 candidates and calculate route scores for them
# _scores: (bsize, beam_size, beam_size)
# done_trans: (bsize, beam_size)
# scores: (bsize, beam_size)
# _wds: (bsize, beam_size, beam_size)
# mask_from_done_trans: (bsize, beam_size) => (bsize, beam_size * beam_size)
# added_scores: (bsize, 1, beam_size) => (bsize, beam_size, beam_size)
_scores, _wds = out.topk(beam_size, dim=-1)
_scores = (
_scores.masked_fill(
done_trans.unsqueeze(2).expand(bsize, beam_size,
beam_size), 0.0) +
sum_scores.unsqueeze(2).expand(bsize, beam_size, beam_size))
if length_penalty > 0.0:
lpv = lpv.masked_fill(1 - done_trans.view(real_bsize, 1),
((step + 6.0)**length_penalty) /
lpv_base)
# clip from k ** 2 candidate and remain the top-k for each path
# scores: (bsize, beam_size * beam_size) => (bsize, beam_size)
# _inds: indexes for the top-k candidate (bsize, beam_size)
if clip_beam and (length_penalty > 0.0):
scores, _inds = (_scores.view(real_bsize, beam_size) /
lpv.expand(real_bsize, beam_size)).view(
bsize, beam_size2).topk(beam_size, dim=-1)
_tinds = (_inds + torch.arange(
0,
bsizeb2,
beam_size2,
dtype=_inds.dtype,
device=_inds.device).unsqueeze(1).expand_as(_inds)
).view(real_bsize)
sum_scores = _scores.view(bsizeb2).index_select(
0, _tinds).view(bsize, beam_size)
else:
scores, _inds = _scores.view(bsize, beam_size2).topk(beam_size,
dim=-1)
_tinds = (_inds + torch.arange(
0,
bsizeb2,
beam_size2,
dtype=_inds.dtype,
device=_inds.device).unsqueeze(1).expand_as(_inds)
).view(real_bsize)
sum_scores = scores
# select the top-k candidate with higher route score and update translation record
# wds: (bsize, beam_size, beam_size) => (bsize * beam_size, 1)
wds = _wds.view(bsizeb2).index_select(0,
_tinds).view(real_bsize, 1)
# reduces indexes in _inds from (beam_size ** 2) to beam_size
# thus the fore path of the top-k candidate is pointed out
# _inds: indexes for the top-k candidate (bsize, beam_size)
_inds = (
_inds / beam_size +
torch.arange(0,
real_bsize,
beam_size,
dtype=_inds.dtype,
device=_inds.device).unsqueeze(1).expand_as(_inds)
).view(real_bsize)
# select the corresponding translation history for the top-k candidate and update translation records
# trans: (bsize * beam_size, nquery) => (bsize * beam_size, nquery + 1)
trans = torch.cat((trans.index_select(0, _inds), wds), 1)
done_trans = (done_trans.view(real_bsize).index_select(0, _inds) +
wds.eq(2).squeeze(1)).gt(0).view(bsize, beam_size)
# check early stop for beam search
# done_trans: (bsize, beam_size)
# scores: (bsize, beam_size)
_done = False
if length_penalty > 0.0:
lpv = lpv.index_select(0, _inds)
elif (not return_all) and done_trans.select(
1, 0).sum().item() == bsize:
_done = True
# check beam states(done or not)
if _done or (done_trans.sum().item() == real_bsize):
break
# update the corresponding hidden states
# states[i]: (bsize * beam_size, nquery, isize)
# _inds: (bsize, beam_size) => (bsize * beam_size)
for key, value in states.items():
states[key] = value.index_select(0, _inds)
# if length penalty is only applied in the last step, apply length penalty
if (not clip_beam) and (length_penalty > 0.0):
scores = scores / lpv.view(bsize, beam_size)
scores, _inds = scores.topk(beam_size, dim=-1)
_inds = (
_inds +
torch.arange(0,
real_bsize,
beam_size,
dtype=_inds.dtype,
device=_inds.device).unsqueeze(1).expand_as(_inds)
).view(real_bsize)
trans = trans.view(real_bsize, -1).index_select(0, _inds).view(
bsize, beam_size, -1)
if return_all:
return trans, scores
else:
return trans.view(bsize, beam_size, -1).select(1, 0)
def repeat_buffer(self, beam_size): if self.real_iK is not None: self.real_iK, self.real_iV = repeat_bsize_for_beam_tensor(self.real_iK, beam_size), repeat_bsize_for_beam_tensor(self.real_iV, beam_size)
def beam_decode(self,
inpute,
inputh,
src_pad_mask=None,
beam_size=8,
max_len=512,
length_penalty=0.0,
return_all=False,
clip_beam=clip_beam_with_lp,
fill_pad=False):
bsize, seql = inpute.size()[:2]
beam_size2 = beam_size * beam_size
bsizeb2 = bsize * beam_size2
real_bsize = bsize * beam_size
sos_emb = self.get_sos_emb(inpute)
isize = sos_emb.size(-1)
sqrt_isize = sqrt(isize)
if length_penalty > 0.0:
lpv = sos_emb.new_ones(real_bsize, 1)
lpv_base = 6.0**length_penalty
out = sos_emb * sqrt_isize
if self.pemb is not None:
out = out + self.pemb.get_pos(0)
if self.drop is not None:
out = self.drop(out)
out = self.out_normer(out)
states = {}
for _tmp, (net, inputu, inputhu) in enumerate(
zip(self.nets, inpute.unbind(dim=-1), inputh.unbind(dim=-1))):
out, _state = net(inputu, inputhu, None, src_pad_mask, None, out,
True)
states[_tmp] = _state
out = self.lsm(self.classifier(out))
scores, wds = out.topk(beam_size, dim=-1)
scores = scores.squeeze(1)
sum_scores = scores
wds = wds.view(real_bsize, 1)
trans = wds
done_trans = wds.view(bsize, beam_size).eq(2)
inputh = repeat_bsize_for_beam_tensor(inputh, beam_size)
self.repeat_cross_attn_buffer(beam_size)
_src_pad_mask = None if src_pad_mask is None else src_pad_mask.repeat(
1, beam_size, 1).view(real_bsize, 1, seql)
states = expand_bsize_for_beam(states, beam_size=beam_size)
for step in range(1, max_len):
out = self.wemb(wds) * sqrt_isize
if self.pemb is not None:
out = out + self.pemb.get_pos(step)
if self.drop is not None:
out = self.drop(out)
out = self.out_normer(out)
for _tmp, (net, inputu, inputhu) in enumerate(
zip(self.nets, inpute.unbind(dim=-1),
inputh.unbind(dim=-1))):
out, _state = net(inputu, inputhu, states[_tmp], _src_pad_mask,
None, out, True)
states[_tmp] = _state
out = self.lsm(self.classifier(out)).view(bsize, beam_size, -1)
_scores, _wds = out.topk(beam_size, dim=-1)
_done_trans_unsqueeze = done_trans.unsqueeze(2)
_scores = (
_scores.masked_fill(
_done_trans_unsqueeze.expand(bsize, beam_size, beam_size),
0.0) +
sum_scores.unsqueeze(2).repeat(1, 1, beam_size).masked_fill_(
select_zero_(_done_trans_unsqueeze.repeat(1, 1, beam_size),
-1, 0), -inf_default))
if length_penalty > 0.0:
lpv.masked_fill_(~done_trans.view(real_bsize, 1),
((step + 6.0)**length_penalty) / lpv_base)
if clip_beam and (length_penalty > 0.0):
scores, _inds = (_scores.view(real_bsize, beam_size) /
lpv.expand(real_bsize, beam_size)).view(
bsize, beam_size2).topk(beam_size, dim=-1)
_tinds = (_inds + torch.arange(
0,
bsizeb2,
beam_size2,
dtype=_inds.dtype,
device=_inds.device).unsqueeze(1).expand_as(_inds)
).view(real_bsize)
sum_scores = _scores.view(bsizeb2).index_select(
0, _tinds).view(bsize, beam_size)
else:
scores, _inds = _scores.view(bsize, beam_size2).topk(beam_size,
dim=-1)
_tinds = (_inds + torch.arange(
0,
bsizeb2,
beam_size2,
dtype=_inds.dtype,
device=_inds.device).unsqueeze(1).expand_as(_inds)
).view(real_bsize)
sum_scores = scores
wds = _wds.view(bsizeb2).index_select(0,
_tinds).view(real_bsize, 1)
_inds = (
_inds // beam_size +
torch.arange(0,
real_bsize,
beam_size,
dtype=_inds.dtype,
device=_inds.device).unsqueeze(1).expand_as(_inds)
).view(real_bsize)
trans = torch.cat(
(trans.index_select(0, _inds),
wds.masked_fill(done_trans.view(real_bsize, 1), pad_id)
if fill_pad else wds), 1)
done_trans = (done_trans.view(real_bsize).index_select(0, _inds)
| wds.eq(2).squeeze(1)).view(bsize, beam_size)
_done = False
if length_penalty > 0.0:
lpv = lpv.index_select(0, _inds)
elif (not return_all) and all_done(done_trans.select(1, 0), bsize):
_done = True
if _done or all_done(done_trans, real_bsize):
break
states = index_tensors(states, indices=_inds, dim=0)
if (not clip_beam) and (length_penalty > 0.0):
scores = scores / lpv.view(bsize, beam_size)
scores, _inds = scores.topk(beam_size, dim=-1)
_inds = (
_inds +
torch.arange(0,
real_bsize,
beam_size,
dtype=_inds.dtype,
device=_inds.device).unsqueeze(1).expand_as(_inds)
).view(real_bsize)
trans = trans.view(real_bsize, -1).index_select(0, _inds)
if return_all:
return trans.view(bsize, beam_size, -1), scores
else:
return trans.view(bsize, beam_size, -1).select(1, 0)
