fix: handle prompt/completion datasets in slow-path BOS detection

The slow-path check_text blocks in rl_replacements.py and
tokenizer_utils.py crash when a prompt/completion dataset is used
because they unconditionally access dataset[0][dataset_text_field]
even when the dataset does not have a text field.

This fixes both files to:
- Default dataset_text_field to None instead of raising when undefined
- Detect prompt/completion columns and concatenate them for BOS check
- Guard with isinstance(str) on both prompt and completion to handle
  conversational format (list of dicts) by setting test_text to None
- Add test_text is not None guard on has_bos_token_already to prevent
  AttributeError on NoneType.startswith()

This is the slow-path complement to unslothai/unsloth-zoo#560 which
fixes the fast-path in sft_prepare_dataset.

Closes #4486
This commit is contained in:
Daniel Han 2026-03-24 09:42:52 +00:00
commit b6ce5786d0
2 changed files with 20 additions and 6 deletions

View file

@ -153,11 +153,18 @@ def sft_trainer_prepare_dataset(function_name, function):
"if 'tokenizer' not in locals(): tokenizer = processing_class\n"
"if 'formatting_func' not in locals(): raise RuntimeError('Unsloth: Please file a bug report - `formatting_func` does not exist!')\n"
"if 'dataset_text_field' not in locals() and 'args' in locals(): dataset_text_field = args.dataset_text_field\n"
"if 'dataset_text_field' not in locals(): raise RuntimeError('Unsloth: Please file a bug report - `dataset_text_field` does not exist!')\n"
"test_text = dataset[0][dataset_text_field] if (formatting_func is None and dataset_text_field is not None) else formatting_func(dataset[0])[0]\n"
"if 'dataset_text_field' not in locals(): dataset_text_field = None\n"
"if formatting_func is None and dataset_text_field is None and 'prompt' in dataset[0] and 'completion' in dataset[0]:\n"
" test_text = (dataset[0]['prompt'] + dataset[0]['completion']) if (isinstance(dataset[0]['prompt'], str) and isinstance(dataset[0]['completion'], str)) else None\n"
"elif formatting_func is None and dataset_text_field is not None:\n"
" test_text = dataset[0][dataset_text_field]\n"
"elif formatting_func is not None:\n"
" test_text = formatting_func(dataset[0])[0]\n"
"else:\n"
" test_text = None\n"
"chat_template = getattr(tokenizer, 'chat_template', None)\n"
"chat_template = '' if chat_template is None else chat_template\n"
"has_bos_token_already = (test_text.startswith(tokenizer.bos_token) or tokenizer.bos_token in chat_template) "
"has_bos_token_already = (test_text is not None and (test_text.startswith(tokenizer.bos_token) or tokenizer.bos_token in chat_template)) "
"if getattr(tokenizer, 'bos_token', None) is not None else False\n"
"if 'add_special_tokens' not in locals() and has_bos_token_already:\n"
" from functools import partial\n"

View file

@ -974,11 +974,18 @@ def patch_sft_trainer_tokenizer():
"if 'tokenizer' not in locals(): tokenizer = processing_class\n"
"if 'formatting_func' not in locals(): raise RuntimeError('Unsloth: Please file a bug report - `formatting_func` does not exist!')\n"
"if 'dataset_text_field' not in locals() and 'args' in locals(): dataset_text_field = args.dataset_text_field\n"
"if 'dataset_text_field' not in locals(): raise RuntimeError('Unsloth: Please file a bug report - `dataset_text_field` does not exist!')\n"
"test_text = dataset[0][dataset_text_field] if (formatting_func is None and dataset_text_field is not None) else formatting_func(dataset[0])[0]\n"
"if 'dataset_text_field' not in locals(): dataset_text_field = None\n"
"if formatting_func is None and dataset_text_field is None and 'prompt' in dataset[0] and 'completion' in dataset[0]:\n"
" test_text = (dataset[0]['prompt'] + dataset[0]['completion']) if (isinstance(dataset[0]['prompt'], str) and isinstance(dataset[0]['completion'], str)) else None\n"
"elif formatting_func is None and dataset_text_field is not None:\n"
" test_text = dataset[0][dataset_text_field]\n"
"elif formatting_func is not None:\n"
" test_text = formatting_func(dataset[0])[0]\n"
"else:\n"
" test_text = None\n"
"chat_template = getattr(tokenizer, 'chat_template', None)\n"
"chat_template = '' if chat_template is None else chat_template\n"
"has_bos_token_already = (test_text.startswith(tokenizer.bos_token) or tokenizer.bos_token in chat_template) "
"has_bos_token_already = (test_text is not None and (test_text.startswith(tokenizer.bos_token) or tokenizer.bos_token in chat_template)) "
"if getattr(tokenizer, 'bos_token', None) is not None else False\n"
"if 'add_special_tokens' not in locals() and has_bos_token_already:\n"
" from functools import partial\n"