Fix block_ngram_repeat w/ deepcopy (#1829)

Author: pltrdy

onmt/tests/test_beam_search.py

@@ -61,21 +61,18 @@ def test_advance_with_all_repeats_gets_blocked(self):
                     # (but it's still the best score, thus we have
                     # [BLOCKED_SCORE, -inf, -inf, -inf, -inf]
                     expected_scores = torch.tensor(
-                        [0] + [-float('inf')] * (beam_sz - 1))\
-                        .repeat(batch_sz, 1)
-                    expected_scores[:, 0] = self.BLOCKED_SCORE
+                        [self.BLOCKED_SCORE] + [-float('inf')] * (beam_sz - 1)
+                    ).repeat(batch_sz, 1)
                     self.assertTrue(beam.topk_log_probs.equal(expected_scores))
                 else:
                     # repetitions keeps maximizing score
                     # index 0 has been blocked, so repeating=>+0.0 score
                     # other indexes are -inf so repeating=>BLOCKED_SCORE
                     # which is higher
                     expected_scores = torch.tensor(
-                        [0] + [-float('inf')] * (beam_sz - 1))\
-                        .repeat(batch_sz, 1)
-                    expected_scores[:, :] = self.BLOCKED_SCORE
-                    expected_scores = torch.tensor(
-                        self.BLOCKED_SCORE).repeat(batch_sz, beam_sz)
+                        [self.BLOCKED_SCORE] + [-float('inf')] * (beam_sz - 1)
+                    ).repeat(batch_sz, 1)
+                    self.assertTrue(beam.topk_log_probs.equal(expected_scores))
 
     def test_advance_with_some_repeats_gets_blocked(self):
         # beam 0 and beam >=2 will repeat (beam >= 2 repeat dummy scores)
@@ -137,7 +134,8 @@ def test_advance_with_some_repeats_gets_blocked(self):
 
                     expected = torch.full([batch_sz, beam_sz], float("-inf"))
                     expected[:, 0] = no_repeat_score
-                    expected[:, 1:] = self.BLOCKED_SCORE
+                    expected[:, 1:3] = self.BLOCKED_SCORE
+                    expected[:, 3:] = float("-inf")
                     self.assertTrue(
                         beam.topk_log_probs.equal(expected))
 

onmt/translate/decode_strategy.py

@@ -1,4 +1,5 @@
 import torch
+from copy import deepcopy
 
 
 class DecodeStrategy(object):
@@ -184,7 +185,7 @@ def maybe_update_forbidden_tokens(self):
             # Reordering forbidden_tokens following beam selection
             # We rebuild a dict to ensure we get the value and not the pointer
             forbidden_tokens.append(
-                dict(self.forbidden_tokens[path_idx]))
+                deepcopy(self.forbidden_tokens[path_idx]))
 
             # Grabing the newly selected tokens and associated ngram
             current_ngram = tuple(seq[-n:].tolist())