diff --git a/studio/backend/tests/test_dataset_preference_format.py b/studio/backend/tests/test_dataset_preference_format.py new file mode 100644 index 0000000000..36ba927491 --- /dev/null +++ b/studio/backend/tests/test_dataset_preference_format.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from datasets import Dataset + +from utils.datasets.dataset_utils import ( + check_dataset_format, + format_and_template_dataset, + format_dataset, +) + + +def _preference_dataset(): + return Dataset.from_list( + [ + { + "prompt": "The sky is", + "chosen": " blue.", + "rejected": " green.", + } + ] + ) + + +def test_check_dataset_format_detects_preference_dataset(): + result = check_dataset_format(_preference_dataset()) + + assert result["detected_format"] == "preference" + assert result["requires_manual_mapping"] is False + + +def test_check_dataset_format_requires_review_when_prompt_is_missing(): + dataset = Dataset.from_list([{"chosen": " blue.", "rejected": " green."}]) + + result = check_dataset_format(dataset) + + assert result["detected_format"] == "unknown" + assert result["requires_manual_mapping"] is True + + +def test_format_dataset_auto_preserves_preference_dataset(): + result = format_dataset(_preference_dataset(), format_type = "auto") + + assert result["detected_format"] == "preference" + assert result["final_format"] == "preference" + assert result["requires_manual_mapping"] is False + assert result["dataset"][0]["chosen"] == " blue." + + +def test_format_dataset_preference_mapping_converts_custom_columns(): + dataset = Dataset.from_list( + [{"question": "The sky is", "winner": " blue.", "loser": " green."}] + ) + + result = format_dataset( + dataset, + format_type = "preference", + custom_format_mapping = { + "question": "prompt", + "winner": "chosen", + "loser": "rejected", + }, + ) + + assert result["final_format"] == "preference" + assert result["dataset"][0] == { + "prompt": "The sky is", + "chosen": " blue.", + "rejected": " green.", + } + + +def test_format_and_template_dataset_skips_chat_template_for_preference(): + result = format_and_template_dataset( + _preference_dataset(), + model_name = "unsloth/Qwen2.5-0.5B-Instruct", + tokenizer = object(), + format_type = "auto", + ) + + assert result["success"] is True + assert result["final_format"] == "preference" + assert result["dataset"][0] == { + "prompt": "The sky is", + "chosen": " blue.", + "rejected": " green.", + }