diff --git a/trl/scripts/dpo.py b/trl/scripts/dpo.py index 88c93398ed4..47d2f3527c4 100644 --- a/trl/scripts/dpo.py +++ b/trl/scripts/dpo.py @@ -159,6 +159,9 @@ def main(script_args, training_args, model_args, dataset_args): # Train the model trainer.train() + # Log training complete + trainer.accelerator.print("✅ Training completed.") + if training_args.eval_strategy != "no": metrics = trainer.evaluate() trainer.log_metrics("eval", metrics) @@ -166,8 +169,11 @@ def main(script_args, training_args, model_args, dataset_args): # Save and push to Hub trainer.save_model(training_args.output_dir) + trainer.accelerator.print(f"💾 Model saved to {training_args.output_dir}.") + if training_args.push_to_hub: trainer.push_to_hub(dataset_name=script_args.dataset_name) + trainer.accelerator.print(f"🤗 Model pushed to the Hub in https://huggingface.co/{trainer.hub_model_id}.") def make_parser(subparsers: Optional[argparse._SubParsersAction] = None): diff --git a/trl/scripts/grpo.py b/trl/scripts/grpo.py index 5373e482db0..de91941d0e8 100644 --- a/trl/scripts/grpo.py +++ b/trl/scripts/grpo.py @@ -141,10 +141,16 @@ def main(script_args, training_args, model_args, dataset_args): # Train the model trainer.train() + # Log training complete + trainer.accelerator.print("✅ Training completed.") + # Save and push to Hub trainer.save_model(training_args.output_dir) + trainer.accelerator.print(f"💾 Model saved to {training_args.output_dir}.") + if training_args.push_to_hub: trainer.push_to_hub(dataset_name=script_args.dataset_name) + trainer.accelerator.print(f"🤗 Model pushed to the Hub in https://huggingface.co/{trainer.hub_model_id}.") def make_parser(subparsers: Optional[argparse._SubParsersAction] = None): diff --git a/trl/scripts/kto.py b/trl/scripts/kto.py index 4c71f0178e1..a17f9f96544 100644 --- a/trl/scripts/kto.py +++ b/trl/scripts/kto.py @@ -135,10 +135,16 @@ def main(script_args, training_args, model_args, dataset_args): # Train the model trainer.train() + # Log training complete + trainer.accelerator.print("✅ Training completed.") + # Save and push to Hub trainer.save_model(training_args.output_dir) + trainer.accelerator.print(f"💾 Model saved to {training_args.output_dir}.") + if training_args.push_to_hub: trainer.push_to_hub(dataset_name=script_args.dataset_name) + trainer.accelerator.print(f"🤗 Model pushed to the Hub in https://huggingface.co/{trainer.hub_model_id}.") def make_parser(subparsers: Optional[argparse._SubParsersAction] = None): diff --git a/trl/scripts/rloo.py b/trl/scripts/rloo.py index 0a1941c556d..701a41746c5 100644 --- a/trl/scripts/rloo.py +++ b/trl/scripts/rloo.py @@ -141,10 +141,16 @@ def main(script_args, training_args, model_args, dataset_args): # Train the model trainer.train() + # Log training complete + trainer.accelerator.print("✅ Training completed.") + # Save and push to Hub trainer.save_model(training_args.output_dir) + trainer.accelerator.print(f"💾 Model saved to {training_args.output_dir}.") + if training_args.push_to_hub: trainer.push_to_hub(dataset_name=script_args.dataset_name) + trainer.accelerator.print(f"🤗 Model pushed to the Hub in https://huggingface.co/{trainer.hub_model_id}.") def make_parser(subparsers: Optional[argparse._SubParsersAction] = None): diff --git a/trl/scripts/sft.py b/trl/scripts/sft.py index a9fa28a9a94..742cf29d741 100644 --- a/trl/scripts/sft.py +++ b/trl/scripts/sft.py @@ -152,10 +152,16 @@ def main(script_args, training_args, model_args, dataset_args): # Train the model trainer.train() + # Log training complete + trainer.accelerator.print("✅ Training completed.") + # Save and push to Hub trainer.save_model(training_args.output_dir) + trainer.accelerator.print(f"💾 Model saved to {training_args.output_dir}.") + if training_args.push_to_hub: trainer.push_to_hub(dataset_name=script_args.dataset_name) + trainer.accelerator.print(f"🤗 Model pushed to the Hub in https://huggingface.co/{trainer.hub_model_id}.") def make_parser(subparsers: Optional[argparse._SubParsersAction] = None):