MBART supports freeze multi-lingual model when dy2sta

paddlenlp/ops/faster_transformer/sample/mbart_export_model_sample.py

@@ -91,8 +91,9 @@ def do_predict(args):
     place = paddle.set_device(place)
 
     model = MBartForConditionalGeneration.from_pretrained(
-        args.model_name_or_path, src_lang="en_XX")
-    tokenizer = MBartTokenizer.from_pretrained(args.model_name_or_path)
+        args.model_name_or_path)
+    tokenizer = MBartTokenizer.from_pretrained(args.model_name_or_path,
+                                               src_lang="en_XX")
 
     bos_id = tokenizer.lang_code_to_id["zh_CN"]
     eos_id = model.mbart.config["eos_token_id"]
@@ -115,7 +116,9 @@ def do_predict(args):
             None,
             # seq_len
             None,
-            bos_id,  # forced_bos_token_id
+            paddle.static.InputSpec(
+                shape=[None, 1], dtype="int32"
+            ),  # forced_bos_token_id can be a Tensor or int (bos_id)
             args.num_beams,  # num_beams.
             args.topk,  # top_k
             args.topp,  # top_p

paddlenlp/ops/faster_transformer/sample/mbart_inference.py

@@ -31,6 +31,18 @@ def setup_args():
                         default="./infer_model/",
                         type=str,
                         help="Path to save inference model of BART. ")
+    parser.add_argument("--batch_size",
+                        default=1,
+                        type=int,
+                        help="Batch size. ")
+    parser.add_argument("--num_return_sequences",
+                        default=1,
+                        type=int,
+                        help="The number of returned sequences. When " \
+                             "applying sampling to decode, this must be " \
+                             "given according to the --num_return_sequences " \
+                             "when you process mbart export scripts. When applying " \
+                             "beam search to decode, this flag should NOT be set. ")
 
     args = parser.parse_args()
 
@@ -53,14 +65,22 @@ def postprocess_response(tokenizer, seq, bos_idx, eos_idx):
 
 def infer(args):
     model_name = "mbart-large-50-many-to-many-mmt"
-    tokenizer = MBartTokenizer.from_pretrained(model_name)
 
+    tokenizer = MBartTokenizer.from_pretrained(model_name, src_lang="en_XX")
     bos_id = tokenizer.lang_code_to_id["zh_CN"]
+    inputs = "PaddleNLP is a powerful NLP library with Awesome pre-trained models and easy-to-use interface, supporting wide-range of NLP tasks from research to industrial applications."
+
     eos_id = tokenizer.eos_token_id
 
-    inputs = "PaddleNLP is a powerful NLP library with Awesome pre-trained models and easy-to-use interface, supporting wide-range of NLP tasks from research to industrial applications."
+    # Input ids
     input_ids = tokenizer(inputs)["input_ids"]
-    input_ids = np.asarray(input_ids, dtype="int32").reshape(1, -1)
+    input_ids = np.asarray(input_ids,
+                           dtype="int32").reshape(1, -1).repeat(args.batch_size,
+                                                                axis=0)
+
+    # Forced bos token ids
+    forced_bos_token = np.ones([args.batch_size * args.num_return_sequences, 1],
+                               dtype="int32") * bos_id
 
     # Load FasterTransformer lib.
     load("FasterTransformer", verbose=True)
@@ -74,20 +94,32 @@ def infer(args):
     predictor = paddle_infer.create_predictor(config)
 
     input_names = predictor.get_input_names()
-    input_handle = predictor.get_input_handle(input_names[0])
-    input_handle.copy_from_cpu(input_ids.astype("int32"))
+
+    # Input ids
+    input_ids_handle = predictor.get_input_handle(input_names[0])
+    input_ids_handle.copy_from_cpu(input_ids.astype("int32"))
+
+    # Forced bos token ids
+    forced_bos_token_handle = predictor.get_input_handle(input_names[1])
+    forced_bos_token_handle.copy_from_cpu(forced_bos_token.astype("int32"))
 
     predictor.run()
 
     output_names = predictor.get_output_names()
     output_handle = predictor.get_output_handle(output_names[0])
     output_data = output_handle.copy_to_cpu()
 
-    result = postprocess_response(
-        tokenizer,
-        output_data.transpose([1, 2, 0]).tolist()[0][0], bos_id, eos_id)
+    # [batch_size, num_beams * 2, sequence_length]
+    output_data = output_data.transpose([1, 2, 0])
+
+    # Only use the best sequence.
+    result = [
+        postprocess_response(tokenizer,
+                             sample.tolist()[0], bos_id, eos_id)
+        for sample in output_data
+    ]
     print("Model input:", inputs)
-    print("Result:", result)
+    print("Result:", "\n".join(result))
 
 
 if __name__ == "__main__":

paddlenlp/ops/faster_transformer/transformer/faster_transformer.py

@@ -1474,22 +1474,24 @@ def forward(self,
         if decoder_start_token_id is not None:
             bos_token_id = decoder_start_token_id
 
-        if forced_bos_token_id is not None:
-            if decode_strategy == "sampling":
-                trg_word = paddle.full([batch_size * num_return_sequences, 1],
-                                       forced_bos_token_id,
-                                       dtype="int32")
+        if not isinstance(forced_bos_token_id, paddle.Tensor):
+            if forced_bos_token_id is not None:
+                if decode_strategy == "sampling":
+                    forced_bos_token_id = paddle.full(
+                        [batch_size * num_return_sequences, 1],
+                        forced_bos_token_id,
+                        dtype="int32")
+                else:
+                    forced_bos_token_id = paddle.full([batch_size, 1],
+                                                      forced_bos_token_id,
+                                                      dtype="int32")
             else:
-                trg_word = paddle.full([batch_size, 1],
-                                       forced_bos_token_id,
-                                       dtype="int32")
-        else:
-            trg_word = paddle.zeros([0])
+                forced_bos_token_id = paddle.zeros([0])
 
         return self.decoding(enc_output=encoder_output,
                              memory_seq_lens=seq_len,
                              beam_size=num_beams,
-                             trg_word=trg_word,
+                             trg_word=forced_bos_token_id,
                              top_k=top_k,
                              top_p=top_p,
                              decoding_strategy=decode_strategy,