Expose beam search diversity rate for ft (PaddlePaddle#1072)

* expose beam search diversity rate for ft

* rename param

* fix generation api

* fix bart

Author: FrostML

examples/machine_translation/transformer/faster_transformer/README.md

@@ -139,6 +139,7 @@ python encoder_decoding_predict.py --config ../configs/transformer.base.yaml --d
   * 当使用 `topk_sampling` 的时候，需要指定 `--topk` 的值
   * 当使用 `topp_sampling` 的时候，需要指定 `topp` 的值，并且需要保证 `--topk` 的值为 0
 * `--beam_size`: 解码策略是 `beam_search` 的时候，beam size 的大小，数据类型是 `int`
+* `--diversity_rate`: 解码策略是 `beam_search` 的时候，设置 diversity rate 的大小，数据类型是 `float`。当设置的 `diversity_rate` 大于 0 的时候，FasterTransformer 仅支持 beam size 为 1，4，16，64
 * `--topk`: 解码策略是 `topk_sampling` 的时候，topk 计算的 k 值的大小，数据类型是 `int`
 * `--topp`: 解码策略是 `topp_sampling` 的时候，p 的大小，数据类型是 `float`
 

examples/machine_translation/transformer/faster_transformer/encoder_decoding_predict.py

@@ -45,6 +45,11 @@ def parse_args():
         help="Decoding strategy. Can be one of ['beam_search', 'topk_sampling', 'topp_sampling']. "
     )
     parser.add_argument("--beam_size", default=5, type=int, help="Beam size. ")
+    parser.add_argument(
+        "--diversity_rate",
+        default=0.0,
+        type=float,
+        help="The diversity rate for beam search. ")
     parser.add_argument(
         "--topk",
         default=4,
@@ -144,6 +149,7 @@ def do_predict(args):
         decoding_strategy=args.decoding_strategy,
         beam_size=args.beam_size,
         max_out_len=args.max_out_len,
+        diversity_rate=args.diversity_rate,
         decoding_lib=args.decoding_lib,
         use_fp16_decoding=args.use_fp16_decoding)
 
@@ -206,6 +212,7 @@ def do_predict(args):
     args.use_fp16_decoding = ARGS.use_fp16_decoding
     args.decoding_strategy = ARGS.decoding_strategy
     args.beam_size = ARGS.beam_size
+    args.diversity_rate = ARGS.diversity_rate
     args.topk = ARGS.topk
     args.topp = ARGS.topp
     args.profile = ARGS.profile

paddlenlp/ops/faster_transformer/sample/bart_decoding_sample.py

@@ -81,7 +81,7 @@ def parse_args():
     parser.add_argument(
         "--max_out_len", default=50, type=int, help="Maximum output length. ")
     parser.add_argument(
-        "--beam_search_diversity_rate",
+        "--diversity_rate",
         default=0.0,
         type=float,
         help="The diversity of beam search. ")
@@ -144,7 +144,7 @@ def do_predict(args):
         topk=args.topk,
         topp=args.topp,
         max_out_len=args.max_out_len,
-        beam_search_diversity_rate=args.beam_search_diversity_rate,
+        diversity_rate=args.diversity_rate,
         decoding_lib=args.decoding_lib,
         use_fp16_decoding=args.use_fp16_decoding,
         rel_len=args.rel_len,

paddlenlp/ops/faster_transformer/transformer/decoding.py

@@ -36,8 +36,8 @@ def infer_transformer_decoding(
         ffn_inter_bias, ffn_out_weight, ffn_out_bias, decoder_ln_weight,
         decoder_ln_bias, linear_weight, linear_bias, pos_emb,
         _decoding_strategy, _beam_size, _topk, _topp, _n_head, _size_per_head,
-        _n_layer, _bos_id, _eos_id, _max_out_len, _beam_search_diversity_rate,
-        _rel_len, _alpha):
+        _n_layer, _bos_id, _eos_id, _max_out_len, _diversity_rate, _rel_len,
+        _alpha):
     helper = LayerHelper('fusion_decoding', **locals())
 
     inputs = {
@@ -88,7 +88,7 @@ def infer_transformer_decoding(
         'bos_id': _bos_id,
         'eos_id': _eos_id,
         'max_len': _max_out_len,
-        'beam_search_diversity_rate': _beam_search_diversity_rate,
+        'beam_search_diversity_rate': _diversity_rate,
         "rel_len": _rel_len,
         "alpha": _alpha
     }
@@ -175,8 +175,8 @@ def infer_unified_decoding(
         trans_weight, trans_bias, lm_ln_weight, lm_ln_bias, linear_weight,
         linear_bias, pos_emb, type_emb, _decoding_strategy, _beam_size, _topk,
         _topp, _n_head, _size_per_head, _n_layer, _bos_id, _eos_id,
-        _max_out_len, _beam_search_diversity_rate, _unk_id, _mask_id,
-        _temperature, _len_penalty, _normalize_before, _pos_bias, _hidden_act):
+        _max_out_len, _diversity_rate, _unk_id, _mask_id, _temperature,
+        _len_penalty, _normalize_before, _pos_bias, _hidden_act):
     helper = LayerHelper('fusion_unified_decoding', **locals())
 
     inputs = {
@@ -225,7 +225,7 @@ def infer_unified_decoding(
         "bos_id": _bos_id,
         "eos_id": _eos_id,
         "max_len": _max_out_len,
-        "beam_search_diversity_rate": _beam_search_diversity_rate,
+        "beam_search_diversity_rate": _diversity_rate,
         "unk_id": _unk_id,
         "mask_id": _mask_id,
         "temperature": _temperature,
@@ -264,8 +264,8 @@ def infer_bart_decoding(
         ffn_inter_bias, ffn_out_weight, ffn_out_bias, decoder_ln_weight,
         decoder_ln_bias, linear_weight, linear_bias, pos_emb,
         _decoding_strategy, _beam_size, _topk, _topp, _n_head, _size_per_head,
-        _n_layer, _bos_id, _eos_id, _max_out_len, _beam_search_diversity_rate,
-        _rel_len, _alpha):
+        _n_layer, _bos_id, _eos_id, _max_out_len, _diversity_rate, _rel_len,
+        _alpha):
 
     helper = LayerHelper('fusion_bart_decoding', **locals())
 
@@ -317,7 +317,7 @@ def infer_bart_decoding(
         'bos_id': _bos_id,
         'eos_id': _eos_id,
         'max_len': _max_out_len,
-        'beam_search_diversity_rate': _beam_search_diversity_rate,
+        'beam_search_diversity_rate': _diversity_rate,
         "rel_len": _rel_len,
         "alpha": _alpha
     }
@@ -391,7 +391,7 @@ def __init__(self,
                  topk=1,
                  topp=0.0,
                  max_out_len=256,
-                 beam_search_diversity_rate=0.0,
+                 diversity_rate=0.0,
                  decoding_lib=None,
                  use_fp16_decoding=False,
                  rel_len=False,
@@ -564,8 +564,8 @@ def forward(self, enc_output, memory_seq_lens):
             self._decoding_strategy, self._beam_size, self._topk, self._topp,
             self._n_head,
             int(self._d_model / self._n_head), self._num_decoder_layers,
-            self._bos_id, self._eos_id, self._max_out_len,
-            self._beam_search_diversity_rate, self._rel_len, self._alpha)
+            self._bos_id, self._eos_id, self._max_out_len, self._diversity_rate,
+            self._rel_len, self._alpha)
 
         ids = finalize(
             self._beam_size,
@@ -1048,7 +1048,7 @@ def forward(self,
                 eos_id=1,
                 temperature=1.0,
                 length_penalty=1.0,
-                beam_search_diversity_rate=0.0,
+                diversity_rate=0.0,
                 pos_bias=True):
         output_ids, parent_ids, sequence_length = infer_unified_decoding(
             cache_k=cache_k,
@@ -1093,7 +1093,7 @@ def forward(self,
             _bos_id=bos_id,
             _eos_id=eos_id,
             _max_out_len=max_out_len,
-            _beam_search_diversity_rate=beam_search_diversity_rate,
+            _diversity_rate=diversity_rate,
             _unk_id=self._unk_id,
             _mask_id=self._mask_id,
             _temperature=temperature,
@@ -1120,7 +1120,7 @@ def __init__(self,
                  topk=1,
                  topp=0.0,
                  max_out_len=256,
-                 beam_search_diversity_rate=0.0,
+                 diversity_rate=0.0,
                  decoding_lib=None,
                  use_fp16_decoding=False,
                  rel_len=False,
@@ -1321,8 +1321,8 @@ def forward(self, enc_output, memory_seq_lens):
             self._decoding_strategy, self._beam_size, self._topk, self._topp,
             self._n_head,
             int(self._d_model / self._n_head), self._num_decoder_layers,
-            self._bos_id, self._eos_id, self._max_out_len,
-            self._beam_search_diversity_rate, self._rel_len, self._alpha)
+            self._bos_id, self._eos_id, self._max_out_len, self._diversity_rate,
+            self._rel_len, self._alpha)
 
         ids = finalize(
             self._beam_size,

paddlenlp/ops/faster_transformer/transformer/faster_transformer.py

@@ -91,6 +91,8 @@ class FasterTransformer(TransformerModel):
             `topp` are kept for top-p sampling. Defaults to 4. 
         max_out_len (int, optional):
             The maximum output length. Defaults to 256.
+        diversity_rate (float, optional):
+            The diversity rate for beam search. Defaults to 0.0.
         use_fp16_decoding(bool, optional): Whether to use fp16 for decoding. 
         rel_len(bool, optional):
             Indicating whether `max_out_len` in is the length relative to that
@@ -123,6 +125,7 @@ def __init__(self,
                  topk=1,
                  topp=0.0,
                  max_out_len=256,
+                 diversity_rate=0.0,
                  decoding_lib=None,
                  use_fp16_decoding=False,
                  rel_len=False,
@@ -141,6 +144,7 @@ def __init__(self,
         self.topk = args.pop("topk")
         self.topp = args.pop("topp")
         self.max_out_len = args.pop("max_out_len")
+        self.diversity_rate = args.pop("diversity_rate")
         self.decoding_lib = args.pop("decoding_lib")
         self.use_fp16_decoding = args.pop("use_fp16_decoding")
         self.rel_len = args.pop("rel_len")
@@ -177,6 +181,7 @@ def __init__(self,
             topk=topk,
             topp=topp,
             max_out_len=max_out_len,
+            diversity_rate=self.diversity_rate,
             decoding_lib=self.decoding_lib,
             use_fp16_decoding=self.use_fp16_decoding,
             rel_len=self.rel_len,
@@ -480,6 +485,9 @@ def __init__(self,
         self.d_model = d_model
         self.max_length = max_length
         self.output_time_major = kwargs.pop("output_time_major", True)
+        # Only works for Faster Transformer.
+        # TODO: original version supports diversity rate.
+        diversity_rate = kwargs.pop("diversity_rate", 0.0)
         use_fp16_decoding = kwargs.pop("use_fp16_decoding", False)
         use_ft = kwargs.pop("use_ft", True)
         beam_search_version = kwargs.pop("beam_search_version", "v1")
@@ -507,6 +515,7 @@ def __init__(self,
                     eos_id=eos_id,
                     beam_size=beam_size,
                     max_out_len=max_out_len,
+                    diversity_rate=diversity_rate,
                     decoding_strategy=decoding_strategy,
                     use_fp16_decoding=use_fp16_decoding,
                     rel_len=rel_len,
@@ -786,7 +795,7 @@ def sample(self,
             temperature=temperature)
 
     def beam_search(self, input_ids, beam_scorer, logits_processors, max_length,
-                    pad_token_id, eos_token_id, **model_kwargs):
+                    diversity_rate, pad_token_id, eos_token_id, **model_kwargs):
         max_length -= input_ids.shape[-1]
         model_inputs = self.prepare_inputs_for_generation(input_ids,
                                                           **model_kwargs)
@@ -796,6 +805,7 @@ def beam_search(self, input_ids, beam_scorer, logits_processors, max_length,
             model_inputs=model_inputs,
             max_length=max_length,
             num_beams=beam_scorer.num_beams,
+            diversity_rate=diversity_rate,
             temperature=temperature)
 
     def forward(self,
@@ -804,6 +814,7 @@ def forward(self,
                 top_k=4,
                 top_p=0.0,
                 num_beams=4,
+                diversity_rate=0.0,
                 temperature=1.0,
                 model_inputs=None,
                 **model_kwargs):
@@ -823,6 +834,7 @@ def forward(self,
             cache_v=cache_v,
             memory_seq_lens=seq_len,
             beam_size=num_beams,
+            diversity_rate=diversity_rate,
             topk=top_k,
             topp=top_p,
             max_out_len=max_length,
@@ -946,7 +958,7 @@ def sample(self,
             temperature=temperature)
 
     def beam_search(self, input_ids, beam_scorer, logits_processors, max_length,
-                    pad_token_id, eos_token_id, **model_kwargs):
+                    diversity_rate, pad_token_id, eos_token_id, **model_kwargs):
         max_length -= input_ids.shape[-1]
         model_inputs = self.prepare_inputs_for_generation(input_ids,
                                                           **model_kwargs)
@@ -956,6 +968,7 @@ def beam_search(self, input_ids, beam_scorer, logits_processors, max_length,
             model_inputs=model_inputs,
             max_length=max_length,
             num_beams=beam_scorer.num_beams,
+            diversity_rate=diversity_rate,
             temperature=temperature)
 
     def forward(self,
@@ -964,6 +977,7 @@ def forward(self,
                 top_k=4,
                 top_p=0.0,
                 num_beams=4,
+                diversity_rate=0.0,
                 temperature=1.0,
                 model_inputs=None,
                 **model_kwargs):
@@ -983,6 +997,7 @@ def forward(self,
             cache_v=cache_v,
             memory_seq_lens=seq_len,
             beam_size=num_beams,
+            diversity_rate=diversity_rate,
             topk=top_k,
             topp=top_p,
             max_out_len=max_length,
@@ -1001,7 +1016,7 @@ def __init__(self,
                  topk=1,
                  topp=0.0,
                  max_out_len=256,
-                 beam_search_diversity_rate=0.0,
+                 diversity_rate=0.0,
                  decoding_lib=None,
                  use_fp16_decoding=False,
                  rel_len=False,
@@ -1023,7 +1038,7 @@ def __init__(self,
             topk=topk,
             topp=topp,
             max_out_len=max_out_len,
-            beam_search_diversity_rate=beam_search_diversity_rate,
+            diversity_rate=diversity_rate,
             decoding_lib=decoding_lib,
             use_fp16_decoding=use_fp16_decoding)
 
@@ -1032,7 +1047,8 @@ def forward(self, input_ids):
         mem_seq_lens = paddle.sum(paddle.cast(
             input_ids != self.pad_id, dtype="int32"),
                                   axis=-1,
-                                  keepdim=True)
+                                  keepdim=True,
+                                  dtype="int32")
         if self.use_fp16_decoding:
             encoder_output = paddle.cast(encoder_output, "float16")
         return self.decoding(encoder_output, mem_seq_lens)