def model_batch_step(self, batch, eval):
enc_batch, enc_padding_mask, enc_lens, enc_batch_extend_vocab, extra_zeros, c_t_1, coverage = \
get_input_from_batch(batch, use_cuda)
dec_batch, dec_padding_mask, max_dec_len, dec_lens_var, target_batch = \
get_output_from_batch(batch, use_cuda)
encoder_outputs, encoder_feature, encoder_hidden = self.model.encoder(enc_batch, enc_lens)
s_t_1 = self.model.reduce_state(encoder_hidden)
step_losses = []
step_decoded_idx = []
for di in range(min(max_dec_len, config.max_dec_steps)):
y_t_1 = dec_batch[:, di] # Teacher forcing
final_dist, s_t_1, c_t_1, attn_dist, p_gen, next_coverage = \
self.model.decoder(y_t_1, s_t_1,
encoder_outputs,
encoder_feature,
enc_padding_mask, c_t_1,
extra_zeros,
enc_batch_extend_vocab,
coverage, di)
if eval:
_, top_idx = final_dist.topk(1)
step_decoded_idx.append(top_idx)
target = target_batch[:, di]
gold_probs = torch.gather(final_dist, 1, target.unsqueeze(1)).squeeze()
step_loss = -torch.log(gold_probs + config.eps)
if config.is_coverage:
step_coverage_loss = torch.sum(torch.min(attn_dist, coverage), 1)
step_loss = step_loss + config.cov_loss_wt * step_coverage_loss
coverage = next_coverage
step_mask = dec_padding_mask[:, di]
step_loss = step_loss * step_mask
step_losses.append(step_loss)
sum_losses = torch.sum(torch.stack(step_losses, 1), 1)
batch_avg_loss = sum_losses / dec_lens_var
loss = torch.mean(batch_avg_loss)
final_decoded_sentences = None
if eval:
final_decoded_sentences = torch.stack(step_decoded_idx, 2).squeeze(1)
print(final_decoded_sentences)
return loss, final_decoded_sentences
def train_one_batch(self, batch):
enc_batch, enc_padding_mask, enc_lens, enc_batch_extend_vocab, extra_zeros, c_t_1, coverage = \
get_input_from_batch(batch)
dec_batch, dec_padding_mask, max_dec_len, dec_lens_var, target_batch = \
get_output_from_batch(batch)
self.optimizer.zero_grad()
encoder_outputs, encoder_feature, encoder_hidden = self.model.encoder(
enc_batch, enc_lens)
s_t_1 = self.model.reduce_state(encoder_hidden)
step_losses = []
for di in range(min(max_dec_len, config.max_dec_steps)):
y_t_1 = dec_batch[:, di] # Teacher forcing
final_dist, s_t_1, c_t_1, attn_dist, p_gen, next_coverage = self.model.decoder(
y_t_1, s_t_1, encoder_outputs, encoder_feature,
enc_padding_mask, c_t_1, extra_zeros, enc_batch_extend_vocab,
coverage, di)
target = target_batch[:, di]
gold_probs = torch.gather(final_dist, 1,
target.unsqueeze(1)).squeeze()
step_loss = -torch.log(gold_probs + config.eps)
if config.is_coverage:
step_coverage_loss = torch.sum(torch.min(attn_dist, coverage),
1)
step_loss = step_loss + config.cov_loss_wt * step_coverage_loss
coverage = next_coverage
step_mask = dec_padding_mask[:, di]
step_loss = step_loss * step_mask
step_losses.append(step_loss)
sum_losses = torch.sum(torch.stack(step_losses, 1), 1)
batch_avg_loss = sum_losses / dec_lens_var
loss = torch.mean(batch_avg_loss)
loss.backward()
self.norm = clip_grad_norm_(self.model.encoder.parameters(),
config.max_grad_norm)
clip_grad_norm_(self.model.decoder.parameters(), config.max_grad_norm)
clip_grad_norm_(self.model.reduce_state.parameters(),
config.max_grad_norm)
self.optimizer.step()
return loss.item()
def eval_one_batch(self, batch):
batch_size = batch.batch_size
enc_batch, enc_padding_mask, enc_lens, enc_batch_extend_vocab, extra_zeros, c_t_1, coverage = \
get_input_from_batch(batch, use_cuda)
dec_batch, dec_padding_mask, max_dec_len, dec_lens_var, target_batch = \
get_output_from_batch(batch, use_cuda)
encoder_outputs, encoder_feature, encoder_hidden = self.model.encoder(
enc_batch, enc_lens)
s_t_1 = self.model.reduce_state(encoder_hidden)
step_losses = []
output_ids = []
y_t_1 = torch.ones(batch_size, dtype=torch.long) * self.vocab.word2id(
data.START_DECODING)
if config.use_gpu:
y_t_1 = y_t_1.cuda()
for _ in range(batch_size):
output_ids.append([])
step_losses.append([])
for di in range(min(max_dec_len, config.max_dec_steps)):
#y_t_1 = dec_batch[:, di] # Teacher forcing
final_dist, s_t_1, c_t_1, attn_dist, p_gen, next_coverage = self.model.decoder(
y_t_1, s_t_1, encoder_outputs, encoder_feature,
enc_padding_mask, c_t_1, extra_zeros, enc_batch_extend_vocab,
coverage, di)
target = target_batch[:, di]
gold_probs = torch.gather(final_dist, 1,
target.unsqueeze(1)).squeeze()
step_loss = -torch.log(gold_probs + config.eps) #NLL
if config.is_coverage:
step_coverage_loss = torch.sum(torch.min(attn_dist, coverage),
1)
step_loss = step_loss + config.cov_loss_wt * step_coverage_loss
coverage = next_coverage
step_mask = dec_padding_mask[:, di]
step_loss = step_loss * step_mask
# Move on to the next token
_, idx = torch.max(final_dist, 1)
idx = idx.reshape(batch_size, -1).squeeze()
y_t_1 = idx
for i, pred in enumerate(y_t_1):
if not pred.item() == data.PAD_TOKEN:
output_ids[i].append(pred.item())
for i, loss in enumerate(step_loss):
step_losses[i].append(step_loss[i])
# Obtain the original and predicted summaries
original_abstracts = batch.original_abstracts_sents
predicted_abstracts = [
data.outputids2words(ids, self.vocab, None) for ids in output_ids
]
# Compute the batched loss
batched_losses = self.compute_batched_loss(step_losses,
original_abstracts,
predicted_abstracts)
losses = torch.stack(batched_losses)
losses = losses / dec_lens_var
loss = torch.mean(losses)
return loss.item()
def beam_search(self, batch):
# batch should have only one example
enc_batch, enc_padding_mask, enc_lens, \
enc_batch_extend_vocab, extra_zeros, \
c_t_0, coverage_t_0 = get_input_from_batch(batch, use_cuda)
encoder_outputs, encoder_feature, encoder_hidden = self.model.encoder(
enc_batch, enc_lens)
s_t_0 = self.model.reduce_state(encoder_hidden)
dec_h, dec_c = s_t_0 # 1 x 2*hidden_size
dec_h, dec_c = dec_h.squeeze(), dec_c.squeeze()
# decoder batch preparation, it has beam_size example initially everything is repeated
beams = [
Beam(tokens=[self.vocab.word2id(data.START_DECODING)],
log_probs=[0.0],
state=(dec_h[0], dec_c[0]),
context=c_t_0[0],
coverage=(coverage_t_0[0] if config.is_coverage else None))
for _ in range(config.beam_size)
]
results = []
steps = 0
while steps < config.max_dec_steps and len(results) < config.beam_size:
latest_tokens = [h.latest_token for h in beams]
latest_tokens = [t if t < self.vocab.size() else self.vocab.word2id(data.UNKNOWN_TOKEN) \
for t in latest_tokens]
y_t_1 = torch.tensor(latest_tokens, dtype=torch.long)
if use_cuda:
y_t_1 = y_t_1.cuda()
all_state_h, all_state_c, all_context = [], [], []
for h in beams:
state_h, state_c = h.state
all_state_h.append(state_h)
all_state_c.append(state_c)
all_context.append(h.context)
s_t_1 = (torch.stack(all_state_h,
0).unsqueeze(0), torch.stack(all_state_c,
0).unsqueeze(0))
c_t_1 = torch.stack(all_context, 0)
coverage_t_1 = None
if config.is_coverage:
all_coverage = [h.coverage for h in beams]
coverage_t_1 = torch.stack(all_coverage, 0)
final_dist, s_t, c_t, attn_dist, p_gen, coverage_t = \
self.model.decoder(y_t_1, s_t_1,
encoder_outputs, encoder_feature,
enc_padding_mask, c_t_1,
extra_zeros, enc_batch_extend_vocab,
coverage_t_1, steps)
log_probs = torch.log(final_dist)
topk_log_probs, topk_ids = torch.topk(log_probs,
config.beam_size * 2)
dec_h, dec_c = s_t
dec_h, dec_c = dec_h.squeeze(), dec_c.squeeze()
all_beams = []
num_orig_beams = 1 if steps == 0 else len(beams)
for i in range(num_orig_beams):
h = beams[i]
state_i = (dec_h[i], dec_c[i])
context_i = c_t[i]
coverage_i = (coverage_t[i] if config.is_coverage else None)
for j in range(config.beam_size *
2): # for each of the top 2*beam_size hyps:
new_beam = h.extend(token=topk_ids[i, j].item(),
log_prob=topk_log_probs[i, j].item(),
state=state_i,
context=context_i,
coverage=coverage_i)
all_beams.append(new_beam)
beams = []
for h in self.sort_beams(all_beams):
if h.latest_token == self.vocab.word2id(data.STOP_DECODING):
if steps >= config.min_dec_steps:
results.append(h)
else:
if h.latest_token != self.vocab.word2id(
data.UNKNOWN_TOKEN):
beams.append(h)
if len(beams) == config.beam_size or len(
results) == config.beam_size:
break
steps += 1
if len(results) == 0:
results = beams
beams_sorted = self.sort_beams(results)
return beams_sorted[0]
def train_one_batch_pg(self, batch):
batch_size = batch.batch_size
enc_batch, enc_padding_mask, enc_lens, enc_batch_extend_vocab, extra_zeros, c_t_1, coverage = \
get_input_from_batch(batch, config.use_gpu)
dec_batch, dec_padding_mask, max_dec_len, dec_lens_var, target_batch = \
get_output_from_batch(batch, config.use_gpu)
self.optimizer.zero_grad()
encoder_outputs, encoder_feature, encoder_hidden = self.model.encoder(
enc_batch, enc_lens)
s_t_1 = self.model.reduce_state(encoder_hidden)
step_losses = []
output_ids = []
# Begin with START symbol
y_t_1 = torch.ones(batch_size, dtype=torch.long) * self.vocab.word2id(
data.START_DECODING)
if config.use_gpu:
y_t_1 = y_t_1.cuda()
for _ in range(batch_size):
output_ids.append([])
step_losses.append([])
for di in range(min(max_dec_len, config.max_dec_steps)):
#y_t_1 = dec_batch[:, di] # Teacher forcing
final_dist, s_t_1, c_t_1, attn_dist, p_gen, next_coverage = self.model.decoder(
y_t_1, s_t_1, encoder_outputs, encoder_feature,
enc_padding_mask, c_t_1, extra_zeros, enc_batch_extend_vocab,
coverage, di)
target = target_batch[:, di]
gold_probs = torch.gather(final_dist, 1,
target.unsqueeze(1)).squeeze()
step_loss = -torch.log(gold_probs + config.eps) # NLL
step_mask = dec_padding_mask[:, di]
step_loss = step_loss * step_mask
# Move on to next token
_, idx = torch.max(final_dist, 1)
idx = idx.reshape(batch_size, -1).squeeze()
y_t_1 = idx
for i, pred in enumerate(y_t_1):
if not pred.item() == data.PAD_TOKEN:
output_ids[i].append(pred.item())
for i, loss in enumerate(step_loss):
step_losses[i].append(step_loss[i])
# Obtain the original and predicted summaries
original_abstracts = batch.original_abstracts_sents
predicted_abstracts = [
data.outputids2words(ids, self.vocab, None) for ids in output_ids
]
# Compute the batched loss
batched_losses = self.compute_batched_sentence_loss(
step_losses, original_abstracts, predicted_abstracts)
#batched_losses = Variable(batched_losses, requires_grad=True)
losses = torch.stack(batched_losses)
losses = losses / dec_lens_var
loss = torch.mean(losses)
loss.backward()
self.norm = clip_grad_norm_(self.model.encoder.parameters(),
config.max_grad_norm)
clip_grad_norm_(self.model.decoder.parameters(), config.max_grad_norm)
clip_grad_norm_(self.model.reduce_state.parameters(),
config.max_grad_norm)
self.optimizer.step()
return loss.item()