fix: setting log level in save() (#304)

* fix: setting log level in save() Signed-off-by: Anh Uong <[email protected]> * add log level doc string to save Signed-off-by: Anh Uong <[email protected]> --------- Signed-off-by: Anh Uong <[email protected]>
foundation-model-stack · Aug 16, 2024 · a6d093e · a6d093e
1 parent 2d1c17c
commit a6d093e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py
@@ -468,7 +468,7 @@ def test_run_causallm_ft_save_with_save_model_dir_save_strategy_no():
  # validate that no checkpoints created
  assert not any(x.startswith("checkpoint-") for x in os.listdir(tempdir))
 
- sft_trainer.save(tempdir, trainer)
+ sft_trainer.save(tempdir, trainer, "debug")
  assert any(x.endswith(".safetensors") for x in os.listdir(tempdir))
  _test_run_inference(checkpoint_path=tempdir)
 

diff --git a/tuning/sft_trainer.py b/tuning/sft_trainer.py
@@ -373,13 +373,15 @@ def save(path: str, trainer: SFTTrainer, log_level="WARNING"):
  Path to save the model to.
  trainer: SFTTrainer
  Instance of SFTTrainer used for training to save the model.
+ log_level: str
+ Optional threshold to set save save logger to, default warning.
  """
  logger = logging.getLogger("sft_trainer_save")
  # default value from TrainingArguments
  if log_level == "passive":
  log_level = "WARNING"
 
- logger.setLevel(log_level)
+ logger.setLevel(log_level.upper())
 
  if not os.path.exists(path):
  os.makedirs(path, exist_ok=True)