scripts/build_clean_training_set.py

#!/usr/bin/env python3
"""Combine strong + repaired canonical rows into clean training set."""

from __future__ import annotations

import json
import sys
from collections import Counter
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from dataset_quality_utils import asks_for_code, asks_single_file_html, has_code, normalize_prompt, score_row

ROOT = Path(__file__).resolve().parent.parent
STRONG_PATH = ROOT / "data" / "audit" / "strong_rows.jsonl"
REPAIRED_PATH = ROOT / "data" / "repaired" / "repaired_rows.jsonl"
OUTPUT_PATH = ROOT / "data" / "train" / "mythos_coder_clean_canonical.jsonl"


def load_rows(path: Path) -> list[dict]:
    rows = []
    if not path.exists():
        return rows
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            entry = json.loads(line)
            rows.append(entry.get("row", entry))
    return rows


def dedupe_best_score(rows: list[dict]) -> list[dict]:
    best: dict[str, tuple[int, dict]] = {}
    for row in rows:
        key = normalize_prompt(row.get("user_prompt", ""))
        score, _ = score_row(row)
        if key not in best or score > best[key][0]:
            best[key] = (score, row)
    return [v[1] for v in best.values()]


def main() -> int:
    OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    combined = load_rows(STRONG_PATH) + load_rows(REPAIRED_PATH)
    clean = dedupe_best_score(combined)

    with OUTPUT_PATH.open("w", encoding="utf-8") as f:
        for row in clean:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    task_types = Counter(row.get("task_type", "?") for row in clean)
    difficulties = Counter(row.get("difficulty", "?") for row in clean)
    code_rows = sum(1 for r in clean if has_code(r.get("solution", "")))
    code_ask = sum(1 for r in clean if asks_for_code(r.get("user_prompt", "")))
    single_html = sum(1 for r in clean if asks_single_file_html(r.get("user_prompt", "")))

    print(f"Strong source rows: {len(load_rows(STRONG_PATH))}")
    print(f"Repaired source rows: {len(load_rows(REPAIRED_PATH))}")
    print(f"Clean canonical rows: {len(clean)}")
    print(f"Code in solution: {code_rows}")
    print(f"Code-request prompts: {code_ask}")
    print(f"Single-file HTML prompts: {single_html}")
    print("Task types:", dict(task_types))
    print("Difficulties:", dict(difficulties))
    print(f"Wrote: {OUTPUT_PATH}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())