Add flexible data source discovery to load_raw_dataset (#661)

[KKothuri]

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

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds magpie and nemotron entries to DATASET_CONFIGS with a _normalize_nemotron helper. Extends load_raw_dataset to support local directory inputs (recursive JSON/JSONL glob) and arbitrary HuggingFace datasets via an hf:<id>[:<subset>:<split>] spec, delegated to a new _load_hf_dataset helper. Integration tests cover all new resolution paths.

Changes

Flexible dataset source discovery

Layer / File(s)	Summary
Nemotron normalizer and new dataset presets src/speculators/data_generation/configs.py	Adds _normalize_nemotron to convert input/output fields into a conversations list, and registers magpie and nemotron entries in DATASET_CONFIGS.
HF spec helper and expanded load_raw_dataset resolution src/speculators/data_generation/preprocessing.py	Adds _load_hf_dataset to parse and validate hf: specs (defaulting split to train, raising on malformed specs or missing conversations). Reworks load_raw_dataset to handle local file, local directory, named preset, and hf: spec inputs in that order.
Integration tests for all resolution paths tests/integration/datagen/test_preprocessing.py	Adds tests for local file/directory loading, empty-directory error, named preset resolution with mocked load_dataset, hf: spec parsing, malformed-spec rejection, missing conversations column validation, DATASET_CONFIGS registration assertions, and _normalize_nemotron output.

Possibly related issues

• #661 [RFC]: Flexible data source discovery — This PR directly implements the resolution chain proposed in the RFC: local file, local directory, named preset, and hf: spec paths, including the Magpie/Nemotron registry additions and integration test coverage described there.

Possibly related PRs

• vllm-project/speculators#495: Introduced normalize_fn hooks and config-driven dataset handling in configs.py/preprocessing.py, which this PR extends with new presets and loading paths.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 70.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately summarizes the main change to flexible data source discovery.
Linked Issues check	✅ Passed	The changes implement the requested file, directory, preset, and hf: resolution chain, add the conversations guard, and cover the new paths with tests.
Out of Scope Changes check	✅ Passed	The added preset and test work stays aligned with the discovery and preset objectives and does not introduce unrelated scope.
Description check	✅ Passed	The description accurately matches the changes to dataset source resolution, new HF loading support, preset additions, and tests.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/speculators/data_generation/preprocessing.py`:
- Around line 671-687: Reject empty hf: suffixes in the preprocessing parser
before load_dataset is called. In preprocessing.py, update the hf: parsing logic
around the spec.removeprefix("hf:").split(":") match in the helper that assigns
hf_id, subset, and split so that trailing-colon forms like hf:org/name: and
hf:org/name:subset: raise the local ValueError instead of falling through with
an empty split. Add validation for empty subset/split segments in this branch
and cover it with a regression test for the trailing-colon cases.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 3da890cb-b47a-4f1f-9148-75e1484e4c32

📥 Commits

Reviewing files that changed from the base of the PR and between 5aa4da9 and def1cfc.

📒 Files selected for processing (3)

• src/speculators/data_generation/configs.py
• src/speculators/data_generation/preprocessing.py
• tests/integration/datagen/test_preprocessing.py

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @KKothuri.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[imargulis]

@KKothuri @shanjiaz
Following up on RFC#583 I mostly had the same approach implemented with a slight difference below:

The local-directory path loads every discovered file in a single pass:

load_dataset("json", data_files=[all .json/.jsonl], split="train")
which forces one Arrow schema across all files, so a bundle whose files/splits don't share an identical schema fails with DatasetGenerationError. This isn't limited to exotic data — even two files that share conversations but differ in one top-level column raise:

# a.jsonl: {"conversations": [...], "idx": 1}
# b.jsonl: {"conversations": [...], "lang": "en"}
load_dataset("json", data_files=["a.jsonl", "b.jsonl"], split="train")
#-> DatasetGenerationError: Couldn't cast ...

Loading each file/source separately and then concatenate_datasets on the common columns (dropping non-shared ones with a warning, and raising if no shared data column remains) handles heterogeneous bundles cleanly, and extends to parquet by grouping on schema. For homogeneous bundles it's identical to the current single-pass.

I have this implemented and tested (directory-bundle slice of RFC #583). Happy to (a) push it as a follow-up on top of this PR, or (b) hand you the diff to fold in — your and @shanjiaz's call. The hf: handling and presets here are great as-is.