refactor(data): consume the shared dataset registry in the regen script

The response-regeneration script kept its own plain-dict DATASET_CONFIGS
(magpie/ultrachat/gsm8k) alongside the DatasetConfig registry in
speculators.data_generation.configs, so adding a dataset meant two edits
in two formats. Since #716 the script imports
speculators.data_generation.vllm_client anyway, so consuming the shared
registry adds no import weight.

DatasetConfig gains an optional prompt_field — the bare-prompt fallback
column used when a row carries no conversation — and presets defining it
are the script's --dataset choices, which stay exactly
magpie/ultrachat/gsm8k with unchanged HF paths, splits, and subsets.

Dataset-registry half of RFC #702.

Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>

Author: WindChimeRan

scripts/response_regeneration/script.py

@@ -12,30 +12,17 @@
 from datasets import load_dataset
 from tqdm import tqdm
 
+from speculators.data_generation.configs import DATASET_CONFIGS
 from speculators.data_generation.vllm_client import (
     DEFAULT_MAX_RETRIES,
     InvalidResponseError,
     with_retries,
 )
 
-DATASET_CONFIGS = {
-    "magpie": {
-        "id": "Magpie-Align/Magpie-Llama-3.1-Pro-300K-Filtered",
-        "prompt_field": "instruction",
-        "default_split": "train",
-    },
-    "ultrachat": {
-        "id": "HuggingFaceH4/ultrachat_200k",
-        "prompt_field": "prompt",
-        "default_split": "train_sft",
-    },
-    "gsm8k": {
-        "id": "openai/gsm8k",
-        "prompt_field": "question",
-        "default_split": "train",
-        "subset": "main",
-    },
-}
+# Presets usable for regeneration: those defining a bare-prompt fallback column.
+REGEN_DATASETS = [
+    name for name, config in DATASET_CONFIGS.items() if config.prompt_field
+]
 
 
 def parse_args():
@@ -56,7 +43,7 @@ def parse_args():
     parser.add_argument(
         "--dataset",
         default="ultrachat",
-        choices=list(DATASET_CONFIGS.keys()),
+        choices=REGEN_DATASETS,
         help="Dataset to process",
     )
     parser.add_argument(
@@ -424,12 +411,12 @@ async def main():
 
     # Get dataset configuration
     dataset_config = DATASET_CONFIGS[args.dataset]
-    dataset_id = dataset_config["id"]
-    prompt_field = dataset_config["prompt_field"]
+    dataset_id = dataset_config.hf_path
+    prompt_field = dataset_config.prompt_field
 
     # Use dataset-specific defaults if not provided
-    split = args.split if args.split is not None else dataset_config["default_split"]
-    subset = args.subset if args.subset is not None else dataset_config.get("subset")
+    split = args.split if args.split is not None else dataset_config.split
+    subset = args.subset if args.subset is not None else dataset_config.subset
 
     # Generate output filename if not specified
     if args.outfile is None:

src/speculators/data_generation/configs.py

@@ -20,6 +20,8 @@ class DatasetConfig:
     split: str
     filter_fn: Callable[[dict], bool] | None = None
     normalize_fn: Callable[[dict], dict] | None = None
+    # Bare user-prompt column; presets defining it are regen --dataset choices.
+    prompt_field: str | None = None
 
 
 def _normalize_ultrachat(example: dict) -> dict:
@@ -112,18 +114,21 @@ def _normalize_sharegpt4v_coco(example: dict) -> dict:
         hf_path="HuggingFaceH4/ultrachat_200k",
         split="train_sft",
         normalize_fn=_normalize_ultrachat,
+        prompt_field="prompt",
     ),
     "gsm8k": DatasetConfig(
         name="gsm8k",
         hf_path="openai/gsm8k",
         subset="main",
         split="train",
         normalize_fn=_normalize_gsm8k,
+        prompt_field="question",
     ),
     "magpie": DatasetConfig(
         name="magpie",
         hf_path="Magpie-Align/Magpie-Llama-3.1-Pro-300K-Filtered",
         split="train",
+        prompt_field="instruction",
     ),
     "nemotron": DatasetConfig(
         name="nemotron",

tests/unit/scripts/test_response_regeneration.py

@@ -15,6 +15,7 @@
 import pytest
 
 from speculators.data_generation import vllm_client
+from speculators.data_generation.configs import DATASET_CONFIGS
 from speculators.data_generation.preprocessing import _preprocess_batch
 from speculators.data_generation.vllm_client import InvalidResponseError
 
@@ -527,3 +528,20 @@ def test_worker_row_identity_and_all_or_nothing_writes(tmp_path):
     error = json.loads(err_path.read_text())
     assert error["id"] == "conv-abc"
     assert error["metadata"]["turns_completed"] == 1
+
+
+# ---------------------------------------------------------------------------
+# 6. Dataset presets come from the shared registry.
+# ---------------------------------------------------------------------------
+
+
+def test_regen_presets_come_from_shared_registry():
+    # Identity, not equality: a re-declared copy in the script would drift.
+    assert regen.DATASET_CONFIGS is DATASET_CONFIGS
+    # The original CLI presets stay regen-eligible, with their fallback columns.
+    eligible = {
+        name: DATASET_CONFIGS[name].prompt_field for name in regen.REGEN_DATASETS
+    }
+    assert eligible["magpie"] == "instruction"
+    assert eligible["ultrachat"] == "prompt"
+    assert eligible["gsm8k"] == "question"