Support coverage penalty for beam search decoder.

PiperOrigin-RevId: 212902045

Author: Anonymous

README.md

@@ -988,7 +988,8 @@ decoder = tf.contrib.seq2seq.BeamSearchDecoder(
         initial_state=decoder_initial_state,
         beam_width=beam_width,
         output_layer=projection_layer,
-        length_penalty_weight=0.0)
+        length_penalty_weight=0.0,
+        coverage_penalty_weight=0.0)
 
 # Dynamic decoding
 outputs, _ = tf.contrib.seq2seq.dynamic_decode(decoder, ...)

nmt/model.py

@@ -534,12 +534,15 @@ def _build_decoder(self, encoder_outputs, encoder_state, hparams):
         start_tokens = tf.fill([self.batch_size], tgt_sos_id)
         end_token = tgt_eos_id
         utils.print_out(
-            "  decoder: infer_mode=%sbeam_width=%d, length_penalty=%f" % (
-                infer_mode, hparams.beam_width, hparams.length_penalty_weight))
+            "  decoder: infer_mode=%sbeam_width=%d, "
+            "length_penalty=%f, coverage_penalty=%f"
+            % (infer_mode, hparams.beam_width, hparams.length_penalty_weight,
+               hparams.coverage_penalty_weight))
 
         if infer_mode == "beam_search":
           beam_width = hparams.beam_width
           length_penalty_weight = hparams.length_penalty_weight
+          coverage_penalty_weight = hparams.coverage_penalty_weight
 
           my_decoder = tf.contrib.seq2seq.BeamSearchDecoder(
               cell=cell,
@@ -549,7 +552,8 @@ def _build_decoder(self, encoder_outputs, encoder_state, hparams):
               initial_state=decoder_initial_state,
               beam_width=beam_width,
               output_layer=self.output_layer,
-              length_penalty_weight=length_penalty_weight)
+              length_penalty_weight=length_penalty_weight,
+              coverage_penalty_weight=coverage_penalty_weight)
         elif infer_mode == "sample":
           # Helper
           sampling_temperature = hparams.sampling_temperature

nmt/nmt.py

@@ -37,8 +37,9 @@
 
 INFERENCE_KEYS = ["src_max_len_infer", "tgt_max_len_infer", "subword_option",
                   "infer_batch_size", "beam_width",
-                  "length_penalty_weight", "sampling_temperature",
-                  "num_translations_per_input", "infer_mode"]
+                  "length_penalty_weight", "coverage_penalty_weight",
+                  "sampling_temperature", "num_translations_per_input",
+                  "infer_mode"]
 
 
 def add_arguments(parser):
@@ -288,6 +289,8 @@ def add_arguments(parser):
       """))
   parser.add_argument("--length_penalty_weight", type=float, default=0.0,
                       help="Length penalty for beam search.")
+  parser.add_argument("--coverage_penalty_weight", type=float, default=0.0,
+                      help="Coverage penalty for beam search.")
   parser.add_argument("--sampling_temperature", type=float,
                       default=0.0,
                       help=("""\
@@ -370,6 +373,7 @@ def create_hparams(flags):
       infer_mode=flags.infer_mode,
       beam_width=flags.beam_width,
       length_penalty_weight=flags.length_penalty_weight,
+      coverage_penalty_weight=flags.coverage_penalty_weight,
       sampling_temperature=flags.sampling_temperature,
       num_translations_per_input=flags.num_translations_per_input,
 

nmt/utils/standard_hparams_utils.py

@@ -95,6 +95,7 @@ def create_standard_hparams():
       # only enable beam search during inference when beam_width > 0.
       beam_width=0,
       length_penalty_weight=0.0,
+      coverage_penalty_weight=0.0,
       override_loaded_hparams=True,
       num_keep_ckpts=5,
       avg_ckpts=False,