File size: 7,955 Bytes
58258b8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | """
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)
|