Add eval rouge metrics

pipelines/tools/finetune.py

@@ -25,6 +25,8 @@ from datasets import Dataset
 from transformers import TrainingArguments
 from trl import SFTTrainer
 
+from .summarization_eval import make_compute_metrics
+
 logger = logging.getLogger(__name__)
 
 
@@ -187,6 +189,9 @@ def main(
         optim="adamw_8bit",
         seed=42,
         report_to="none",
+        metric_for_best_model="eval_composite",
+        greater_is_better=True,
+        predict_with_generate=True,
     )
 
     trainer = SFTTrainer(
@@ -197,6 +202,7 @@ def main(
         args=training_args,
         max_seq_length=config.training.max_seq_length,
         packing=True,
+        compute_metrics=make_compute_metrics(tokenizer),
     )
 
     logger.info(