From 3c00f406d64448ca060059e3202431db01f1ab06 Mon Sep 17 00:00:00 2001 From: David Meikle Date: Wed, 21 Feb 2024 15:57:18 +0000 Subject: [PATCH] Allow load_best_model_at_end to be configured for early stopping on custom evaluation datasets (#1291) * Allow load_best_model_at_end when using test_datasets and val_set_size is zero for custom evaluation datasets * Fixed formatting following failed Lint check --- src/axolotl/core/trainer_builder.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py index 7d39aefea..c69f6cf5a 100644 --- a/src/axolotl/core/trainer_builder.py +++ b/src/axolotl/core/trainer_builder.py @@ -883,8 +883,10 @@ class HFCausalTrainerBuilder(TrainerBuilderBase): self.cfg.load_best_model_at_end is not False or self.cfg.early_stopping_patience ) - and not self.cfg.test_datasets - and self.cfg.val_set_size > 0 + and ( + (not self.cfg.test_datasets and self.cfg.val_set_size > 0) + or (self.cfg.test_datasets and self.cfg.val_set_size == 0) + ) and self.cfg.save_steps and self.cfg.eval_steps and self.cfg.save_steps % self.cfg.eval_steps == 0