From 6530a66a823ba13545615f00c67bccd726d7e6b7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 9 Nov 2024 17:01:33 -0800 Subject: [PATCH] Update tokenizer_utils.py --- unsloth/tokenizer_utils.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index c639dbf1a0..6d0ee548c1 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -1001,13 +1001,14 @@ def patch_sft_trainer_tokenizer(): # Also DPO weirdly tokenizes non numeric columns? Delete them! check_text += \ "\n"\ - "column_names = set(self.train_dataset.column_names)\n"\ - "check = ['chosen', 'rejected', 'prompt', 'chosen_input_ids', 'chosen_attention_mask',\n"\ - " 'chosen_labels', 'rejected_input_ids', 'rejected_attention_mask', 'rejected_labels',\n"\ - " 'prompt_input_ids', 'prompt_attention_mask']\n"\ - "if all(x in column_names for x in check):\n"\ - " self.train_dataset = self.train_dataset.remove_columns(['chosen', 'rejected', 'prompt'])\n"\ - "del check, column_names\n"\ + "if hasattr(self.train_dataset, 'column_names''):\n" + " column_names = set(self.train_dataset.column_names)\n"\ + " check = ['chosen', 'rejected', 'prompt', 'chosen_input_ids', 'chosen_attention_mask',\n"\ + " 'chosen_labels', 'rejected_input_ids', 'rejected_attention_mask', 'rejected_labels',\n"\ + " 'prompt_input_ids', 'prompt_attention_mask']\n"\ + " if all(x in column_names for x in check):\n"\ + " self.train_dataset = self.train_dataset.remove_columns(['chosen', 'rejected', 'prompt'])\n"\ + " del check, column_names\n"\ "\n" check_text = check_text.split("\n")