diff --git a/docs/cli/response_regeneration.md b/docs/cli/response_regeneration.md index 389200da0..8df2be845 100644 --- a/docs/cli/response_regeneration.md +++ b/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 diff --git a/scripts/response_regeneration/script.py b/scripts/response_regeneration/script.py index a3689323b..54d24cb3d 100644 --- a/scripts/response_regeneration/script.py +++ b/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 diff --git a/src/speculators/data_generation/configs.py b/src/speculators/data_generation/configs.py index 3f4ef4951..8cfc39e18 100644 --- a/src/speculators/data_generation/configs.py +++ b/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,6 +114,7 @@ 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", @@ -119,11 +122,13 @@ def _normalize_sharegpt4v_coco(example: dict) -> dict: 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", diff --git a/tests/unit/scripts/test_response_regeneration.py b/tests/unit/scripts/test_response_regeneration.py index d097a0dfd..b9b924ff9 100644 --- a/tests/unit/scripts/test_response_regeneration.py +++ b/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": "", + } + 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"