From fca83182afa44ab69bf79e65078f2dc4456cc8de Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 24 Mar 2026 05:27:59 -0700 Subject: [PATCH] fix: handle prompt/completion datasets in slow-path BOS detection (#4548) * 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 * fix: preserve chat_template BOS check when test_text is None The has_bos_token_already guard wrapped both test_text.startswith() and bos_token in chat_template with test_text is not None, which disabled the chat_template BOS detection for conversational datasets where test_text is set to None. Split the guard so test_text is not None only applies to the startswith() call, while bos_token in chat_template is always checked. --- unsloth/models/rl_replacements.py | 13 ++++++++++--- unsloth/tokenizer_utils.py | 13 ++++++++++--- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 805086324e..9f555416d4 100755 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -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" diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index c445879df7..96c22f62ff 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -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"