Update llama datacollator (PaddlePaddle#6301)

Author: w5688414

examples/language_model/llama/finetune_generation.py

@@ -17,12 +17,7 @@
 from functools import partial
 
 import paddle
-from data import (
-    DataCollatorForSupervisedDataset,
-    convert_example,
-    custom_instruction_convert_example,
-    reader,
-)
+from data import convert_example, custom_instruction_convert_example, reader
 from modeling_pp import LlamaForCausalLMPipe
 from utils import (
     LlamaTrainer,
@@ -31,6 +26,7 @@
     save_infer_result,
 )
 
+from paddlenlp.data import DataCollatorForSeq2Seq
 from paddlenlp.datasets import load_dataset
 from paddlenlp.peft import LoRAConfig, LoRAModel, PrefixConfig, PrefixModelForCausalLM
 from paddlenlp.peft.prefix import llama_postprocess_past_key_value
@@ -221,8 +217,10 @@ def main():
         dev_ds = dev_ds.map(partial(trans_func, is_test=is_test))
 
     model_max_length = 1024 if not training_args.benchmark else 512
-    collate_fn = DataCollatorForSupervisedDataset(
-        tokenizer, max_length=model_max_length if data_args.always_pad_to_max_length else -1
+    collate_fn = DataCollatorForSeq2Seq(
+        return_tensors="pd",
+        tokenizer=tokenizer,
+        max_length=model_max_length if data_args.always_pad_to_max_length else -1,
     )
 
     def compute_metrics_trainer(eval_preds, tokenizer):