diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 129052e9ae0..4329f15488f 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -32,6 +32,11 @@ if TYPE_CHECKING: from transformers import PreTrainedTokenizerBase + +def _join_messages_content(sample: dict) -> str: + return "\n".join(turn["content"] for turn in sample["messages"]) + + # Use dict to store the config for each dataset. # If we want to export more options to user like target languages, we need more standardized approach like dataclass. SUPPORTED_DATASET_CONFIG: dict[str, Any] = { @@ -61,7 +66,7 @@ "path": "nvidia/Nemotron-Post-Training-Dataset-v2", "split": ["stem", "chat", "math", "code"], }, - "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "preprocess": _join_messages_content, "chat_key": "messages", }, "nemotron-post-training-dataset-v1": { @@ -69,7 +74,91 @@ "path": "nvidia/Nemotron-Post-Training-Dataset-v1", "split": ["stem", "chat", "math", "code", "tool_calling"], }, - "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "preprocess": _join_messages_content, + "chat_key": "messages", + }, + "nemotron-sft-instruction-following-chat-v2": { + # Skips ``reasoning_on`` split: heterogeneous messages schema fails streaming cast. + "config": { + "path": "nvidia/Nemotron-SFT-Instruction-Following-Chat-v2", + "split": ["reasoning_off"], + }, + "preprocess": _join_messages_content, + "chat_key": "messages", + }, + "nemotron-science-v1": { + "config": { + "path": "nvidia/Nemotron-Science-v1", + "split": ["MCQ", "RQA"], + }, + "preprocess": _join_messages_content, + "chat_key": "messages", + }, + "nemotron-competitive-programming-v1": { + # Skips ``infinibyte_part0[0|1]``: heterogeneous schema fails streaming cast. + "config": { + "path": "nvidia/Nemotron-Competitive-Programming-v1", + "split": [ + "competitive_coding_cpp_part00", + "competitive_coding_cpp_part01", + "competitive_coding_python_part00", + "competitive_coding_python_part01", + ], + }, + "preprocess": _join_messages_content, + "chat_key": "messages", + }, + "nemotron-sft-agentic-v2": { + # Skips ``search`` split: heterogeneous messages schema fails streaming cast. + "config": { + "path": "nvidia/Nemotron-SFT-Agentic-v2", + "split": ["interactive_agent", "tool_calling"], + }, + "preprocess": _join_messages_content, + "chat_key": "messages", + }, + "nemotron-math-v2": { + "config": { + "path": "nvidia/Nemotron-Math-v2", + "split": ["high_part00", "high_part01", "high_part02", "medium", "low"], + }, + "preprocess": _join_messages_content, + "chat_key": "messages", + }, + "nemotron-sft-swe-v2": { + # Skips ``openhands_swe`` split: heterogeneous schema fails streaming cast. + "config": { + "path": "nvidia/Nemotron-SFT-SWE-v2", + "split": ["agentless"], + }, + "preprocess": _join_messages_content, + "chat_key": "messages", + }, + "nemotron-sft-multilingual-v1": { + "config": { + "path": "nvidia/Nemotron-SFT-Multilingual-v1", + "split": [ + "code_de", + "code_es", + "code_fr", + "code_it", + "code_ja", + "code_zh", + "math_de", + "math_es", + "math_fr", + "math_it", + "math_ja", + "math_zh", + "stem_de", + "stem_es", + "stem_fr", + "stem_it", + "stem_ja", + "stem_zh", + ], + }, + "preprocess": _join_messages_content, "chat_key": "messages", }, "magpie": { diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index f89663d89b5..7566088a1ad 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -17,6 +17,7 @@ import pytest import torch +from huggingface_hub import get_token from torch.utils.data import DataLoader from modelopt.torch.utils.dataset_utils import ( @@ -689,3 +690,57 @@ def test_dataloader_mixing_hf_and_local_jsonl(self, tmp_path, pad_tokenizer): ) batches = list(loader) assert sum(b["input_ids"].shape[0] for b in batches) == 5 + + +_NEW_NEMOTRON_KEYS = [ + "nemotron-sft-instruction-following-chat-v2", + "nemotron-science-v1", + "nemotron-competitive-programming-v1", + "nemotron-sft-agentic-v2", + "nemotron-math-v2", + "nemotron-sft-swe-v2", + "nemotron-sft-multilingual-v1", +] + + +@pytest.mark.parametrize("dataset_key", _NEW_NEMOTRON_KEYS) +def test_new_nemotron_registry_shape(dataset_key): + """Always-on shape check on the 7 newly registered nvidia/Nemotron-* entries. + + Complements the gated smoke test below — catches typos in dataset paths or + split names even when the runner has no HF credentials. + """ + from modelopt.torch.utils.dataset_utils import SUPPORTED_DATASET_CONFIG + + assert dataset_key in SUPPORTED_DATASET_CONFIG + entry = SUPPORTED_DATASET_CONFIG[dataset_key] + config = entry["config"] + assert config["path"].startswith("nvidia/Nemotron-") + splits = config["split"] + assert isinstance(splits, list) and splits + assert all(isinstance(s, str) and s for s in splits) + assert len(set(splits)) == len(splits) + assert callable(entry["preprocess"]) + assert entry["chat_key"] == "messages" + + +@pytest.mark.integration +@pytest.mark.parametrize("dataset_key", _NEW_NEMOTRON_KEYS) +def test_get_dataset_samples_new_nemotron(dataset_key): + """Smoke-test the 7 newly registered nvidia/Nemotron-* calibration datasets. + + Skipped when no HF token is available because these datasets live behind the HF Hub. + ``huggingface_hub.get_token()`` covers both the ``HF_TOKEN`` env var and tokens + cached by ``hf auth login``. + """ + pytest.importorskip("datasets") + if not get_token(): + pytest.skip( + "No HF token (env HF_TOKEN or `hf auth login`); skipping gated Nemotron smoke test" + ) + + samples = get_dataset_samples(dataset_key, num_samples=2) + + assert isinstance(samples, list) + assert len(samples) == 2 + assert all(isinstance(s, str) and len(s) > 0 for s in samples)