From de0eb21c70e234fb55379466949bc56938a3072f Mon Sep 17 00:00:00 2001 From: "Wu, Gangsheng" Date: Mon, 8 Jul 2024 16:30:59 +0000 Subject: [PATCH] update --- llm_on_ray/finetune/finetune.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/llm_on_ray/finetune/finetune.py b/llm_on_ray/finetune/finetune.py index 5af3c86c..84b7700b 100644 --- a/llm_on_ray/finetune/finetune.py +++ b/llm_on_ray/finetune/finetune.py @@ -374,10 +374,23 @@ def train_func(config: Dict[str, Any]): trainer = get_trainer(config, training_args, model, tokenizer, tokenized_dataset) - common.logger.info("train start") - trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint) - trainer.save_model() - common.logger.info("train finish") + if training_args.do_train: + common.logger.info("train start") + results = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint) + trainer.save_model() + + metrics = results.metrics + trainer.log_metrics("train", metrics) + trainer.save_metrics("train", metrics) + common.logger.info("train finish") + + if training_args.do_eval: + common.logger.info("eval start") + metrics = trainer.evaluate() + + trainer.log_metrics("eval", metrics) + trainer.save_metrics("eval", metrics) + common.logger.info("eval finish") def get_finetune_config():