#!/usr/bin/env python3 """ Phase 1 of CODER_TRAINING_PLAN.md: turn the scraped duels_index.jsonl into training-ready datasets — no reconstruction needed, the labels are already real per-prompt judge outcomes (see scrape_duels.py / CODER_TRAINING_PLAN.md for how that was confirmed). Produces three JSONL files under --out-dir, in the row shape TRL's SFTTrainer / DPOTrainer / KTOTrainer expect for vision-language models (see train_sft.py / train_dpo.py / train_kto.py): sft.jsonl {"images": [path], "prompt": [...chat msgs...], "completion": [...msgs...]} preference.jsonl {"images": [path], "prompt": [...chat msgs...], "chosen": [...], "rejected": [...]} kto.jsonl {"images": [path], "prompt": [...chat msgs...], "completion": [...], "label": bool} A held-out split is carved out by ROUND, not by individual prompt — rounds share `carryover_prompts` with their neighbors (config.json), so a per-prompt random split would leak near-duplicate examples across train/eval. """ from __future__ import annotations import argparse import json import random from collections import defaultdict from pathlib import Path SYSTEM_PROMPT = ( "You are a Three.js code generator. Given a reference image of an object, " "output a single JavaScript module: `export default function generate(THREE) { ... }` " "that procedurally constructs a Three.js Object3D matching the reference. " "Respond with JavaScript code only." ) USER_TEXT = "Generate a Three.js module for the object in this reference image." def load_manifest(path: Path) -> list[dict]: rows = [] with open(path) as f: for line in f: line = line.strip() if line: rows.append(json.loads(line)) return rows def read_js(path: str | None) -> str | None: if not path: return None p = Path(path) if not p.exists() or p.stat().st_size == 0: return None try: return p.read_text(encoding="utf-8", errors="replace") except OSError: return None def chat_prompt(image_path: str) -> list[dict]: return [ {"role": "system", "content": SYSTEM_PROMPT}, { "role": "user", "content": [ {"type": "image"}, {"type": "text", "text": USER_TEXT}, ], }, ] def chat_completion(js_code: str) -> list[dict]: return [{"role": "assistant", "content": js_code}] def held_out_rounds(rows: list[dict], eval_fraction: float, seed: int) -> set[int]: rounds = sorted({r["round"] for r in rows}) rng = random.Random(seed) n_eval = max(1, round(len(rounds) * eval_fraction)) return set(rng.sample(rounds, n_eval)) def dedup_key(row: dict, side: str) -> str: """Same generation (round, hotkey, prompt hash) can appear as the winner across many duel records in a round (one leader gen faces every challenger) — dedup SFT examples on this, keep preference pairs as-is since each pair is still a distinct (chosen, rejected) comparison.""" hotkey = row["left"] if side == "left" else row["right"] return f"{row['round']}::{hotkey}::{row['name']}" def build(rows: list[dict], out_dir: Path, eval_fraction: float, seed: int, min_margin: float) -> None: eval_rounds = held_out_rounds(rows, eval_fraction, seed) print(f"Held-out (eval) rounds: {sorted(eval_rounds)}") sft_seen: set[str] = set() counts = defaultdict(lambda: defaultdict(int)) with ( open(out_dir / "sft.jsonl", "w") as sft_fh, open(out_dir / "preference.jsonl", "w") as pref_fh, open(out_dir / "kto.jsonl", "w") as kto_fh, ): for row in rows: winner = row["winner"] if winner not in ("left", "right"): continue # scrape_duels.py already filters draws, defensive re-check if abs(row.get("match_margin", 0.0) or 0.0) < min_margin: # CAUTION: match_margin is a whole-MATCH aggregate (net score # over all ~128 prompts in that round's leader-vs-challenger # pairing), not a per-DUEL confidence signal — this filter # can only drop or keep entire matches by their overall # lopsidedness, it cannot tell you whether THIS SPECIFIC duel # was a decisive win or a squeaker (that lives in the # per-duel `detail` trace, not here). CODER_TRAINING_PLAN.md's # open question — that near-threshold (informative) examples # may be underrepresented relative to blowout (easy, # low-information) ones — is about per-duel informativeness, # which this coarse match-level knob does not actually # operationalize; treat it as a blunt, opt-in experiment # (default 0.0 = off, keep everything), not a precise fix for # that open question. continue split = "eval" if row["round"] in eval_rounds else "train" chosen_js = read_js(row["left_js"] if winner == "left" else row["right_js"]) rejected_js = read_js(row["right_js"] if winner == "left" else row["left_js"]) if not chosen_js or not rejected_js: continue image_path = row["prompt_image"] prompt = chat_prompt(image_path) # --- SFT: winner only, deduped per unique generation --- key = dedup_key(row, winner) if key not in sft_seen: sft_seen.add(key) sft_fh.write(json.dumps({ "images": [image_path], "prompt": prompt, "completion": chat_completion(chosen_js), "split": split, "round": row["round"], }) + "\n") counts[split]["sft"] += 1 # --- Preference (DPO/SimPO/ORPO): full pair, every duel is distinct --- # "detail" (the judge's S1-S4 trace) has a DIFFERENT nested shape # depending on which stage decided the duel (decided_by: "S1" has # only an angles list; "S2 consensus" adds bv/ac/cl sub-judges; # "S3"/"S4" add yet other keys) — heterogeneous nested dicts # across rows are exactly what trips up datasets/PyArrow's # automatic JSON schema inference (error, or silent null-coercion # of mismatched fields). Serializing to a JSON string keeps the # column type uniform; json.loads() it back out if inspecting. pref_fh.write(json.dumps({ "images": [image_path], "prompt": prompt, "chosen": chat_completion(chosen_js), "rejected": chat_completion(rejected_js), "margin": row.get("match_margin"), "detail_json": json.dumps(row.get("detail")) if row.get("detail") is not None else None, "split": split, "round": row["round"], }) + "\n") counts[split]["preference_pairs"] += 1 # --- KTO: same pair, split into two independently-labeled examples --- for js_code, label in ((chosen_js, True), (rejected_js, False)): kto_fh.write(json.dumps({ "images": [image_path], "prompt": prompt, "completion": chat_completion(js_code), "label": label, "split": split, "round": row["round"], }) + "\n") counts[split]["kto_examples"] += 1 print("\nCounts:") for split in ("train", "eval"): c = counts[split] print(f" {split}: sft={c['sft']} preference_pairs={c['preference_pairs']} kto_examples={c['kto_examples']}") def main() -> None: p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument("--manifest", default="./archive/duels_index.jsonl") p.add_argument("--out-dir", default="./datasets") p.add_argument("--eval-fraction", type=float, default=0.15, help="Fraction of ROUNDS held out for eval") p.add_argument("--seed", type=int, default=42) p.add_argument( "--min-margin", type=float, default=0.0, help="Drop entire MATCHES (not individual duels) with |match_margin| below this " "(0.0 = keep everything). Coarse — see the comment above the filter in build() " "for why this isn't the same as filtering by per-duel informativeness.", ) args = p.parse_args() out_dir = Path(args.out_dir) out_dir.mkdir(parents=True, exist_ok=True) rows = load_manifest(Path(args.manifest)) print(f"Loaded {len(rows)} usable duel records from {args.manifest}") if not rows: print("Nothing to build — run scrape_duels.py first.") return build(rows, out_dir, args.eval_fraction, args.seed, args.min_margin) print(f"\nWrote sft.jsonl, preference.jsonl, kto.jsonl to {out_dir}") if __name__ == "__main__": main()