reimplement beam_accum, but with a bug

onmt/translate/translator.py

@@ -303,8 +303,8 @@ def translate(
 
         if self.dump_beam:
             import json
-            json.dump(self.translator.beam_accum,
-                      codecs.open(self.dump_beam, 'w', 'utf-8'))
+            with codecs.open(self.dump_beam, 'w', 'utf-8') as f:
+                json.dump(self.beam_accum, f)
         return all_scores, all_predictions
 
     def sample_with_temperature(self, logits, sampling_temp, keep_topk):
@@ -833,8 +833,7 @@ def _translate_batch(self, batch, data):
             select_indices_array = []
             # Loop over the batch_size number of beam
             for j, b in enumerate(beam):
-                b.advance(out[j, :],
-                          beam_attn.data[j, :, :memory_lengths[j]])
+                b.advance(out[j, :], beam_attn.data[j, :, :memory_lengths[j]])
                 select_indices_array.append(
                     b.get_current_origin() + j * beam_size)
             select_indices = torch.cat(select_indices_array)
@@ -854,6 +853,17 @@ def _translate_batch(self, batch, data):
             results["scores"].append(scores)
             results["attention"].append(attn)
 
+            if self.beam_accum is not None:
+                self.beam_accum["beam_parent_ids"].append(
+                    [t.tolist() for t in b.prev_ks])
+                self.beam_accum["scores"].append([
+                    ["%4f" % s for s in t.tolist()]
+                    for t in b.all_scores][1:])
+                # ok, so what was the tgt_dict?
+                self.beam_accum["predicted_ids"].append(
+                    [[vocab.itos[i] for i in t.tolist()]
+                     for t in b.next_ys][1:])
+
         return results
 
     def _score_target(self, batch, memory_bank, src_lengths, data, src_map):