Fix: Add support for TRL native dataset formats

unsloth/models/rl_replacements.py

@@ -147,22 +147,35 @@ def sft_trainer_prepare_dataset(function_name, function):
         "if 'skip_prepare_dataset' in locals() and skip_prepare_dataset:\n"
         "    return dataset\n"
         "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"
-        "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) "
+        "tokenizer_call = None\n"
+        "unsloth_has_native_format = False\n"
+        "if hasattr(dataset, 'column_names'):\n"
+        "    columns = set(dataset.column_names)\n"
+        "    unsloth_has_native_format = (\n"
+        "        {'prompt', 'completion'}.issubset(columns) or\n"
+        "        {'prompt', 'chosen', 'rejected'}.issubset(columns) or\n"
+        "        'messages' in columns\n"
+        "    )\n"
+        "if not unsloth_has_native_format:\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"
+        "    try:\n"
+        "        _first_item = dataset[0]\n"
+        "    except Exception:\n"
+        "        _first_item = next(iter(dataset))\n"
+        "    test_text = _first_item[dataset_text_field] if (formatting_func is None and dataset_text_field is not None) else formatting_func(_first_item)[0]\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) "
         "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"
-        "    tokenizer_call = tokenizer.__call__\n"
-        "    tokenizer.__call__ = partial(tokenizer_call, add_special_tokens = False)\n"
-        "    processing_class = tokenizer\n"
-        "else:\n"
-        "    tokenizer_call = None\n"
-        "    add_special_tokens = False if has_bos_token_already else locals().get('add_special_tokens', False)\n"
+        "    if 'add_special_tokens' not in locals() and has_bos_token_already:\n"
+        "        from functools import partial\n"
+        "        tokenizer_call = tokenizer.__call__\n"
+        "        tokenizer.__call__ = partial(tokenizer_call, add_special_tokens = False)\n"
+        "        processing_class = tokenizer\n"
+        "    else:\n"
+        "        add_special_tokens = False if has_bos_token_already else locals().get('add_special_tokens', False)\n"
     )
 
     check_text = check_text.split("\n")