fix context calculation at decode stesp

Author: atulkum

training_ptr_gen/decode.py

@@ -160,7 +160,7 @@ def beam_search(self, batch):
 
             final_dist, s_t, c_t, attn_dist, p_gen, coverage_t = self.model.decoder(y_t_1, s_t_1,
                                                         encoder_outputs, enc_padding_mask, c_t_1,
-                                                        extra_zeros, enc_batch_extend_vocab, coverage_t_1)
+                                                        extra_zeros, enc_batch_extend_vocab, coverage_t_1, steps)
 
             topk_log_probs, topk_ids = torch.topk(final_dist, config.beam_size * 2)
 

training_ptr_gen/eval.py

@@ -49,7 +49,7 @@ def eval_one_batch(self, batch):
             y_t_1 = dec_batch[:, di]  # Teacher forcing
             final_dist, s_t_1, c_t_1,attn_dist, p_gen, coverage = self.model.decoder(y_t_1, s_t_1,
                                                                 encoder_outputs, enc_padding_mask, c_t_1,
-                                                                extra_zeros, enc_batch_extend_vocab, coverage)
+                                                                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)

training_ptr_gen/model.py

@@ -145,7 +145,15 @@ def __init__(self):
         init_linear_wt(self.out2)
 
     def forward(self, y_t_1, s_t_1, encoder_outputs, enc_padding_mask,
-                c_t_1, extra_zeros, enc_batch_extend_vocab, coverage):
+                c_t_1, extra_zeros, enc_batch_extend_vocab, coverage, step):
+
+        if not self.training and step == 0:
+            h_decoder, c_decoder = s_t_1
+            s_t_hat = torch.cat((h_decoder.view(-1, config.hidden_dim),
+                                 c_decoder.view(-1, config.hidden_dim)), 1)  # B x 2*hidden_dim
+            c_t, _, coverage_next = self.attention_network(s_t_hat, encoder_outputs,
+                                                              enc_padding_mask, coverage)
+            coverage = coverage_next
 
         y_t_1_embd = self.embedding(y_t_1)
         x = self.x_context(torch.cat((c_t_1, y_t_1_embd), 1))
@@ -154,9 +162,12 @@ def forward(self, y_t_1, s_t_1, encoder_outputs, enc_padding_mask,
         h_decoder, c_decoder = s_t
         s_t_hat = torch.cat((h_decoder.view(-1, config.hidden_dim),
                              c_decoder.view(-1, config.hidden_dim)), 1)  # B x 2*hidden_dim
-        c_t, attn_dist, coverage = self.attention_network(s_t_hat, encoder_outputs,
+        c_t, attn_dist, coverage_next = self.attention_network(s_t_hat, encoder_outputs,
                                                           enc_padding_mask, coverage)
 
+        if self.training or step > 0:
+            coverage = coverage_next
+
         p_gen = None
         if config.pointer_gen:
             p_gen_input = torch.cat((c_t, s_t_hat, x), 1)  # B x (2*2*hidden_dim + emb_dim)

training_ptr_gen/train.py

@@ -91,7 +91,7 @@ def train_one_batch(self, batch):
             final_dist, s_t_1,  c_t_1, attn_dist, p_gen, coverage = self.model.decoder(y_t_1, s_t_1,
                                                         encoder_outputs, enc_padding_mask, c_t_1,
                                                         extra_zeros, enc_batch_extend_vocab,
-                                                                           coverage)
+                                                                           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)