diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index e5cb226433..24cd5e60af 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -107,7 +107,7 @@ import torch import numpy as np from contextlib import nullcontext from torch.nn import functional as F -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling +from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling torch_compile_options = {{ "epilogue_fusion" : True, @@ -358,8 +358,8 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): "from unsloth_zoo.vision_utils import UnslothVisionDataCollator\n"\ "if not isinstance(data_collator, UnslothVisionDataCollator):\n"\ " if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in train_dataset.column_names:\n"\ - " data_collator = DataCollatorForLanguageModeling(__tokenizer, mlm = False)\n"\ - " elif isinstance(data_collator, DataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names:\n"\ + " data_collator = TransformersDataCollatorForLanguageModeling(__tokenizer, mlm = False)\n"\ + " elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names:\n"\ " data_collator = DataCollatorForSeq2Seq(__tokenizer)\n"\ "else:\n"\ " if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False\n"\ @@ -374,7 +374,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): " if isinstance(data_collator, DataCollatorForSeq2Seq):\n"\ " data_collator = DataCollatorForSeq2Seq(__tokenizer.tokenizer)\n"\ " else:\n"\ - " data_collator = DataCollatorForLanguageModeling(__tokenizer.tokenizer, mlm = False)\n" + " data_collator = TransformersDataCollatorForLanguageModeling(__tokenizer.tokenizer, mlm = False)\n" extra_args += pad_check pass