Update rl.py

This commit is contained in:
Daniel Han 2025-03-10 05:00:45 -07:00
commit 440b3c8723

View file

@ -342,12 +342,12 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
# Check data collator if it's correct!
if "data_collator" in call_args and "train_dataset" in call_args:
data_collator_check = \
"if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in train_dataset.column_names):\n"\
" print('Unsloth: Changing data collator to `DataCollatorForLanguageModeling` since `labels` not found.)\n"\
"if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in train_dataset.column_names:\n"\
" print('Unsloth: Changing data collator to `DataCollatorForLanguageModeling` since `labels` not found.')\n"\
" data_collator = DataCollatorForLanguageModeling("\
"tokenizer = processing_class if 'processing_class' in locals() else tokenizer, mlm = False)\n"\
"elif isinstance(data_collator, DataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names):\n"\
" print('Unsloth: Changing data collator to `DataCollatorForSeq2Seq` since `labels` found.)\n"\
"elif isinstance(data_collator, DataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names:\n"\
" print('Unsloth: Changing data collator to `DataCollatorForSeq2Seq` since `labels` found.')\n"\
" data_collator = DataCollatorForSeq2Seq("\
"tokenizer = processing_class if 'processing_class' in locals() else tokenizer)\n"
extra_args += data_collator_check