| 457 | ) |
| 458 | |
| 459 | def postprocess_text(preds, labels): |
| 460 | preds = [pred.strip() for pred in preds] |
| 461 | labels = [label.strip() for label in labels] |
| 462 | |
| 463 | # rougeLSum expects newline after each sentence |
| 464 | preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in preds] |
| 465 | labels = ["\n".join(nltk.sent_tokenize(label)) for label in labels] |
| 466 | |
| 467 | return preds, labels |
| 468 | |
| 469 | train_dataloader = DataLoader( |
| 470 | train_dataset, shuffle=True, collate_fn=data_collator, batch_size=args.per_device_train_batch_size, num_workers=0, |