def greedy_decode(model, src, src_mask, max_len):
memory = model.encode(src, src_mask)
ys = torch.ones(1,
1).fill_(src[0][-1].cpu().numpy().item()).type_as(src.data)
for i in range(max_len):
out = model.decode(
memory, src_mask, Variable(ys),
Variable(subsequent_mask(ys.size(1)).type_as(src.data)))
prob = model.generator(out[:, -1])
_, next_word = torch.max(prob, dim=1)
next_word = next_word.data[0]
ys = torch.cat(
[ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=1)
return ys
def choose_options(model, memory, src, src_mask, ys):
out = model.decode(
memory, src_mask, Variable(ys[1]),
Variable(subsequent_mask(ys[1].size(1)).type_as(src.data)))
prob = model.generator(out[:, -1])
dict = {}
for j in range(prob.size()[-1]):
dict[j] = prob[0][j].item()
sort_dict = sorted(zip(dict.values(), dict.keys()), reverse=True)
options = sort_dict[:beam_search_number]
result = []
for i in range(beam_search_number):
result.append([
ys[0] + options[i][0],
torch.cat([
ys[1],
torch.ones(1, 1).type_as(src.data).fill_(options[i][1])
],
dim=1)
])
return result
def make_std_mask(tgt, pad):
"创建一个mask来隐藏填充和将来的单词"
tgt_mask = (tgt != pad).unsqueeze(-2)
tgt_mask = tgt_mask & Variable(
subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data))
return tgt_mask