def sample(self, batch, top_p=1, return_pair_level=False, return_preds=False):
inputs = unpack_batch(batch, self.use_cuda)
src, tgt_in, tgt_out, turn_ids = \
inputs['src'], inputs['tgt_in'], inputs['tgt_out'], inputs['turn_ids']
bg = inputs.get('bg', None)
src_mask = src.eq(constant.PAD_ID)
bg_mask = bg.eq(constant.PAD_ID) if bg is not None else None
if not return_preds:
self.model.eval()
batch_size = src.size(0)
preds = self.model.sample(src, src_mask, turn_ids, top_p=top_p, bg=bg, bg_mask=bg_mask, return_pair_level=return_pair_level)
preds, nll = preds
if return_preds:
return preds, nll
pred_seqs = [[self.vocab['id2word'][id_] for id_ in ids] for ids in preds] # unmap to tokens
pred_seqs = utils.prune_decoded_seqs(pred_seqs)
if self.args['lambda2'] > 0:
pred_in, pred_out, pred_out_char, pred_text = samples_to_in_out(preds, self.vocab, src.device, src.size(0))
with torch.no_grad():
t_r, s_r = get_reward(tgt_in, tgt_out, inputs['tgt_out_char'], inputs['tgt_text'], pred_in, pred_out, pred_out_char, pred_text, inputs, self.teacher)
else:
t_r = s_r = 0
return pred_seqs, t_r, s_r
def predict(self, batch, unsort=True):
inputs, orig_idx = unpack_batch(batch, self.use_cuda)
src, src_mask, tgt, tgt_mask = inputs
self.model.eval()
batch_size = src.size(0)
preds, _ = self.model.predict(src, src_mask, self.args['beam_size'])
pred_seqs = [self.vocab.unmap(ids) for ids in preds] # unmap to tokens
pred_seqs = utils.prune_decoded_seqs(pred_seqs)
pred_tokens = ["".join(seq)
for seq in pred_seqs] # join chars to be tokens
if unsort:
pred_tokens = utils.unsort(pred_tokens, orig_idx)
return pred_tokens
def predict(self, batch, beam_size=1):
inputs, orig_idx = unpack_batch(batch, self.use_cuda)
src, src_mask, tgt, tgt_mask, pos, edits = inputs
self.model.eval()
batch_size = src.size(0)
preds, edit_logits = self.model.predict(src, src_mask, pos=pos, beam_size=beam_size)
pred_seqs = [self.vocab['char'].unmap(ids) for ids in preds] # unmap to tokens
pred_seqs = utils.prune_decoded_seqs(pred_seqs)
pred_tokens = ["".join(seq) for seq in pred_seqs] # join chars to be tokens
pred_tokens = utils.unsort(pred_tokens, orig_idx)
if self.args.get('edit', False):
assert edit_logits is not None
edits = np.argmax(edit_logits.data.cpu().numpy(), axis=1).reshape([batch_size]).tolist()
edits = utils.unsort(edits, orig_idx)
else:
edits = None
return pred_tokens, edits