File size: 4,283 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 | """
Preprocess OpenThoughts3-1.2M (code domain only).
Source: https://huggingface.co/datasets/open-thoughts/OpenThoughts3-1.2M
This script filters the code domain, deduplicates by user prompt
(picking one random rollout per prompt via reservoir sampling),
asserts correctness, and uploads to HuggingFace.
"""
from __future__ import annotations
import json
import random
import tempfile
from pathlib import Path
from datasets import load_dataset
from huggingface_hub import HfApi, create_repo
DATASET_ID = "open-thoughts/OpenThoughts3-1.2M"
EXPECTED_UNIQUE_PROMPTS = 5_714
HF_REPO_ID = "teetone/openthoughts3-code-deduped"
def extract_prompt(conversations: list[dict]) -> str | None:
"""Return the first human turn from a conversation."""
for turn in conversations:
if turn["from"] == "human":
return turn["value"]
return None
def preprocess_and_upload(seed: int = 42) -> None:
random.seed(seed)
print(f"Loading {DATASET_ID} (streaming) ...")
ds = load_dataset(DATASET_ID, split="train", streaming=True)
# Single streaming pass: reservoir-sample one row per unique prompt.
# For each prompt we track (chosen_row, count_seen) so we can do
# uniform random selection without storing all rollouts.
prompt_reservoir: dict[str, tuple[dict, int]] = {}
total = 0
code_count = 0
for row in ds:
total += 1
if row["domain"] != "code":
continue
code_count += 1
prompt = extract_prompt(row["conversations"])
if prompt is None:
continue
rec = {
"difficulty": row["difficulty"],
"source": row["source"],
"domain": row["domain"],
"conversations": row["conversations"],
}
if prompt not in prompt_reservoir:
prompt_reservoir[prompt] = (rec, 1)
else:
_, count = prompt_reservoir[prompt]
count += 1
# Reservoir sampling: replace with probability 1/count
if random.randint(1, count) == 1:
prompt_reservoir[prompt] = (rec, count)
else:
prompt_reservoir[prompt] = (prompt_reservoir[prompt][0], count)
if code_count % 100_000 == 0:
print(f" Processed {code_count:,} code rows, {len(prompt_reservoir):,} unique prompts ...")
print(f" Total rows streamed: {total:,}")
print(f" Code rows: {code_count:,}")
print(f" Unique prompts: {len(prompt_reservoir):,}")
records = [row for row, _ in prompt_reservoir.values()]
print(f" Deduplicated rows: {len(records):,}")
# ── Assertions ────────────────────────────────────────────────────
assert len(records) == EXPECTED_UNIQUE_PROMPTS, (
f"Expected {EXPECTED_UNIQUE_PROMPTS} rows, got {len(records)}"
)
seen_prompts: set[str] = set()
for rec in records:
assert rec["domain"] == "code", f"Non-code row found: domain={rec['domain']}"
prompt = extract_prompt(rec["conversations"])
assert prompt not in seen_prompts, f"Duplicate prompt found: {prompt[:80]}..."
seen_prompts.add(prompt)
assert len(seen_prompts) == EXPECTED_UNIQUE_PROMPTS, (
f"Expected {EXPECTED_UNIQUE_PROMPTS} unique prompts, got {len(seen_prompts)}"
)
print(" All assertions passed.")
# ── Upload to HuggingFace ─────────────────────────────────────────
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(f"Done! Dataset available at: https://huggingface.co/datasets/{HF_REPO_ID}")
if __name__ == "__main__":
preprocess_and_upload()
|