length pen param (facebookresearch#2372)

jeremyzzzhj · Jan 31, 2020 · 80904aa · 80904aa
1 parent 1791379
commit 80904aa
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/parlai/core/torch_generator_agent.py b/parlai/core/torch_generator_agent.py
@@ -302,6 +302,12 @@ def add_cmdline_args(cls, argparser):
             default=-1,
             help='Size n-grams to block in beam search. val <= 0 implies no blocking',
         )
+        agent.add_argument(
+            '--beam-length-penalty',
+            type=float,
+            default=0.65,
+            help='Applies a length penalty. Set to 0 for no penalty.',
+        )
         agent.add_argument(
             '--skip-generation',
             type='bool',
@@ -804,6 +810,7 @@ def _treesearch_factory(self, device):
                 min_length=0,
                 block_ngram=self.beam_block_ngram,
                 context_block_ngram=self.beam_context_block_ngram,
+                length_penalty=self.opt.get('beam_length_penalty', 0.65),
                 padding_token=self.NULL_IDX,
                 bos_token=self.START_IDX,
                 eos_token=self.END_IDX,
@@ -815,6 +822,7 @@ def _treesearch_factory(self, device):
                 min_length=self.beam_min_length,
                 block_ngram=self.beam_block_ngram,
                 context_block_ngram=self.beam_context_block_ngram,
+                length_penalty=self.opt.get('beam_length_penalty', 0.65),
                 padding_token=self.NULL_IDX,
                 bos_token=self.START_IDX,
                 eos_token=self.END_IDX,
@@ -827,6 +835,7 @@ def _treesearch_factory(self, device):
                 min_length=self.beam_min_length,
                 block_ngram=self.beam_block_ngram,
                 context_block_ngram=self.beam_context_block_ngram,
+                length_penalty=self.opt.get('beam_length_penalty', 0.65),
                 padding_token=self.NULL_IDX,
                 bos_token=self.START_IDX,
                 eos_token=self.END_IDX,
@@ -839,6 +848,7 @@ def _treesearch_factory(self, device):
                 min_length=self.beam_min_length,
                 block_ngram=self.beam_block_ngram,
                 context_block_ngram=self.beam_context_block_ngram,
+                length_penalty=self.opt.get('beam_length_penalty', 0.65),
                 padding_token=self.NULL_IDX,
                 bos_token=self.START_IDX,
                 eos_token=self.END_IDX,
@@ -976,6 +986,7 @@ def __init__(
         eos_token=2,
         min_length=3,
         device='cpu',
+        length_penalty=0.65,
     ):
         """
         Instantiate Beam object.
@@ -998,6 +1009,7 @@ def __init__(
             What device to use for computations
         """
         self.beam_size = beam_size
+        self.length_penalty = length_penalty
         self.block_ngram = block_ngram
         self.min_length = min_length
         self.eos = eos_token
@@ -1237,7 +1249,7 @@ def get_rescored_finished(self, n_best=None):
         for finished_item in self.finished:
             current_length = finished_item.timestep + 1
             # these weights are from Google NMT paper
-            length_penalty = math.pow((1 + current_length) / 6, 0.65)
+            length_penalty = math.pow((1 + current_length) / 6, self.length_penalty)
             rescored_finished.append(
                 _HypothesisTail(
                     timestep=finished_item.timestep,