ros/LLaMA-Factory
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[3rdparty] support swanlab lark notification (#7481)

Browse Source
This commit is contained in:
Xu-pixel
2025-03-27 01:52:01 +08:00
committed by GitHub
parent 24afceddb7
commit b578a7d5b6
2 changed files with 17 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/llamafactory/train/trainer_utils.py
View File

@@ -599,6 +599,15 @@ def get_swanlab_callback(finetuning_args: "FinetuningArguments") -> "TrainerCall
if finetuning_args.swanlab_api_key is not None:
swanlab.login(api_key=finetuning_args.swanlab_api_key)
if finetuning_args.swanlab_lark_webhook_url is not None:
from swanlab.plugin.notification import LarkCallback # type: ignore
lark_callback = LarkCallback(
webhook_url=finetuning_args.swanlab_lark_webhook_url,
secret=finetuning_args.swanlab_lark_secret,
)
swanlab.register_callbacks([lark_callback])
class SwanLabCallbackExtension(SwanLabCallback):
def setup(self, args: "TrainingArguments", state: "TrainerState", model: "PreTrainedModel", **kwargs):
if not state.is_world_process_zero: