sdg/preprocessing/preprocess_rstarcoder.py

"""
Preprocess rStar-Coder (seed_sft subset).

Source: https://huggingface.co/datasets/microsoft/rStar-Coder

Dataset fields:
  - question_id:  Unique identifier linking problems with their test cases
  - question:     The original competitive programming problem description
  - starter_code: Initial code template provided with the problem
  - response:     The long CoT solution
  - code:         The implementation code for the solution
  - verified:     Whether the solution has been verified with test cases
  - is_passed:    Whether the solution has passed the test cases

This script deduplicates by `question_id` (picking one random rollout per
question via reservoir sampling), asserts correctness, and uploads to
HuggingFace in the `openthoughts4` conversations format so the SDG configs
can consume it.
"""

from __future__ import annotations

import argparse
import json
import random
import tempfile
from pathlib import Path

from datasets import load_dataset
from huggingface_hub import HfApi, create_repo
from tqdm import tqdm


DATASET_ID = "microsoft/rStar-Coder"
DATASET_SUBSET = "seed_sft"
# Set to an int to assert the expected number of unique question_ids; leave as
# None on first run, then update after observing the printed count.
EXPECTED_UNIQUE_PROMPTS: int | None = None
HF_REPO_ID = "teetone/rstarcoder-seedsft-deduped"


def build_prompt(question: str, starter_code: str | None) -> str:
    """Combine the problem description with the starter code template."""
    question = (question or "").strip()
    starter = (starter_code or "").strip()
    if starter:
        return f"{question}\n\n{starter}"
    return question


def preprocess_and_upload(seed: int = 42, dry_run: bool = False) -> None:
    random.seed(seed)

    print(f"Loading {DATASET_ID} (subset={DATASET_SUBSET}, streaming) ...")
    ds = load_dataset(DATASET_ID, DATASET_SUBSET, split="train", streaming=True)

    # Single streaming pass: reservoir-sample one row per unique question_id.
    # For each question we track (chosen_row, count_seen) so we can do uniform
    # random selection without storing all rollouts.
    qid_reservoir: dict[str, tuple[dict, int]] = {}
    total = 0
    skipped_no_qid = 0

    pbar = tqdm(ds, desc="Streaming rStar-Coder", unit="row", smoothing=0.05)
    for row in pbar:
        total += 1

        qid = row.get("question_id")
        if qid is None:
            skipped_no_qid += 1
            continue
        qid = str(qid)

        prompt = build_prompt(row.get("question"), row.get("starter_code"))
        if not prompt:
            continue

        rec = {
            "conversations": [
                {"from": "human", "value": prompt},
                {"from": "gpt", "value": row.get("response", "")},
            ],
        }

        if qid not in qid_reservoir:
            qid_reservoir[qid] = (rec, 1)
        else:
            _, count = qid_reservoir[qid]
            count += 1
            # Reservoir sampling: replace with probability 1/count
            if random.randint(1, count) == 1:
                qid_reservoir[qid] = (rec, count)
            else:
                qid_reservoir[qid] = (qid_reservoir[qid][0], count)

        pbar.set_postfix(unique=len(qid_reservoir), skipped=skipped_no_qid)
    pbar.close()

    print(f"  Total rows streamed: {total:,}")
    print(f"  Skipped (no question_id): {skipped_no_qid:,}")
    print(f"  Unique question_ids: {len(qid_reservoir):,}")

    records_by_qid = [row for row, _ in qid_reservoir.values()]
    print(f"  After question_id dedup: {len(records_by_qid):,}")

    # ── Second-pass dedup on question text ────────────────────────────
    # Some rows can share an identical question while having different
    # question_ids (e.g. the same problem reused under multiple sources).
    # Drop those so only one record per unique question survives.
    records: list[dict] = []
    seen_questions: set[str] = set()
    duplicate_questions = 0
    for rec in records_by_qid:
        prompt = rec["conversations"][0]["value"]
        if prompt in seen_questions:
            duplicate_questions += 1
            continue
        seen_questions.add(prompt)
        records.append(rec)
    print(f"  Removed by question-text dedup: {duplicate_questions:,}")
    print(f"  Final unique questions: {len(records):,}")

    # ── Assertions ────────────────────────────────────────────────────
    if EXPECTED_UNIQUE_PROMPTS is not None:
        assert len(records) == EXPECTED_UNIQUE_PROMPTS, (
            f"Expected {EXPECTED_UNIQUE_PROMPTS} rows, got {len(records)}"
        )
    assert len(seen_questions) == len(records), "Question-text dedup invariant failed"
    print("  All assertions passed.")

    # ── Sample preview ────────────────────────────────────────────────
    if dry_run:
        n_preview = min(5, len(records))
        print()
        print(f"──────── Sample preview ({n_preview} records that WOULD upload) ────────")
        for i, rec in enumerate(random.sample(records, n_preview)):
            print(f"\n[Sample {i + 1}/{n_preview}]")
            print(json.dumps(rec, ensure_ascii=False, indent=2)[:2000])
            if len(json.dumps(rec, ensure_ascii=False)) > 2000:
                print("  ... (truncated to 2000 chars)")
        print("─────────────────────────────────────────────────────────────────")

    # ── Upload to HuggingFace ─────────────────────────────────────────
    if dry_run:
        print("[dry-run] Skipping upload to HuggingFace.")
    else:
        print(f"Uploading to {HF_REPO_ID} ...")
        api = HfApi()
        create_repo(HF_REPO_ID, repo_type="dataset", exist_ok=True)

        with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f:
            temp_path = f.name
            for rec in records:
                f.write(json.dumps(rec, ensure_ascii=False) + "\n")

        api.upload_file(
            path_or_fileobj=temp_path,
            path_in_repo="data/train-00000-of-00001.jsonl",
            repo_id=HF_REPO_ID,
            repo_type="dataset",
        )
        Path(temp_path).unlink()

    print()
    print("──────── Final stats ────────")
    print(f"  Total rows streamed:           {total:,}")
    print(f"  Skipped (no question_id):      {skipped_no_qid:,}")
    print(f"  Unique question_ids:           {len(qid_reservoir):,}")
    print(f"  Removed by question-text dedup: {duplicate_questions:,}")
    print(f"  Final unique questions:        {len(records):,}")
    if dry_run:
        print(f"  Rows that WOULD upload:        {len(records):,} (dry-run, not uploaded)")
        print(f"  HF dataset (target):           https://huggingface.co/datasets/{HF_REPO_ID}")
    else:
        print(f"  Rows uploaded to HF:           {len(records):,}")
        print(f"  HF dataset:                    https://huggingface.co/datasets/{HF_REPO_ID}")
    print("─────────────────────────────")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Run preprocessing and assertions but skip the HuggingFace upload.",
    )
    parser.add_argument("--seed", type=int, default=42, help="Random seed for reservoir sampling.")
    args = parser.parse_args()
    preprocess_and_upload(seed=args.seed, dry_run=args.dry_run)