add swanlab

Former-commit-id: c85a77c8a8824a56a67d56b97b4877fcd6edeb3d

src/llamafactory/train/trainer_utils.py

@@ -40,7 +40,7 @@ if is_galore_available():
 
 
 if TYPE_CHECKING:
-    from transformers import PreTrainedModel, Seq2SeqTrainingArguments
+    from transformers import PreTrainedModel, Seq2SeqTrainingArguments, TrainerCallback
     from trl import AutoModelForCausalLMWithValueHead
 
     from ..hparams import DataArguments
@@ -457,3 +457,12 @@ def get_batch_logps(
     labels[labels == label_pad_token_id] = 0  # dummy token
     per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)
     return (per_token_logps * loss_mask).sum(-1), loss_mask.sum(-1)
+
+
+def get_swanlab_callback(finetuning_args: "FinetuningArguments") -> "TrainerCallback":
+    r"""
+    Gets the callback for logging to SwanLab.
+    """
+    from swanlab.integration.huggingface import SwanLabCallback
+
+    return SwanLabCallback()