wehub-resource-sync
113 lines
3.7 KiB
Python
113 lines
3.7 KiB
Python
"""
|
|
Build preference data for the reward model and DPO from real public datasets:
|
|
- Anthropic/hh-rlhf (helpful/harmless human preferences)
|
|
- HuggingFaceH4/ultrafeedback_binarized (LLM-judged preference pairs)
|
|
|
|
Emits JSONL of ``{"prompt", "chosen", "rejected"}`` (train + held-out test). The held-out
|
|
test split is what reward-model preference accuracy is measured on.
|
|
|
|
Example:
|
|
PYTHONPATH=. HF_HOME=/ephemeral/hf_cache python scripts/prepare_preference_data.py \
|
|
--source both --max_per_source 40000 --out_dir /ephemeral/data
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
|
|
os.environ.setdefault("HF_HOME", "/ephemeral/hf_cache")
|
|
|
|
_ASSISTANT_MARKER = "\n\nAssistant:"
|
|
|
|
|
|
def _split_hh(text: str) -> tuple[str, str] | None:
|
|
"""Split an HH-RLHF conversation string into (prompt_context, final_response)."""
|
|
idx = text.rfind(_ASSISTANT_MARKER)
|
|
if idx == -1:
|
|
return None
|
|
prompt = text[:idx].strip()
|
|
response = text[idx + len(_ASSISTANT_MARKER):].strip()
|
|
if not prompt or not response:
|
|
return None
|
|
return prompt, response
|
|
|
|
|
|
def from_hh(max_n: int, split: str) -> list[dict]:
|
|
from datasets import load_dataset
|
|
|
|
ds = load_dataset("Anthropic/hh-rlhf", split=split)
|
|
if max_n:
|
|
ds = ds.select(range(min(max_n, len(ds))))
|
|
out = []
|
|
for ex in ds:
|
|
c = _split_hh(ex["chosen"]); r = _split_hh(ex["rejected"])
|
|
if not c or not r:
|
|
continue
|
|
prompt, chosen = c
|
|
_, rejected = r
|
|
if chosen == rejected:
|
|
continue
|
|
out.append({"prompt": prompt, "chosen": chosen, "rejected": rejected})
|
|
return out
|
|
|
|
|
|
def from_ultrafeedback(max_n: int, split: str) -> list[dict]:
|
|
from datasets import load_dataset
|
|
|
|
hf_split = "train_prefs" if split == "train" else "test_prefs"
|
|
ds = load_dataset("HuggingFaceH4/ultrafeedback_binarized", split=hf_split)
|
|
if max_n:
|
|
ds = ds.select(range(min(max_n, len(ds))))
|
|
out = []
|
|
for ex in ds:
|
|
prompt = ex["prompt"].strip()
|
|
chosen = ex["chosen"][-1]["content"].strip()
|
|
rejected = ex["rejected"][-1]["content"].strip()
|
|
if not prompt or not chosen or not rejected or chosen == rejected:
|
|
continue
|
|
out.append({"prompt": prompt, "chosen": chosen, "rejected": rejected})
|
|
return out
|
|
|
|
|
|
def collect(source: str, max_n: int, split: str) -> list[dict]:
|
|
rows: list[dict] = []
|
|
if source in ("hh", "both"):
|
|
print(f"Loading Anthropic/hh-rlhf [{split}] ...")
|
|
rows += from_hh(max_n, split)
|
|
if source in ("ultrafeedback", "both"):
|
|
print(f"Loading ultrafeedback_binarized [{split}] ...")
|
|
try:
|
|
rows += from_ultrafeedback(max_n, split)
|
|
except Exception as e: # noqa: BLE001
|
|
print(f" (skipping ultrafeedback: {e})")
|
|
return rows
|
|
|
|
|
|
def write_jsonl(rows: list[dict], path: str) -> None:
|
|
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
|
|
with open(path, "w") as f:
|
|
for r in rows:
|
|
f.write(json.dumps(r) + "\n")
|
|
print(f" wrote {len(rows)} pairs -> {path}")
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--source", choices=["hh", "ultrafeedback", "both"], default="both")
|
|
p.add_argument("--max_per_source", type=int, default=40000)
|
|
p.add_argument("--out_dir", default="/ephemeral/data")
|
|
args = p.parse_args()
|
|
|
|
train = collect(args.source, args.max_per_source, "train")
|
|
test = collect(args.source, max(2000, args.max_per_source // 20), "test")
|
|
import random
|
|
random.Random(0).shuffle(train)
|
|
write_jsonl(train, os.path.join(args.out_dir, "preferences.jsonl"))
|
|
write_jsonl(test, os.path.join(args.out_dir, "preferences_test.jsonl"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|