diff --git a/src/llamafactory/hparams/finetuning_args.py b/src/llamafactory/hparams/finetuning_args.py index faf786d5..59194329 100644 --- a/src/llamafactory/hparams/finetuning_args.py +++ b/src/llamafactory/hparams/finetuning_args.py @@ -308,11 +308,31 @@ class BAdamArgument: class SwanLabArguments: use_swanlab: bool = field( default=False, - metadata={"help": ""}, + metadata={"help": "Whether or not to use the SwanLab (an experiment tracking and visualization tools)."}, ) - swanlab_name: str = field( + swanlab_project: str = field( default="", - metadata={}, + metadata={"help": "The project name in SwanLab."}, + ) + swanlab_workspace: str = field( + default="", + metadata={"help": "The workspace name in SwanLab."}, + ) + swanlab_experiment_name: str = field( + default="", + metadata={"help": "The experiment name in SwanLab."}, + ) + swanlab_description: str = field( + default="", + metadata={"help": "The experiment description in SwanLab."}, + ) + swanlab_mode: Literal["cloud", "local", "disabled"] = field( + default="cloud", + metadata={"help": "The mode of SwanLab."}, + ) + swanlab_api_key: str = field( + default="", + metadata={"help": "The API key for SwanLab."}, ) diff --git a/src/llamafactory/train/trainer_utils.py b/src/llamafactory/train/trainer_utils.py index d33bc538..e13fe552 100644 --- a/src/llamafactory/train/trainer_utils.py +++ b/src/llamafactory/train/trainer_utils.py @@ -463,6 +463,18 @@ def get_swanlab_callback(finetuning_args: "FinetuningArguments") -> "TrainerCall r""" Gets the callback for logging to SwanLab. """ - from swanlab.integration.huggingface import SwanLabCallback + import swanlab + from swanlab.integration.transformers import SwanLabCallback + + if finetuning_args.swanlab_api_key is not None: + swanlab.login(api_key=finetuning_args.swanlab_api_key) - return SwanLabCallback() + swanlab_callback = SwanLabCallback( + project=finetuning_args.swanlab_project, + workspace=finetuning_args.swanlab_workspace, + experiment_name=finetuning_args.swanlab_experiment_name, + description=finetuning_args.swanlab_description, + mode=finetuning_args.swanlab_mode, + ) + + return swanlab_callback \ No newline at end of file