Fix generate api bug and add diverse sibling search (PaddlePaddle#1041)

* Add offset mapping doc

* fix eval hang because of unique endpoint

* generate api support encoder-decoder

* Add lightseq beam_search

* optimize performence

* add blockroughk kernel

* optimize

* minor fix

* Fix generate api bug and add diverse sibling search

* minor fix

Author: smallv0221

paddlenlp/transformers/generation_utils.py

@@ -277,11 +277,14 @@ class GenerationMixin(object):
     """
 
     @staticmethod
-    def prepare_input_ids_for_generation(bos_token_id):
+    def prepare_input_ids_for_generation(bos_token_id, encoder_output=None):
+        batch_size = 1
         if bos_token_id is None:
             raise ValueError("`bos_token_id` should be defined when no "
                              "`input_ids` are provided.")
-        return paddle.ones([1, 1], dtype="int64") * bos_token_id
+        if encoder_output is not None:
+            batch_size = encoder_output.shape[0]
+        return paddle.ones([batch_size, 1], dtype="int64") * bos_token_id
 
     @staticmethod
     def prepare_attention_mask_for_generation(input_ids, pad_token_id,
@@ -338,6 +341,11 @@ def expand_inputs_for_generation(input_ids,
             seq_len = model_kwargs["seq_len"]
             model_kwargs["seq_len"] = paddle.index_select(seq_len, index)
 
+        if "encoder_output" in model_kwargs:
+            encoder_output = model_kwargs["encoder_output"]
+            model_kwargs["encoder_output"] = paddle.index_select(encoder_output,
+                                                                 index)
+
         return input_ids, model_kwargs
 
     @staticmethod
@@ -441,6 +449,7 @@ def generate(self,
                  eos_token_id=None,
                  pad_token_id=None,
                  num_return_sequences=1,
+                 diversity_rate=0.0,
                  use_cache=True,
                  **model_kwargs):
         r"""
@@ -489,6 +498,9 @@ def generate(self,
                 None.
             num_return_sequences (int, optional): The number of returned 
                 sequences for each sequence in the batch. Default to 1.
+            diversity_rate (float, optional): The diversity_rate for diverse 
+                siblings search. See this paper for more details. 
+                `https://arxiv.org/abs/1611.08562`.
             use_cache: (bool, optional): Whether or not use the model cache to 
                 speed up decoding. Default to True.
             model_kwargs (dict): It can be used to specify additional kwargs 
@@ -617,7 +629,8 @@ def generate(self,
             if "decoder_input_ids" in model_kwargs:
                 input_ids = model_kwargs.pop("decoder_input_ids")
             else:
-                input_ids = self.prepare_input_ids_for_generation(bos_token_id)
+                input_ids = self.prepare_input_ids_for_generation(
+                    bos_token_id, model_kwargs["encoder_output"])
 
         if pad_token_id is None and eos_token_id is not None:
             print("Setting `pad_token_id` to `eos_token_id`:{} for "
@@ -673,8 +686,8 @@ def generate(self,
                 input_ids, expand_size=num_beams, **model_kwargs)
 
             return self.beam_search(input_ids, beam_scorer, logits_processors,
-                                    max_length, pad_token_id, eos_token_id,
-                                    **model_kwargs)
+                                    max_length, diversity_rate, pad_token_id,
+                                    eos_token_id, **model_kwargs)
 
         else:
             raise ValueError(
@@ -835,7 +848,7 @@ def TopPProcess(probs, top_p, min_tokens_to_keep):
         return input_ids[:, origin_len:], scores
 
     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):
         batch_size = len(beam_scorer._beam_hyps)
         num_beams = beam_scorer.num_beams
 
@@ -871,15 +884,50 @@ def beam_search(self, input_ids, beam_scorer, logits_processors, max_length,
             next_scores = paddle.log(next_scores)
 
             next_scores = next_scores + beam_scores.unsqueeze(-1)
-            # reshape for beam search
+
             vocab_size = next_scores.shape[-1]
-            next_scores = next_scores.reshape(
-                [batch_size, num_beams * vocab_size])
+            if diversity_rate == 0.0:
+                # reshape for beam search
+                next_scores = next_scores.reshape(
+                    [batch_size, num_beams * vocab_size])
 
-            next_scores, next_tokens = paddle.topk(
-                next_scores, 2 * num_beams, axis=1)
+                next_scores, next_tokens = paddle.topk(
+                    next_scores, 2 * num_beams, axis=1)
+
+                next_indices = next_tokens // vocab_size
+
+            else:
+                next_scores, next_tokens = paddle.topk(
+                    next_scores, 2 * num_beams, axis=1)
+
+                sibling_score = paddle.tile(
+                    paddle.arange(1, 2 * num_beams + 1),
+                    repeat_times=[batch_size * num_beams, 1]) * diversity_rate
+
+                diversed_score = next_scores - sibling_score
+                next_scores = next_scores.reshape(
+                    [batch_size, 2 * num_beams * num_beams])
+                next_tokens = next_tokens.reshape(
+                    [batch_size, 2 * num_beams * num_beams])
+
+                diversed_score = diversed_score.reshape(
+                    [batch_size, 2 * num_beams * num_beams])
+                diversed_score, diversed_tokens = paddle.topk(
+                    diversed_score, 2 * num_beams, axis=1)
+
+                # TODO
+                # Use gather_nd() to select origan token and score
+                next_scores = paddle.stack([
+                    paddle.index_select(next_scores[i], diversed_tokens[i])
+                    for i in range(next_scores.shape[0])
+                ])
+                next_tokens = paddle.stack([
+                    paddle.index_select(next_tokens[i], diversed_tokens[i])
+                    for i in range(next_tokens.shape[0])
+                ])
+
+                next_indices = next_tokens // (2 * num_beams)
 
-            next_indices = next_tokens // vocab_size
             next_tokens = next_tokens % vocab_size
 
             # stateless