Generate api support encoder-decoder (PaddlePaddle#966)

* Add offset mapping doc

* fix eval hang because of unique endpoint

* generate api support encoder-decoder

Author: smallv0221

paddlenlp/transformers/bart/modeling.py

@@ -310,6 +310,12 @@ def __init__(self,
             max_position_embeddings, init_std)
         self.apply(self.init_weights)
 
+    def get_encoder(self):
+        return self.encoder
+
+    def get_decoder(self):
+        return self.decoder
+
     def forward(self,
                 input_ids,
                 attention_mask=None,
@@ -323,7 +329,6 @@ def forward(self,
         if input_ids is None and encoder_output is None:
             raise ValueError(
                 "You have to specify either input_ids or encoder_output")
-
         if decoder_input_ids is None:
             assert input_ids is not None, "input_ids should be " \
                                           "specified when generating decoder_input_ids"
@@ -450,6 +455,12 @@ def __init__(self, bart):
                              paddle.zeros((1, self.bart.config['vocab_size'])))
         self.apply(self.init_weights)
 
+    def get_encoder(self):
+        return self.bart.get_encoder()
+
+    def get_decoder(self):
+        return self.bart.get_decoder()
+
     def forward(self,
                 input_ids,
                 attention_mask=None,
@@ -465,5 +476,46 @@ def forward(self,
             output[0] if use_cache else output,
             self.lm_head_weight,
             transpose_y=True) + self.final_logits_bias
+        if use_cache:
+            cache = output[1]
+            return lm_logits, cache
+        else:
+            return lm_logits
+
+    def prepare_inputs_for_generation(self,
+                                      decoder_input_ids,
+                                      attention_mask=None,
+                                      decoder_attention_mask=None,
+                                      cache=None,
+                                      use_cache=False,
+                                      encoder_output=None,
+                                      **kwargs):
+        # cut decoder_input_ids if past is used
+        if cache is not None:
+            decoder_input_ids = decoder_input_ids[:, -1].unsqueeze(-1)
+            if decoder_attention_mask is not None:
+                decoder_attention_mask = decoder_attention_mask[:, :,
+                                                                -1, :].unsqueeze(
+                                                                    2)
+
+        return {
+            "input_ids": None,
+            "decoder_input_ids": decoder_input_ids,
+            "encoder_output": encoder_output,
+            "decoder_attention_mask": decoder_attention_mask,
+            "attention_mask": attention_mask,
+            "use_cache": use_cache,
+            "cache": cache
+        }
 
-        return lm_logits
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError as e:
+            try:
+                return getattr(getattr(self, self.base_model_prefix), name)
+            except AttributeError:
+                try:
+                    return getattr(self, self.base_model_prefix).config[name]
+                except KeyError:
+                    raise e

paddlenlp/transformers/generation_utils.py

@@ -281,7 +281,7 @@ def prepare_input_ids_for_generation(bos_token_id):
         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]) * bos_token_id
+        return paddle.ones([1, 1], dtype="int64") * bos_token_id
 
     @staticmethod
     def prepare_attention_mask_for_generation(input_ids, pad_token_id,
@@ -341,7 +341,9 @@ def expand_inputs_for_generation(input_ids,
         return input_ids, model_kwargs
 
     @staticmethod
-    def update_model_kwargs_for_generation(outputs, model_kwargs):
+    def update_model_kwargs_for_generation(outputs,
+                                           model_kwargs,
+                                           is_encoder_decoder=False):
         # Update the model inputs during generation. 
         # Note that If `token_type_ids` and `attention_mask` in `model_kwargs` 
         # and they contain pad value, the result vectors updated by this method 
@@ -366,7 +368,7 @@ def update_model_kwargs_for_generation(outputs, model_kwargs):
                 axis=-1)
 
         # update attention_mask
-        if "attention_mask" in model_kwargs:
+        if not is_encoder_decoder and "attention_mask" in model_kwargs:
             attention_mask = model_kwargs["attention_mask"]
             # nn.Pad2D don't support the data type `bool`
             if convert_dtype(attention_mask.dtype) == 'bool':
@@ -395,6 +397,22 @@ def update_scores_for_generation(scores, next_scores, length,
         scores = paddle.where(unfinished_flag, unfinished_scores, scores)
         return scores
 
+    def prepare_encoder_decoder_kwargs_for_generation(self, input_ids,
+                                                      model_kwargs):
+        if "encoder_output" not in model_kwargs:
+            # retrieve encoder hidden states
+            encoder = self.get_encoder()
+            encoder_kwargs = {
+                argument: value
+                for argument, value in model_kwargs.items()
+                if not (argument.startswith("decoder_") or argument.startswith(
+                    "cross_attn"))
+            }
+
+            model_kwargs["encoder_output"] = encoder(input_ids,
+                                                     **encoder_kwargs)
+        return model_kwargs
+
     def prepare_inputs_for_generation(self, input_ids, **kwargs):
         # Implement in subclasses for custom behavior to prepare inputs in the
         # generate method.
@@ -590,14 +608,22 @@ def generate(self,
             model_kwargs[
                 "attention_mask"] = self.prepare_attention_mask_for_generation(
                     input_ids, pad_token_id, eos_token_id)
+        self.is_encoder_decoder = hasattr(self, 'encoder') and hasattr(
+            self, 'decoder')
+        if self.is_encoder_decoder:
+            model_kwargs = self.prepare_encoder_decoder_kwargs_for_generation(
+                input_ids, model_kwargs)
+            # set input_ids as decoder_input_ids
+            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)
 
         if pad_token_id is None and eos_token_id is not None:
             print("Setting `pad_token_id` to `eos_token_id`:{} for "
                   "open-end generation.".format(eos_token_id))
             pad_token_id = eos_token_id
 
-        # TODO Add relevant processing for encoder_decoder model.
-
         model_kwargs["use_cache"] = use_cache
         max_length += input_ids.shape[-1]
         min_length += input_ids.shape[-1]
@@ -671,7 +697,6 @@ def greedy_search(self, input_ids, logits_processors, max_length,
             logits = outputs[0] if isinstance(outputs, tuple) else outputs
             # [batch_size, vocab_size]
             logits = logits[:, -1, :]
-
             # pre-process distribution
             logits = self.adjust_logits_during_generation(logits)
             logits = logits_processors(input_ids, logits)
@@ -700,8 +725,10 @@ def greedy_search(self, input_ids, logits_processors, max_length,
             if not paddle.any(unfinished_flag):
                 break
 
-            model_kwargs = self.update_model_kwargs_for_generation(outputs,
-                                                                   model_kwargs)
+            model_kwargs = self.update_model_kwargs_for_generation(
+                outputs,
+                model_kwargs,
+                is_encoder_decoder=self.is_encoder_decoder)
         return input_ids[:, origin_len:], scores
 
     def sample(self,
@@ -801,8 +828,10 @@ def TopPProcess(probs, top_p, min_tokens_to_keep):
             # Stop when there is a </s> in all sentences
             if not paddle.any(unfinished_flag):
                 break
-            model_kwargs = self.update_model_kwargs_for_generation(outputs,
-                                                                   model_kwargs)
+            model_kwargs = self.update_model_kwargs_for_generation(
+                outputs,
+                model_kwargs,
+                is_encoder_decoder=self.is_encoder_decoder)
         return input_ids[:, origin_len:], scores
 
     def beam_search(self, input_ids, beam_scorer, logits_processors, max_length,
@@ -876,8 +905,10 @@ def beam_search(self, input_ids, beam_scorer, logits_processors, max_length,
 
             if beam_scorer.is_done:
                 break
-            model_kwargs = self.update_model_kwargs_for_generation(outputs,
-                                                                   model_kwargs)
+            model_kwargs = self.update_model_kwargs_for_generation(
+                outputs,
+                model_kwargs,
+                is_encoder_decoder=self.is_encoder_decoder)
             if model_kwargs["cache"] is not None:
                 # reorder the cache
                 model_kwargs["cache"] = map_structure(