Add option to check dataset labels in SFTTrainer

trl/trainer/sft_trainer.py

@@ -36,6 +36,7 @@
 from transformers.modeling_utils import unwrap_model
 from transformers.trainer_callback import TrainerCallback
 from transformers.trainer_utils import EvalPrediction
+from transformers.utils import logging
 
 from ..extras.dataset_formatting import get_formatting_func_from_dataset
 from ..import_utils import is_peft_available
@@ -47,6 +48,7 @@
     trl_sanitze_kwargs_for_tagging,
 )
 
+logger = logging.get_logger(__name__)
 
 if is_peft_available():
     from peft import PeftConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training
@@ -118,6 +120,8 @@ class SFTTrainer(Trainer):
             Dict of Optional kwargs to pass when creating packed or non-packed datasets
         eval_packing: (`Optional[bool]`, *optional*):
             Whether to pack the eval dataset as well. Defaults to `packing` if `None` is passed.
+        check_dataset_labels (`Optional[bool]`):
+            Flag to enable debugging of dataset labels and tokenization. If set to True, the trainer will print the tokens, decoded tokens, and their corresponding labels for the first item in the training dataset during initialization. Defaults to False.
     """
 
     _tag_names = ["trl", "sft"]
@@ -149,6 +153,7 @@ def __init__(
         model_init_kwargs: Optional[Dict] = None,
         dataset_kwargs: Optional[Dict] = None,
         eval_packing: Optional[bool] = None,
+        check_dataset_labels: Optional[bool] = False,
     ):
         if model_init_kwargs is None:
             model_init_kwargs = {}
@@ -302,6 +307,15 @@ def make_inputs_require_grad(module, input, output):
                 "overflow issues when training a model in half-precision. You might consider adding `tokenizer.padding_side = 'right'` to your code."
             )
 
+
+        if check_dataset_labels:
+            if train_dataset is not None and len(train_dataset) > 0:
+                input_ids, attention_mask, labels = data_collator([train_dataset[0]]).values()
+                logger.info("check_dataset_labels:")  # noqa
+                logger.info(tokenizer.decode(input_ids[0]))  # noqa
+                for token, label in zip(input_ids[0], labels[0]):
+                    logger.info(f"{token.item()}, '{tokenizer.decode(token)}', {label.item()}")  # noqa
+
         super().__init__(
             model=model,
             args=args,