refactor(regen): one dataset registry, every preset usable both on&off-policy

[WindChimeRan]

Purpose

Dataset-registry half of RFC #702. The regen script kept its own 3-entry dataset dict alongside the shared registry in data_generation/configs.py, so adding a dataset means editing both files in two formats (e.g. #769, #771). Unblocked since #716: the script already imports the package.

• Commit 1 (pure refactor): the script consumes the shared registry; DatasetConfig gains prompt_field, the bare-prompt fallback column. No behavior change.
• Commit 2: new prepare_row() applies a preset's filter_fn/normalize_fn the same way prepare-data does, so every registry preset is a valid --dataset — datasets work in both pipelines by construction. --resume identity and the original three presets' extracted turns are unchanged. Multimodal payloads (sharegpt4v_coco) and tool-call regen semantics are follow-ups.

Tests

• pytest tests/unit/scripts tests/unit/data_generation: 48 passed (new: nemotron normalize-parity, filter_fn, map-merge fallback)
• ruff check, ruff format --check, mypy --check-untyped-defs: clean
• CLI: --help offers all presets; --dataset sharegpt flows through endpoint auto-detection

Checklist

I have filled in:

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan/results, such as providing test command and pasting the results.
• (Optional) The necessary documentation update.
• I (a human) have written or reviewed the code in this pr to the best of my ability.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 958fcddc-08e5-4daf-8d6f-c934b0e1b424

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Changes

The response regeneration script now consumes shared dataset presets, applies preset filtering and normalization before turn extraction, and uses attribute-based metadata. Dataset prompt fields, tests, and CLI documentation were updated accordingly.

Dataset-configured response regeneration

Layer / File(s)	Summary
Dataset configuration contract src/speculators/data_generation/configs.py	DatasetConfig supports an optional prompt field, and the ultrachat, gsm8k, and magpie presets define their source fields.
Response regeneration flow scripts/response_regeneration/script.py, tests/unit/scripts/test_response_regeneration.py	The script uses shared configs and helpers, prepares rows with filtering and normalization, reads config attributes, and tests the new behavior.
CLI dataset documentation docs/cli/response_regeneration.md	CLI arguments and supported dataset presets are documented, including HuggingFace identifiers, default splits, and the sharegpt4v_coco requirement.

Possibly related issues

• vllm-project/speculators#584 — Concerns the same response-regeneration script and dataset filtering/normalization flow.
• vllm-project/speculators#583 — Covers the shared DatasetConfig and DATASET_CONFIGS changes used by this update.

Possibly related PRs

• vllm-project/speculators#675 — Updates the same dataset configuration registry and DatasetConfig infrastructure.
• vllm-project/speculators#693 — Modifies the same response-regeneration turn extraction path and its unit tests.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.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 clearly summarizes the refactor toward a shared dataset registry and broader preset support.
Description check	✅ Passed	The description matches the refactor, added prompt_field, prepare_row, docs, and the stated validation/tests.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ 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]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

docs/cli/response_regeneration.md (1)

72-78: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Document --subset in the Data Arguments section The CLI exposes --subset in scripts/response_regeneration/script.py, but it’s missing from this list between --split and --limit. Add a short entry for completeness.

🤖 Prompt for 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.

In `@docs/cli/response_regeneration.md` around lines 72 - 78, Add a `--subset`
entry to the Data Arguments section between `--split` and `--limit`, briefly
describing its purpose and matching the existing option-format style in the
response regeneration documentation.

🧹 Nitpick comments (1)

scripts/response_regeneration/script.py (1)

144-154: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Consider adding type annotations to prepare_row for mypy consistency.

Adjacent functions like _primary_identifier annotate row: dict[str, Any], while prepare_row and extract_turns are unannotated. Annotating prepare_row would let mypy verify config.filter_fn, config.normalize_fn, and config.prompt_field attribute access at call sites.

As per coding guidelines, verify type annotations are consistent with mypy requirements for **/*.py files.

♻️ Optional: add type annotations

-from speculators.data_generation.configs import DatasetConfig
+from typing import Any
+
+from speculators.data_generation.configs import DatasetConfig

-def prepare_row(row, config):
+def prepare_row(row: dict[str, Any], config: DatasetConfig) -> list[dict[str, str]]:
     """Extract regeneration turns from a raw dataset row, ``[]`` to skip it.

🤖 Prompt for 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.

In `@scripts/response_regeneration/script.py` around lines 144 - 154, Annotate
prepare_row with the repository’s established row and config types, and add the
appropriate return type for its extracted turns or skip sentinel. Reuse existing
type aliases or definitions used by _primary_identifier, extract_turns, and the
filter_fn/normalize_fn configuration fields so mypy can validate attribute
access without introducing duplicate types.

🤖 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.

Outside diff comments:
In `@docs/cli/response_regeneration.md`:
- Around line 72-78: Add a `--subset` entry to the Data Arguments section
between `--split` and `--limit`, briefly describing its purpose and matching the
existing option-format style in the response regeneration documentation.

---

Nitpick comments:
In `@scripts/response_regeneration/script.py`:
- Around line 144-154: Annotate prepare_row with the repository’s established
row and config types, and add the appropriate return type for its extracted
turns or skip sentinel. Reuse existing type aliases or definitions used by
_primary_identifier, extract_turns, and the filter_fn/normalize_fn configuration
fields so mypy can validate attribute access without introducing duplicate
types.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 60b004a2-b340-413e-b550-0681b7625c9f

📥 Commits

Reviewing files that changed from the base of the PR and between 6030a44 and b0be6d8.

📒 Files selected for processing (4)

• docs/cli/response_regeneration.md
• scripts/response_regeneration/script.py
• src/speculators/data_generation/configs.py
• tests/unit/scripts/test_response_regeneration.py

[orestis-z]

LGTM. Clean dedup of the dataset registry per RFC #702 — verified all three original presets produce identical results through the new prepare_row path, resume identity is preserved, and import weight is unchanged (script already imports from speculators). Agree with @coderabbitai's suggestion to document --subset in the Data Arguments section. Recommend approving.

🤖 Generated with Claude Code using the /pr-review skill

[orestis-z]

LGTM. Both prior items addressed: --subset is documented and extract_turns/prepare_row are annotated. The multimodal gating (MULTIMODAL_DATASETS + _dataset_choice) is sound — the test exhaustively verifies every preset is accounted for. Recommend approving.

🤖 Generated with Claude Code using the /pr-review skill

[rahul-tuli]

LGTM!