Add flexible data source discovery to load_raw_dataset (#661) (#675)

Resolves #661.

## Summary

Extends `load_raw_dataset()` in
`src/speculators/data_generation/preprocessing.py` from two source types
to the four-step resolution chain requested in RFC #661, and promotes
Magpie/Nemotron to shared presets. Scope is intentionally limited to
**data source discovery** — schema/role normalization is left to a
separate effort.

## Changes

- `load_raw_dataset` now resolves a source in order:
  1. Local `.json`/`.jsonl` file (existing).
2. **Local directory** — recursively glob `*.json`/`*.jsonl` (sorted for
determinism) and load as one dataset; an empty directory raises an
actionable error.
  3. Named preset from `DATASET_CONFIGS` (existing).
4. **`hf:<id>[:<subset>:<split>]`** — load an arbitrary HuggingFace
dataset.
- New `_load_hf_dataset` helper parses the `hf:` spec (split defaults to
`train`). It **raises if the loaded dataset has no `conversations`
column**, since `_preprocess_batch` otherwise only warns and silently
yields zero samples — satisfying the RFC's "no silent data corruption"
requirement.
- Adds `magpie` (already in ShareGPT `conversations` format, no
`normalize_fn`) and `nemotron`
(`nvidia/Llama-Nemotron-Post-Training-Dataset`, `SFT`/`chat`, with
`_normalize_nemotron` building conversations from `input` + `output`)
presets. Both schemas were verified against the live HF datasets-server
before hardcoding.

All source types compose across mixed `--data` usage, e.g.:
```
--data sharegpt --data /my/shards/ --data hf:nvidia/HelpSteer2 --data extra.jsonl
```

### Note for the maintainer

Beyond the RFC's literal `hf:<id>:<subset>:<split>`, I also accept the
2-part form `hf:<id>:<split>` (subset omitted), since single-config
datasets with a non-`train` split (e.g.
`HuggingFaceH4/ultrachat_200k:train_sft`) are common. Happy to make it
strict if you'd prefer.

## Why this is not a duplicate

#637 (RFC #583) touches the same function but bundles directory/HF
discovery **together with** role/schema normalization. Per @shanjiaz's
scoping on #583 and the assignment of #661, this PR delivers the
**discovery slice only**, keeping it separate from normalization. The
`hf:` handling here is also a different, more explicit design: an
explicit `hf:` prefix (rather than treating any unmatched string as an
HF path) plus a hard failure on non-conversations data.

## Tests

New tests in `tests/integration/datagen/test_preprocessing.py`: local
file, recursive directory bundle (incl. nested `.json`), empty-directory
error, preset dispatch, unsupported source, `hf:` spec parsing
(parametrized), the conversations-format guard, malformed `hf:` specs,
and the new presets. The `hf:`/preset cases mock `load_dataset` so they
are deterministic and network-free.

Commands run (in the project venv):

```
python -m pytest tests/integration/datagen/test_preprocessing.py -q
# 51 passed, 2 skipped

ruff check  src/speculators/data_generation/{preprocessing,configs}.py tests/integration/datagen/test_preprocessing.py   # All checks passed
ruff format --check <same files>                                                                                          # already formatted
python -m mypy src/speculators/data_generation/{preprocessing,configs}.py                                                 # Success: no issues
```

## AI assistance disclosure

AI assistance (Claude Code) was used to help write this change. I (the
submitter) have reviewed every changed line, understand the change
end-to-end, and ran the tests above.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Signed-off-by: Karthik Kothuri <karthikkothuri2009@gmail.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: KKothuri

src/speculators/data_generation/configs.py

@@ -37,6 +37,16 @@ def _normalize_gsm8k(example: dict) -> dict:
     }
 
 
+def _normalize_nemotron(example: dict) -> dict:
+    """Build a conversation from Nemotron ``input`` turns plus ``output``."""
+    return {
+        "conversations": [
+            *example["input"],
+            {"role": "assistant", "content": example["output"]},
+        ]
+    }
+
+
 def get_coco_dir():
     return os.getenv("COCO_DIR") or "coco/"
 
@@ -110,6 +120,18 @@ def _normalize_sharegpt4v_coco(example: dict) -> dict:
         split="train",
         normalize_fn=_normalize_gsm8k,
     ),
+    "magpie": DatasetConfig(
+        name="magpie",
+        hf_path="Magpie-Align/Magpie-Llama-3.1-Pro-300K-Filtered",
+        split="train",
+    ),
+    "nemotron": DatasetConfig(
+        name="nemotron",
+        hf_path="nvidia/Llama-Nemotron-Post-Training-Dataset",
+        subset="SFT",
+        split="chat",
+        normalize_fn=_normalize_nemotron,
+    ),
     # NOTE: You need to serve vLLM with `--allowed-local-media-path /path/to/coco`
     "sharegpt4v_coco": DatasetConfig(
         name="sharegpt4v_coco",

src/speculators/data_generation/preprocessing.py

@@ -651,26 +651,114 @@ def build_eagle3_dataset(
     return dataset
 
 
+def _load_hf_dataset(spec: str) -> tuple[HFDataset, None]:
+    """Load an arbitrary HuggingFace dataset from an ``hf:`` spec.
+
+    Args:
+        spec: ``hf:<dataset_id>[:<subset>:<split>]``. The split defaults to
+            ``train``. A single suffix (``hf:<id>:<split>``) selects a split
+            without a subset; both can be given as ``hf:<id>:<subset>:<split>``.
+
+    Returns:
+        Tuple of (raw_dataset, None). No normalize_fn is applied: the dataset
+        must already be in conversations format.
+
+    Raises:
+        ValueError: If the spec is malformed or the loaded dataset has no
+            ``conversations`` column.
+    """
+    subset: str | None
+    match spec.removeprefix("hf:").split(":"):
+        case [hf_id]:
+            subset, split = None, "train"
+        case [hf_id, split]:
+            subset = None
+        case [hf_id, subset, split]:
+            pass
+        case _:
+            raise ValueError(
+                f"Invalid hf: spec '{spec}'. "
+                f"Expected hf:<dataset_id>[:<subset>:<split>]."
+            )
+
+    if not hf_id:
+        raise ValueError(f"Invalid hf: spec '{spec}': missing dataset id.")
+    if subset == "":
+        raise ValueError(f"Invalid hf: spec '{spec}': empty subset.")
+    if not split:
+        raise ValueError(f"Invalid hf: spec '{spec}': empty split.")
+
+    raw_dataset = load_dataset(hf_id, name=subset, split=split)
+
+    if "conversations" not in raw_dataset.column_names:
+        raise ValueError(
+            f"HuggingFace dataset '{hf_id}' (split '{split}') is not in "
+            f"conversations format: expected a 'conversations' column but found "
+            f"{raw_dataset.column_names}. Pass a dataset already in conversations "
+            f"format, or add a preset to DATASET_CONFIGS with a normalize_fn."
+        )
+
+    return raw_dataset, None
+
+
 def load_raw_dataset(
     train_data_path: str,
 ) -> tuple[HFDataset, Callable[[dict], dict] | None]:
-    """Load raw dataset from local file or HuggingFace."""
+    """Load a raw dataset from one of several source types.
+
+    Resolution order:
+        1. Local ``.json``/``.jsonl`` file.
+        2. Local directory: recursively load all ``*.json``/``*.jsonl`` files
+           as a single dataset.
+        3. Named preset from ``DATASET_CONFIGS``.
+        4. ``hf:<id>[:<subset>:<split>]`` for an arbitrary HuggingFace dataset.
+
+    Args:
+        train_data_path: File path, directory path, preset name, or ``hf:`` spec.
+
+    Returns:
+        Tuple of (raw_dataset, normalize_fn). normalize_fn is None for sources
+        already in conversations format.
+
+    Raises:
+        ValueError: If the source cannot be resolved or a local directory
+            contains no ``.json``/``.jsonl`` files.
+    """
+    # 1. Local file
     if train_data_path.endswith((".jsonl", ".json")):
         return load_dataset("json", data_files=train_data_path, split="train"), None
 
-    if train_data_path not in DATASET_CONFIGS:
-        raise ValueError(
-            f"Unsupported dataset: {train_data_path}. "
-            f"Supported: local .json/.jsonl files or {list(DATASET_CONFIGS.keys())}"
+    # 2. Local directory
+    path = Path(train_data_path)
+    if path.is_dir():
+        data_files = sorted(
+            str(p) for p in (*path.rglob("*.json"), *path.rglob("*.jsonl"))
         )
+        if not data_files:
+            raise ValueError(
+                f"No .json/.jsonl files found in directory: {train_data_path}"
+            )
+        return load_dataset("json", data_files=data_files, split="train"), None
 
-    config = DATASET_CONFIGS[train_data_path]
-    raw_dataset = load_dataset(config.hf_path, name=config.subset, split=config.split)
-
-    if config.filter_fn is not None:
-        raw_dataset = raw_dataset.filter(config.filter_fn)
-
-    return raw_dataset, config.normalize_fn
+    # 3. Named preset
+    if train_data_path in DATASET_CONFIGS:
+        config = DATASET_CONFIGS[train_data_path]
+        raw_dataset = load_dataset(
+            config.hf_path, name=config.subset, split=config.split
+        )
+        if config.filter_fn is not None:
+            raw_dataset = raw_dataset.filter(config.filter_fn)
+        return raw_dataset, config.normalize_fn
+
+    # 4. Arbitrary HuggingFace dataset
+    if train_data_path.startswith("hf:"):
+        return _load_hf_dataset(train_data_path)
+
+    raise ValueError(
+        f"Unsupported dataset: {train_data_path}. Supported: local .json/.jsonl "
+        f"file, local directory of .json/.jsonl files, hf:<id>[:<subset>:<split>], "
+        f"or a preset {list(DATASET_CONFIGS.keys())}."
+    )
 
 
 def get_tokenizer(processor: ProcessorLike):

tests/integration/datagen/test_preprocessing.py

@@ -5,25 +5,32 @@
 import json
 import re
 from typing import Any
+from unittest.mock import patch
 
 import pytest
 import torch
 from datasets import Dataset as HFDataset
 from PIL import Image
 from transformers import AutoTokenizer
 
+from speculators.data_generation.configs import (
+    DATASET_CONFIGS,
+    _normalize_nemotron,
+)
 from speculators.data_generation.preprocessing import (
     _adapt_conv_for_hf,
     _adapt_conv_for_vllm,
     _create_loss_mask_from_offsets,
     _detect_assistant_pattern,
+    _load_hf_dataset,
     _normalize_conversation,
     _preprocess_batch,
     _supports_assistant_mask,
     build_eagle3_dataset,
     get_tokenizer,
     load_and_preprocess_dataset,
     load_processor,
+    load_raw_dataset,
 )
 
 # Test model from HuggingFace with chat template
@@ -1397,6 +1404,187 @@ def test_normalize_conversation_tool_calls_with_empty_content():
     assert assistant_tool_call_turn["tool_calls"] == tool_calls
 
 
+# Tests for load_raw_dataset resolution chain (issue #661)
+
+PREFIX = "speculators.data_generation.preprocessing"
+
+
+def _write_jsonl(path, rows):
+    """Write rows as newline-delimited JSON to path."""
+    path.write_text("\n".join(json.dumps(r) for r in rows))
+
+
+def _conv_row(prompt: str) -> dict:
+    """Return one conversations-format row for the given prompt."""
+    return {
+        "conversations": [
+            {"from": "human", "value": prompt},
+            {"from": "gpt", "value": f"answer to {prompt}"},
+        ]
+    }
+
+
+@pytest.mark.sanity
+def test_load_raw_dataset_local_file(tmp_path):
+    """A local .jsonl file loads with no normalize_fn."""
+    data_file = tmp_path / "data.jsonl"
+    _write_jsonl(data_file, [_conv_row("a"), _conv_row("b")])
+
+    dataset, normalize_fn = load_raw_dataset(str(data_file))
+
+    assert len(dataset) == 2
+    assert normalize_fn is None
+
+
+@pytest.mark.sanity
+def test_load_raw_dataset_local_directory(tmp_path):
+    """A directory of .json/.jsonl shards loads as one combined dataset."""
+    _write_jsonl(tmp_path / "shard1.jsonl", [_conv_row("a"), _conv_row("b")])
+    _write_jsonl(tmp_path / "shard2.jsonl", [_conv_row("c")])
+    # Nested file is discovered recursively; .json extension also matched.
+    nested = tmp_path / "nested"
+    nested.mkdir()
+    _write_jsonl(nested / "shard3.json", [_conv_row("d")])
+
+    dataset, normalize_fn = load_raw_dataset(str(tmp_path))
+
+    assert len(dataset) == 4
+    assert normalize_fn is None
+
+
+@pytest.mark.sanity
+def test_load_raw_dataset_empty_directory_raises(tmp_path):
+    """A directory with no .json/.jsonl files raises an actionable error."""
+    empty_dir = tmp_path / "empty"
+    empty_dir.mkdir()
+
+    with pytest.raises(ValueError, match="No .json/.jsonl files found"):
+        load_raw_dataset(str(empty_dir))
+
+
+@pytest.mark.sanity
+def test_load_raw_dataset_named_preset():
+    """A named preset resolves through DATASET_CONFIGS to load_dataset."""
+    sentinel = HFDataset.from_list([_conv_row("x")])
+    with patch(f"{PREFIX}.load_dataset", return_value=sentinel) as mock_load:
+        dataset, normalize_fn = load_raw_dataset("sharegpt")
+
+    config = DATASET_CONFIGS["sharegpt"]
+    mock_load.assert_called_once_with(
+        config.hf_path, name=config.subset, split=config.split
+    )
+    assert dataset is sentinel
+    assert normalize_fn is config.normalize_fn
+
+
+@pytest.mark.sanity
+def test_load_raw_dataset_unsupported_source_raises():
+    """An unknown source that is not a file/dir/preset/hf: spec raises."""
+    with pytest.raises(ValueError, match="Unsupported dataset"):
+        load_raw_dataset("not_a_real_preset")
+
+
+@pytest.mark.sanity
+@pytest.mark.parametrize(
+    ("spec", "expected_id", "expected_name", "expected_split"),
+    [
+        ("hf:org/name", "org/name", None, "train"),
+        ("hf:org/name:custom_split", "org/name", None, "custom_split"),
+        ("hf:org/name:subset:custom_split", "org/name", "subset", "custom_split"),
+    ],
+)
+def test_load_hf_dataset_spec_parsing(spec, expected_id, expected_name, expected_split):
+    """hf: specs parse into (id, subset, split) and call load_dataset."""
+    sentinel = HFDataset.from_list([_conv_row("x")])
+    with patch(f"{PREFIX}.load_dataset", return_value=sentinel) as mock_load:
+        dataset, normalize_fn = load_raw_dataset(spec)
+
+    mock_load.assert_called_once_with(
+        expected_id, name=expected_name, split=expected_split
+    )
+    assert dataset is sentinel
+    assert normalize_fn is None
+
+
+# A small public dataset already in conversations format, used to exercise the
+# real hf: download path end to end (the RFC asks for a small public dataset).
+HF_CONV_DATASET = "philschmid/guanaco-sharegpt-style"
+
+
+@pytest.mark.sanity
+def test_load_raw_dataset_hf_real_download():
+    """End-to-end hf: load of a small public conversations dataset.
+
+    Skipped when the dataset cannot be fetched, e.g. network-restricted CI
+    runners that only have a pre-baked model cache. The hf: parsing and the
+    conversations guard are covered deterministically by the mocked tests above.
+    """
+    try:
+        dataset, normalize_fn = load_raw_dataset(f"hf:{HF_CONV_DATASET}")
+    except Exception as exc:  # noqa: BLE001
+        pytest.skip(f"Could not fetch {HF_CONV_DATASET}: {exc}")
+
+    assert normalize_fn is None
+    assert "conversations" in dataset.column_names
+    assert len(dataset) > 0
+
+
+@pytest.mark.sanity
+def test_load_hf_dataset_non_conversations_raises():
+    """An hf: dataset without a conversations column fails loudly."""
+    non_conv = HFDataset.from_list([{"prompt": "hi", "response": "there"}])
+    with (
+        patch(f"{PREFIX}.load_dataset", return_value=non_conv),
+        pytest.raises(ValueError, match="conversations format"),
+    ):
+        _load_hf_dataset("hf:org/name")
+
+
+@pytest.mark.sanity
+@pytest.mark.parametrize(
+    "spec",
+    [
+        "hf:org/name:a:b:c",  # too many segments
+        "hf:",  # missing id
+        "hf:org/name:",  # trailing colon -> empty split
+        "hf:org/name:subset:",  # empty split with subset
+        "hf:org/name::split",  # empty subset
+    ],
+)
+def test_load_hf_dataset_malformed_spec_raises(spec):
+    """Malformed hf: specs raise locally without touching the network."""
+    with (
+        patch(f"{PREFIX}.load_dataset") as mock_load,
+        pytest.raises(ValueError, match="Invalid hf: spec"),
+    ):
+        _load_hf_dataset(spec)
+    mock_load.assert_not_called()
+
+
+@pytest.mark.sanity
+def test_dataset_configs_has_magpie_and_nemotron():
+    """magpie and nemotron presets are registered."""
+    assert "magpie" in DATASET_CONFIGS
+    assert "nemotron" in DATASET_CONFIGS
+    # magpie ships conversations already, so needs no normalize_fn.
+    assert DATASET_CONFIGS["magpie"].normalize_fn is None
+
+
+@pytest.mark.sanity
+def test_normalize_nemotron_builds_conversations():
+    """nemotron normalize_fn appends the output as an assistant turn."""
+    example = {
+        "input": [{"role": "user", "content": "hi"}],
+        "output": "hello",
+    }
+    result = _normalize_nemotron(example)
+
+    assert result["conversations"] == [
+        {"role": "user", "content": "hi"},
+        {"role": "assistant", "content": "hello"},
+    ]
+
+
 # Tests for load_and_preprocess_dataset