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

docs/cli/response_regeneration.md

@@ -69,9 +69,11 @@ python scripts/response_regeneration/script.py --dataset magpie
 
 #### Data Arguments
 
-- **`--dataset`** (str, default: `ultrachat`, choices: `magpie`, `ultrachat`) Dataset to process.
+- **`--dataset`** (str, default: `ultrachat`) Dataset preset to process (see [Supported Datasets](#supported-datasets)).
 
-- **`--split`** (str, default: dataset-specific) Dataset split. Defaults to `train` for magpie and `train_sft` for ultrachat.
+- **`--split`** (str, default: preset-specific) Dataset split. Defaults to the preset's split.
+
+- **`--subset`** (str, default: preset-specific) Dataset subset/config name. Defaults to the preset's subset.
 
 - **`--limit`** (int, default: `None`) Stop after N rows.
 
@@ -112,10 +114,18 @@ python scripts/response_regeneration/script.py \
 
 ## Supported Datasets
 
-| Dataset   | HuggingFace ID                                    | Prompt Field  | Default Split |
-| --------- | ------------------------------------------------- | ------------- | ------------- |
-| Magpie    | `Magpie-Align/Magpie-Llama-3.1-Pro-300K-Filtered` | `instruction` | `train`       |
-| UltraChat | `HuggingFaceH4/ultrachat_200k`                    | `prompt`      | `train_sft`   |
+The text presets from the shared dataset registry (`DATASET_CONFIGS` in `speculators/data_generation/configs.py`) — the same ones `prepare-data` accepts:
+
+| Dataset             | HuggingFace ID                                    | Default Split |
+| ------------------- | ------------------------------------------------- | ------------- |
+| `sharegpt`          | `Aeala/ShareGPT_Vicuna_unfiltered`                | `train`       |
+| `ultrachat`         | `HuggingFaceH4/ultrachat_200k`                    | `train_sft`   |
+| `gsm8k`             | `openai/gsm8k`                                    | `train`       |
+| `magpie`            | `Magpie-Align/Magpie-Llama-3.1-Pro-300K-Filtered` | `train`       |
+| `nemotron`          | `nvidia/Llama-Nemotron-Post-Training-Dataset`     | `chat`        |
+| `open-perfectblend` | `mlabonne/open-perfectblend`                      | `train`       |
+
+The registry's multimodal preset, `sharegpt4v_coco`, is **off-policy only** and `--dataset` rejects it. Its turns carry image content parts, which the Chat Completions API rejects, and the pre-tokenized output row has nowhere to keep pixel data. Use it with `prepare-data`.
 
 ## Output Format
 

scripts/response_regeneration/script.py

@@ -12,36 +12,33 @@
 from datasets import load_dataset
 from tqdm import tqdm
 
+from speculators.data_generation.configs import DATASET_CONFIGS, DatasetConfig
 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",
-    },
-}
+# On-policy regeneration has no multimodal support yet; off-policy `prepare-data`
+# does, so these presets are gated here rather than dropped from the registry.
+MULTIMODAL_DATASETS = {"sharegpt4v_coco"}
+REGEN_DATASETS = [name for name in DATASET_CONFIGS if name not in MULTIMODAL_DATASETS]
+
+
+def _dataset_choice(name: str) -> str:
+    """Reject multimodal presets with a reason, not a bare invalid choice."""
+    if name in MULTIMODAL_DATASETS:
+        raise argparse.ArgumentTypeError(
+            f"{name!r} is multimodal; on-policy regeneration does not support "
+            "images yet. Use it off-policy with `prepare-data`."
+        )
+    return name
 
 
 def parse_args():
     """Parse command-line arguments for the script."""
     parser = argparse.ArgumentParser(
-        description="Regenerate responses from Magpie instructions via vLLM Chat API."
+        description="Regenerate dataset responses via a vLLM Chat API endpoint."
     )
     parser.add_argument(
         "--endpoint",
@@ -56,7 +53,8 @@ def parse_args():
     parser.add_argument(
         "--dataset",
         default="ultrachat",
-        choices=list(DATASET_CONFIGS.keys()),
+        type=_dataset_choice,
+        choices=REGEN_DATASETS,
         help="Dataset to process",
     )
     parser.add_argument(
@@ -138,7 +136,9 @@ def sanitize_filename(name: str) -> str:
     return name.strip("._")
 
 
-def extract_turns(row, prompt_field):
+def extract_turns(
+    row: dict[str, Any], prompt_field: str | None
+) -> list[dict[str, Any]]:
     """Extract ordered system/user turns from a dataset row.
 
     Multi-turn conversations are read from a ``messages`` or ``conversations``
@@ -175,6 +175,19 @@ def extract_turns(row, prompt_field):
     return []
 
 
+def prepare_row(row: dict[str, Any], config: DatasetConfig) -> list[dict[str, Any]]:
+    """Extract regeneration turns from a raw dataset row, ``[]`` to skip it.
+
+    Mirrors off-policy ingestion: ``filter_fn`` sees the raw row, and
+    ``normalize_fn`` is merged over it (HF ``map`` semantics keep raw columns).
+    """
+    if config.filter_fn is not None and not config.filter_fn(row):
+        return []
+    if config.normalize_fn is not None:
+        row = {**row, **config.normalize_fn(row)}
+    return extract_turns(row, config.prompt_field)
+
+
 def _is_present(value: Any) -> bool:
     """Return True for a usable identifier (not None / not empty string)."""
     return value not in (None, "")
@@ -442,12 +455,11 @@ 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
 
     # 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:
@@ -462,7 +474,7 @@ async def main():
 
     print(f"Using dataset: {dataset_id}")
     print(f"Split: {split}")
-    print(f"Prompt field: {prompt_field}")
+    print(f"Prompt field: {dataset_config.prompt_field}")
     print(f"Output file: {args.outfile}")
     print(f"Error file: {error_outfile}")
     print()
@@ -519,8 +531,7 @@ async def main():
                 if args.language_filter and row.get("language") != args.language_filter:
                     continue
 
-                turns = extract_turns(row, prompt_field)
-                # extract_turns returns [] when there is no usable user turn.
+                turns = prepare_row(row, dataset_config)
                 if not turns:
                     continue
 

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, used when a row has no conversation.
+    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

@@ -7,6 +7,7 @@
 The script is not a package, so it is imported by path.
 """
 
+import argparse
 import asyncio
 import importlib.util
 import json
@@ -16,6 +17,7 @@
 import pytest
 
 from speculators.data_generation import vllm_client
+from speculators.data_generation.configs import DATASET_CONFIGS, DatasetConfig
 from speculators.data_generation.preprocessing import _preprocess_batch
 from speculators.data_generation.vllm_client import InvalidResponseError
 
@@ -529,3 +531,53 @@ 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. Every shared-registry preset works on-policy (off-policy parity).
+# ---------------------------------------------------------------------------
+
+
+def test_prepare_row_normalizes_like_off_policy():
+    # nemotron rows only become extractable through the preset's normalize_fn.
+    row = {
+        "input": [{"role": "user", "content": "Hi"}],
+        "output": "<original answer to drop>",
+    }
+    assert regen.prepare_row(row, DATASET_CONFIGS["nemotron"]) == [
+        {"role": "user", "content": "Hi"}
+    ]
+
+
+def test_prepare_row_applies_filter_fn():
+    config = DatasetConfig(
+        name="t",
+        hf_path="t",
+        split="train",
+        filter_fn=lambda row: row["keep"],
+    )
+    row = {"keep": False, "conversations": [{"role": "user", "content": "Hi"}]}
+    assert regen.prepare_row(row, config) == []
+    assert regen.prepare_row(row | {"keep": True}, config) == [
+        {"role": "user", "content": "Hi"}
+    ]
+
+
+def test_prepare_row_merges_normalize_output_over_raw_row():
+    # HF map merges columns: normalize output must not clobber the raw fallback.
+    config = DatasetConfig(
+        name="t",
+        hf_path="t",
+        split="train",
+        normalize_fn=lambda row: {"conversations": []},
+        prompt_field="prompt",
+    )
+    assert regen.prepare_row({"prompt": "Hi"}, config) == [
+        {"role": "user", "content": "Hi"}
+    ]
+
+
+def test_dataset_choice_rejects_multimodal_with_a_reason():
+    with pytest.raises(argparse.ArgumentTypeError, match="does not support images"):
+        regen._dataset_choice("sharegpt4v_coco")
+    assert regen._dataset_choice("ultrachat") == "ultrachat"