feat(regen): make every registry preset usable on-policy

The regen script previously accepted only presets with a bare-prompt
column (magpie/ultrachat/gsm8k) — a historical artifact of its
Magpie-specific origin, not an architectural limit. Extraction now
mirrors the off-policy ingestion via prepare_row(): filter_fn sees the
raw row, normalize_fn is merged over it with HF map semantics, and
system/user turns are extracted from the result. Every DATASET_CONFIGS
preset is therefore a valid --dataset choice, and any dataset that
works off-policy works on-policy by construction (nemotron's
input/output schema, sharegpt's native conversations, ...).

Resume identity is untouched: _primary_identifier still hashes the raw
row, so --resume stays valid for outputs from earlier runs. For the
three previous presets the extracted turns are unchanged (gsm8k's
normalizer yields the same single user turn its prompt_field fallback
did); prompt_field remains as a per-row fallback for rows whose
conversation is unusable.

Known limits deferred to follow-ups rather than gated: sharegpt4v_coco
needs an image-capable endpoint and local COCO files (same requirement
and failure mode as off-policy), and tool-use datasets get plain-text
regeneration until tool-call regen lands.

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

Author: WindChimeRan

scripts/response_regeneration/script.py

@@ -19,16 +19,11 @@
     with_retries,
 )
 
-# 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():
     """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",
@@ -43,7 +38,7 @@ def parse_args():
     parser.add_argument(
         "--dataset",
         default="ultrachat",
-        choices=REGEN_DATASETS,
+        choices=list(DATASET_CONFIGS),
         help="Dataset to process",
     )
     parser.add_argument(
@@ -146,6 +141,19 @@ def extract_turns(row, prompt_field):
     return []
 
 
+def prepare_row(row, config):
+    """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, "")
@@ -412,7 +420,6 @@ async def main():
     # Get dataset configuration
     dataset_config = DATASET_CONFIGS[args.dataset]
     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.split
@@ -431,7 +438,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()
@@ -488,8 +495,8 @@ 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)
+                # prepare_row returns [] for filtered rows and unusable turns.
                 if not turns:
                     continue
 

src/speculators/data_generation/configs.py

@@ -20,7 +20,7 @@ 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.
+    # Bare user-prompt fallback column for the response-regeneration script.
     prompt_field: str | None = None
 
 

tests/unit/scripts/test_response_regeneration.py

@@ -15,7 +15,7 @@
 import pytest
 
 from speculators.data_generation import vllm_client
-from speculators.data_generation.configs import DATASET_CONFIGS
+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
 
@@ -531,17 +531,49 @@ def test_worker_row_identity_and_all_or_nothing_writes(tmp_path):
 
 
 # ---------------------------------------------------------------------------
-# 6. Dataset presets come from the shared registry.
+# 6. Every shared-registry preset works on-policy (off-policy parity).
 # ---------------------------------------------------------------------------
 
 
-def test_regen_presets_come_from_shared_registry():
+def test_script_consumes_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
+
+
+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 eligible["magpie"] == "instruction"
-    assert eligible["ultrachat"] == "prompt"
-    assert eligible["gsm8k"] == "question"
+    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"}
+    ]