def hidden_dep_dec_func(self, prev_word_ids, hidden):
t1 = [
[1., 1.1, 0., 0.0, 0.0, 0.0],
[10., 0., 50., 0., 0., 0.] # this will be ignored
]
t2 = [[0.2, 0., 0.1, 0., 0., 0.], [0., 0., 0., 1.1, 0., 2.]]
t3 = [
[4., 99., 331., 10., 133., 53.], # should be ignored
[0., 0., 0., -3.0000, 0., 0.001]
]
self.state += 1
if self.state == 0:
hidden[0, 0] += 1.
return DecState(T.tensor(t1), rec_vals={"hidden": hidden})
if self.state == 1:
t2 = T.tensor(t2)
t2[0, :] += hidden[0, 0]
return DecState(t2, rec_vals={"hidden": T.tensor([[1.2], [2.]])})
if self.state == 2:
t3 = T.tensor(t3)
t3[0, :] += hidden[0, 0]
t3[1, :] += hidden[1, 0]
return DecState(t3, rec_vals={"hidden": T.tensor([[0.], [0.]])})
raise ValueError("The decoding func supports only 3 steps!")
def decode_beam(self, seqs, hidden, att_word_ids, init_z=None, **kwargs):
"""Function to be used in the beam search process.
:param seqs: [batch_size, 1]
:param hidden: [batch_size, hidden_dim]
:param att_word_ids: [batch_size, cat_rev_len]
:param init_z: [batch_size, z_dim]
"""
embds = self._embds(seqs)
mask = T.ones_like(seqs, dtype=T.float32)
if init_z is None:
init_z = hidden
word_log_probs, att_wts, \
hidden, cont, ptr_probs = self._decode(embds=embds, mask=mask,
extra_feat=init_z.unsqueeze(1),
hidden=hidden,
att_word_ids=att_word_ids,
**kwargs)
out = DecState(word_scores=word_log_probs,
rec_vals={"hidden": hidden, "cont": cont,
"init_z": init_z},
coll_vals={'copy_probs': ptr_probs.squeeze(-1),
'att_wts': att_wts.squeeze(1),
"att_word_ids": att_word_ids})
return out
def generate_summaries(self, batch, **kwargs):
"""Generates only summaries; simplified script for inference."""
self.model.eval()
revs = batch[ModelF.REV].to(self.device)
rev_lens = batch[ModelF.REV_LEN].to(self.device)
revs_mask = batch[ModelF.REV_MASK].to(self.device)
group_rev_indxs = batch[ModelF.GROUP_REV_INDXS].to(self.device)
if ModelF.GROUP_REV_INDXS_MASK in batch:
summ_rev_indxs_mask = batch[ModelF.GROUP_REV_INDXS_MASK].to(
self.device)
else:
summ_rev_indxs_mask = None
max_rev_len = revs.size(1)
summs_nr = group_rev_indxs.size(0)
with T.no_grad():
rev_embds = self.model._embds(revs)
rev_encs, rev_hiddens = self.model.encode(rev_embds, rev_lens)
att_keys = self.model.create_att_keys(rev_hiddens)
contxt_states = self.model.get_contxt_states(
rev_hiddens, rev_embds)
summ_att_keys, \
summ_att_vals, \
summ_att_mask = group_att_over_input(inp_att_keys=att_keys,
inp_att_vals=rev_hiddens,
inp_att_mask=revs_mask,
att_indxs=group_rev_indxs,
att_indxs_mask=summ_rev_indxs_mask)
summ_att_word_ids = revs[group_rev_indxs].view(summs_nr, -1)
c_mu_q, c_sigma_q, _ = self.model.get_q_c_mu_sigma(
contxt_states, revs_mask, group_rev_indxs, summ_rev_indxs_mask)
if self.min_sen_seq_len is not None:
min_lens = [self.min_sen_seq_len] * summs_nr
else:
min_lens = None
z_mu_p, z_sigma_p = self.model.get_p_z_mu_sigma(c_mu_q)
init_summ_dec_state = DecState(rec_vals={"hidden": z_mu_p})
summ_word_ids, \
summ_coll_vals = self.beamer(init_summ_dec_state,
min_lens=min_lens,
max_steps=max_rev_len,
att_keys=summ_att_keys,
att_values=summ_att_vals,
att_mask=summ_att_mask,
att_word_ids=summ_att_word_ids,
minimum=1, **kwargs)
return summ_word_ids
def test_hidden_dependent_output(self):
beam_size = 2
max_steps = 3
vocab = {0: "a", 1: "b", 2: "c", 3: "<pad>", 4: "<s>", 5: "<e>"}
exp_seqs = [[4, 1, 0, 5]]
init_hidden = T.tensor([[0.]], dtype=T.float32)
dec = Dec()
beam_decoder = Beamer(decoding_func=dec.hidden_dep_dec_func,
start_id=4,
end_id=5,
validate_dec_out=False,
n_best=beam_size,
beam_size=beam_size)
init_dec_state = DecState(rec_vals={"hidden": init_hidden})
act_seqs, _ = beam_decoder(init_dec_state, max_steps=max_steps)
self.assertTrue((exp_seqs == act_seqs))
def test_simple_output(self):
"""Hidden state independent test."""
beam_size = 2
max_steps = 3
vocab = {0: "a", 1: "b", 2: "c", 3: "<pad>", 4: "<s>", 5: "<e>"}
exp_seqs = [[4, 2, 5], [4, 0, 0, 5]]
init_hidden = T.tensor([[0., 0., 0.], [0., 0., 0.]], dtype=T.float32)
init_dec_state = DecState(rec_vals={"hidden": init_hidden})
dec = Dec()
beam_decoder = Beamer(decoding_func=dec.dummy_dec_func,
start_id=4,
beam_size=beam_size,
end_id=5,
validate_dec_out=False)
act_seqs, _ = beam_decoder(init_dec_state=init_dec_state,
max_steps=max_steps)
self.assertTrue((exp_seqs == act_seqs))
def decode(self, seq, tr_state=None, dummy=None, **kwargs):
"""BeamSearch or Sampler specific decoding function. Performs one step
decoding based on the current state in `tr_state`.
Args:
seq: [batch_size, 1]
tr_state: [batch_size, num_layers, curr_seq_len, model_dim]
dummy: it used for consistency with the beam search.
Returns:
DecState
"""
pos_offset = 0 if tr_state is None else tr_state.size(2)
tr_state = tr_state.transpose(1, 0) if tr_state is not None else None
word_scores, new_tr_state, \
mem_att_wts = self._decode(tgt=seq, tr_state=tr_state,
pos_offset=pos_offset, **kwargs)
# new_state: [num_layers, batch_size, curr_seq_len, model_dim]
tr_state = new_tr_state if tr_state is None else\
T.cat((tr_state, new_tr_state), dim=2)
tr_state = tr_state.transpose(1, 0)
out = DecState(word_scores=word_scores,
rec_vals={"tr_state": tr_state})
return out
def dummy_dec_func(self, prev_word_ids, hidden):
t1 = [
[1.1, 0., 1., 0., 0., 0.],
[0., 0., 0., 10., 10., 10.], # this will be ignored
[1.1, 1., 0., 0., 0., 0.],
[0., 0., 0., 10., 0., 0.] # this will be ignored
]
t2 = [[0., 1., 2., 0., 0., 0.], [0., 0., 0., 0., 0., 6.],
[1.1, 0., 0., 0., 0., 0.], [0., 1., 0., 0., 0., 0.]]
t3 = [[4., 9999., 3., 10., 133., 5.], [0., 0., 0., 0., 0., 1.],
[0., 0., 0., 0., 0., 1.], [0., 0., 0., 0., 0., 1.]]
self.state += 1
if self.state == 0:
word_scores = T.tensor(t1)
elif self.state == 1:
word_scores = T.tensor(t2)
elif self.state == 2:
word_scores = T.tensor(t3)
else:
raise ValueError("The decoding func supports only 3 steps!")
return DecState(word_scores=word_scores, rec_vals={"hidden": hidden})
def generate(self, batch, min_seq_len, max_seq_len, use_true_props=True):
"""Generates/decodes sequences from a conditional pmf.
Args:
batch: self-explanatory.
min_seq_len: minimum length of the generated sequence.
max_seq_len: maximum allowed length for generated sequences.
use_true_props: whether to use true property values or the ones
inferred by the plug-in network.
Returns:
word_ids: list of word ids.
prop_values: dict of lists.
"""
self.model.eval()
rev = batch[ModelF.REV].to(self.device)
rev_mask = batch[ModelF.REV_MASK].to(self.device)
group_rev_indxs = batch[ModelF.GROUP_REV_INDXS].to(self.device)
group_rev_indxs_mask = batch[ModelF.GROUP_REV_INDXS_MASK].to(
self.device)
bs = group_rev_indxs.size(0)
mem,\
mem_bin_mask = self.model.create_mem(rev=rev, rev_mask=rev_mask,
group_rev_indxs=group_rev_indxs,
group_rev_indxs_mask=group_rev_indxs_mask)
if use_true_props:
len_prop = batch[ModelF.LEN_PROP].to(self.device)
rating_prop = batch[ModelF.RATING_PROP].to(self.device)
rouge_prop = batch[ModelF.ROUGE_PROP].to(self.device)
pov_prop = batch[ModelF.POV_PROP].to(self.device)
prop_vals = {
ModelF.LEN_PROP: len_prop,
ModelF.RATING_PROP: rating_prop,
ModelF.ROUGE_PROP: rouge_prop,
ModelF.POV_PROP: pov_prop
}
else:
if hasattr(self.model, 'plugin'):
prop_vals, _ = self.model.plugin(mem, mem_bin_mask)
else:
prop_vals = {}
dummy = T.zeros(group_rev_indxs.size(0), device=self.device)
if min_seq_len is not None:
min_lens = [min_seq_len] * bs
else:
min_lens = None
with T.no_grad():
init_dec_state = DecState(rec_vals={"dummy": dummy})
word_ids, _ = self.gen_func(init_dec_state=init_dec_state,
max_steps=max_seq_len,
log_normalize=True,
min_lens=min_lens,
mem=mem,
mem_bin_mask=mem_bin_mask,
minimum=1,
**prop_vals)
prop_vals = {k: v.tolist() for k, v in prop_vals.items()}
return word_ids, prop_vals
def predict(self, batch, **kwargs):
"""Predicts summaries and reviews."""
self.model.eval()
revs = batch[ModelF.REV].to(self.device)
rev_lens = batch[ModelF.REV_LEN].to(self.device)
revs_mask = batch[ModelF.REV_MASK].to(self.device)
summ_rev_indxs = batch[ModelF.GROUP_REV_INDXS].to(self.device)
summ_rev_indxs_mask = batch[ModelF.GROUP_REV_INDXS_MASK].to(
self.device)
other_revs = batch[ModelF.OTHER_REV_INDXS].to(self.device)
other_revs_mask = batch[ModelF.OTHER_REV_INDXS_MASK].to(self.device)
rev_to_group_indx = batch[ModelF.REV_TO_GROUP_INDX].to(self.device)
bs = revs.size(0)
max_rev_len = revs.size(1)
summs_nr = summ_rev_indxs.size(0)
with T.no_grad():
rev_embds = self.model._embds(revs)
rev_encs, rev_hiddens = self.model.encode(rev_embds, rev_lens)
att_keys = self.model.create_att_keys(rev_hiddens)
contxt_states = self.model.get_contxt_states(
rev_hiddens, rev_embds)
summ_att_keys, \
summ_att_vals, \
summ_att_mask = group_att_over_input(inp_att_keys=att_keys,
inp_att_vals=rev_hiddens,
inp_att_mask=revs_mask,
att_indxs=summ_rev_indxs,
att_indxs_mask=summ_rev_indxs_mask)
summ_att_word_ids = revs[summ_rev_indxs].view(summs_nr, -1)
c_mu_q, c_sigma_q, _ = self.model.get_q_c_mu_sigma(
contxt_states, revs_mask, summ_rev_indxs, summ_rev_indxs_mask)
# DECODING OF SUMMARIES #
if self.min_sen_seq_len is not None:
min_lens = [self.min_sen_seq_len] * summs_nr
else:
min_lens = None
z_mu_p, z_sigma_p = self.model.get_p_z_mu_sigma(c_mu_q)
init_summ_dec_state = DecState(rec_vals={"hidden": z_mu_p})
summ_word_ids, \
summ_coll_vals = self.beamer(init_summ_dec_state,
min_lens=min_lens,
max_steps=max_rev_len,
att_keys=summ_att_keys,
att_values=summ_att_vals,
att_mask=summ_att_mask,
att_word_ids=summ_att_word_ids,
minimum=1, **kwargs)
# DECODING OF REVIEWS #
z_mu_q, \
z_sigma_q = self.model.get_q_z_mu_sigma(rev_encs,
c=c_mu_q[rev_to_group_indx])
# creating attention values for the reviewer
rev_att_keys, \
rev_att_vals, \
rev_att_vals_mask = group_att_over_input(inp_att_keys=att_keys,
inp_att_vals=rev_hiddens,
inp_att_mask=revs_mask,
att_indxs=other_revs,
att_indxs_mask=other_revs_mask)
rev_att_word_ids = revs[other_revs].view(bs, -1)
if self.min_sen_seq_len is not None:
min_lens = [self.min_sen_seq_len] * bs
else:
min_lens = None
init_rev_dec_state = DecState(rec_vals={"hidden": z_mu_q})
rev_word_ids, \
rev_coll_vals = self.beamer(init_rev_dec_state,
min_lens=min_lens,
max_steps=max_rev_len,
att_keys=rev_att_keys,
att_values=rev_att_vals,
att_mask=rev_att_vals_mask,
att_word_ids=rev_att_word_ids,
minimum=1, **kwargs)
return rev_word_ids, rev_coll_vals, summ_word_ids, summ_coll_vals
def gather(self, outputs, output_device):
"""Gathers outputs from replicas, aggregates them generically.
Gathers outputs from replicas, assumes that each element of the output
is a PyTorch tensor(1+dim), DecStates, or dictionary of scalars.
The first and second are concatenated along 0-dim, the last are
added and then divided by the number of replicas (mean of means).
It does not check that the outputs are consistent across replicas. At
the moment.
Args:
outputs: replica outputs.
output_device:
Returns:
aggregated output.
"""
def _coll_tensor(coll_list, out_tensor):
out_tensor = gather(outputs=[out_tensor],
target_device=self.output_device,
dim=0)
coll_list.append(out_tensor)
def _coll_dict(coll_dict, out_dict):
for k, v in out_dict.items():
if k not in coll_dict:
coll_dict[k] = []
coll_dict[k].append(v)
def _coll_dec_state(coll_dstate, out_dstate):
for attr_name, attr_val in out_dstate.__dict__.items():
if isinstance(attr_val, Tensor):
if getattr(coll_dstate, attr_name) is None:
setattr(coll_dstate, attr_name, [])
_coll_tensor(getattr(coll_dstate, attr_name), attr_val)
elif isinstance(attr_val, dict):
if getattr(coll_dstate, attr_name) is None:
setattr(coll_dstate, attr_name, dict())
_coll_dict(getattr(coll_dstate, attr_name), attr_val)
elif attr_val is None:
continue
else:
raise NotImplementedError
# minor reformatting
for indx, o in enumerate(outputs):
if not isinstance(o, (list, tuple)):
outputs[indx] = [o]
coll_outputs = [None for _ in range(len(outputs[0]))]
# collecting outputs
for output in outputs:
for indx, o in enumerate(output):
# statistics, such as loss, assumed to be of simple types
if isinstance(o, dict):
if coll_outputs[indx] is None:
coll_outputs[indx] = dict()
_coll_dict(coll_outputs[indx], o)
# tensors
elif isinstance(o, Tensor):
if coll_outputs[indx] is None:
coll_outputs[indx] = []
_coll_tensor(coll_outputs[indx], o)
# decoder state
elif isinstance(o, DecState):
if coll_outputs[indx] is None:
coll_outputs[indx] = DecState()
coll_dec_state = coll_outputs[indx]
_coll_dec_state(coll_dec_state, o)
else:
raise NotImplementedError
# aggregating outputs
for indx, o in enumerate(coll_outputs):
if isinstance(o, list):
coll_outputs[indx] = T.cat(o)
elif isinstance(o, dict):
coll_outputs[indx] = {k: sum(v) / len(v) for k, v in o.items()}
elif isinstance(o, DecState):
for attr_name, attr_val in o.__dict__.items():
if isinstance(attr_val, list):
setattr(o, attr_name, T.cat(attr_val))
elif isinstance(attr_val, dict):
for k, v in attr_val.items():
attr_val[k] = T.cat(v)
elif attr_val is None:
continue
else:
raise NotImplementedError
else:
raise NotImplementedError
if len(coll_outputs) == 1:
coll_outputs = coll_outputs[0]
return coll_outputs