data/build_datasets.py

#!/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()