""" 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)