fill_placeholders.py

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.9"
# dependencies = ["datasets>=2.19.0"]
# ///
"""Reconstruct masked DAPO / Skywork math questions and answers from the public sources.

The released Ultra math data masks the question and the expected answer for rows
that originate from two public datasets:

  - BytedTsinghua-SIA/DAPO-Math-17k
  - Skywork/Skywork-OR1-RL-Data

Each masked row carries an `_hf_question_placeholder` with the source row index
and the information needed to rebuild it. This script downloads those datasets
from Hugging Face and restores the question text and `expected_answer` (the
latter from the source row's `reward_model.ground_truth`).

Usage:
    ./fill_placeholders.py --input-dir masked/ --output-dir restored/
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

# DAPO wraps each question in a fixed instruction prompt; Skywork stores the bare
# question. Stripping the DAPO wrapper yields the question text our blend used.
DAPO = "BytedTsinghua-SIA/DAPO-Math-17k"
SKYWORK = "Skywork/Skywork-OR1-RL-Data"
HF_SOURCES = [(DAPO, "train"), (SKYWORK, "math")]

DAPO_PREFIX = (
    "Solve the following math problem step by step. The last line of your response "
    "should be of the form Answer: $Answer (without quotes) where $Answer is the "
    "answer to the problem."
)
DAPO_SUFFIX = 'Remember to put your answer on its own line after "Answer:".'

PLACEHOLDER_KEY = "_hf_question_placeholder"


def strip_dapo_wrapper(text: str) -> str:
    t = text
    if DAPO_PREFIX in t:
        t = t.split(DAPO_PREFIX, 1)[1]
    if DAPO_SUFFIX in t:
        t = t.rsplit(DAPO_SUFFIX, 1)[0]
    return t.strip()


def bare_question(dataset: str, content: str) -> str:
    if dataset == DAPO:
        return strip_dapo_wrapper(content)
    return content.strip()


def unwrap_answer(raw) -> str:
    """Bare answer from an HF row's reward_model.ground_truth (Skywork stores a
    JSON list-string like '["5"]'; DAPO stores it bare)."""
    if not isinstance(raw, str):
        if isinstance(raw, list) and raw:
            return str(raw[0])
        return str(raw)
    s = raw.strip()
    if (s.startswith("[") and s.endswith("]")) or (s.startswith("{") and s.endswith("}")):
        try:
            v = json.loads(s)
        except Exception:
            return s
        if isinstance(v, list) and v:
            return str(v[0])
        return str(v)
    return s


def reconstruct_question(ph: dict, bare: str) -> str:
    """Rebuild the question from the placeholder recipe and the public bare text."""
    if ph.get("mode") == "canonical":
        # The original text was reformatted, so we wrap the public bare with the
        # stored NVIDIA-added scaffolding (instruction wrapper + reasoning tag).
        return ph.get("lead", "") + bare + ph.get("trail", "")
    # "exact" (default): literal prefix/suffix reproduce the original text.
    return ph.get("prefix", "") + bare + ph.get("suffix", "")


def restore_row(row: dict, hf: dict) -> dict:
    ph = row.get(PLACEHOLDER_KEY)
    if not ph:
        return row
    ds = hf[(ph["dataset"], ph["split"])]
    src = ds[int(ph["row"])]
    question = reconstruct_question(ph, bare_question(ph["dataset"], src["prompt"][0]["content"]))
    answer = unwrap_answer((src.get("reward_model") or {}).get("ground_truth"))

    restored = dict(row)
    restored.pop(PLACEHOLDER_KEY, None)
    restored["question"] = question
    restored["expected_answer"] = answer
    rcp = restored.get("responses_create_params") or {}
    inp = rcp.get("input") if isinstance(rcp, dict) else None
    if isinstance(inp, list) and inp and isinstance(inp[0], dict):
        inp[0]["content"] = question
    # Restore the answer echoed in `matched_sources` provenance, if present.
    for ms in restored.get("matched_sources") or []:
        if isinstance(ms, dict) and "expected_answer" in ms:
            ms["expected_answer"] = answer
    return restored


def main() -> None:
    ap = argparse.ArgumentParser(description="Reconstruct masked DAPO/Skywork math questions.")
    ap.add_argument("--input-dir", required=True, type=Path, help="Dir of masked .jsonl files.")
    ap.add_argument("--output-dir", required=True, type=Path, help="Dir for restored .jsonl files.")
    args = ap.parse_args()

    files = sorted(args.input_dir.glob("*.jsonl"))
    if not files:
        ap.error(f"No .jsonl files in {args.input_dir}")

    from datasets import load_dataset

    hf = {(d, s): load_dataset(d, split=s) for d, s in HF_SOURCES}

    args.output_dir.mkdir(parents=True, exist_ok=True)
    for in_path in files:
        out_path = args.output_dir / in_path.name
        restored = 0
        total = 0
        with open(in_path) as fin, open(out_path, "w") as fout:
            for line in fin:
                line = line.strip()
                if not line:
                    continue
                total += 1
                row = json.loads(line)
                if PLACEHOLDER_KEY in row:
                    row = restore_row(row, hf)
                    restored += 1
                fout.write(json.dumps(row) + "\n")
        print(f"{in_path.name}: {restored}/{total} questions restored -> {out_path}")


if __name__ == "__main__":
    main()