scripts/audit_sft_templates.py

#!/usr/bin/env python3
"""
audit_sft_templates.py

Scan canonical train rows for repetitive templates, fake verification phrases,
and oversized solution blocks. Writes data/eval/sft_template_audit.md.
"""

from __future__ import annotations

import json
import re
from collections import Counter
from pathlib import Path


ROOT = Path(__file__).resolve().parent.parent
TRAIN_PATH = ROOT / "datasets" / "mythos_coder_train.jsonl"
SFT_PATH = ROOT / "data" / "train" / "mythos_sft_messages.jsonl"
REPORT_PATH = ROOT / "data" / "eval" / "sft_template_audit.md"

FAKE_VERIFY_PATTERNS = [
    r"screenshots?\s+or\s+recordings?",
    r"before-and-after\s+notes",
    r"devtools\s+performance\s+panels",
    r"Chrome,\s*Firefox,\s*and\s*Safari",
    r"Verify with screenshots",
]

REPETITIVE_SOLUTION_PREFIXES = [
    "Scan ",
    "Read the active game loop",
    "Reproduce the reported",
    "Inspect state holders",
    "Implement the smallest change",
]

NUMBERED_ITEM_RE = re.compile(r"\d+\)\s*")


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


def solution_item_count(text: str) -> int:
    return len([part for part in NUMBERED_ITEM_RE.split(str(text)) if part.strip()])


def classify_source(row: dict) -> str:
    row_id = str(row.get("id", ""))
    source_repo = str(row.get("source_repo", ""))
    if source_repo or "-game-" in row_id or any(
        slug in row_id
        for slug in (
            "cloud9c_taro",
            "elkwizard_hengine",
            "labystudio_js_minecraft",
            "icurtis_third_person",
            "phaser",
            "kaboom",
            "excalibur",
            "melonjs",
            "littlejs",
            "kaplay",
        )
    ):
        return "game_repo_batch"
    if "html5up" in row_id:
        return "html5up"
    if "bedim-restaurant" in row_id:
        return "bedim_restaurant"
    if "bedim-portfolio" in row_id or "bedimcode-portfolio" in row_id:
        return "bedim_portfolio"
    if "50projects" in row_id:
        return "50projects50days"
    return "other"


def audit_rows(rows: list[dict]) -> dict:
    by_source: Counter[str] = Counter()
    numbered_solutions = 0
    long_solutions = 0
    fake_verify_hits = 0
    solution_prefix_hits: Counter[str] = Counter()
    duplicate_failure_log: Counter[str] = Counter()
    duplicate_lesson: Counter[str] = Counter()
    user_prompt_lengths: list[int] = []
    solution_lengths: list[int] = []

    for row in rows:
        source = classify_source(row)
        by_source[source] += 1

        solution = str(row.get("solution", ""))
        verification = str(row.get("verification", ""))
        failure_log = str(row.get("failure_log", ""))
        lesson = str(row.get("lesson", ""))
        user_prompt = str(row.get("user_prompt", ""))

        user_prompt_lengths.append(len(user_prompt))
        solution_lengths.append(len(solution))

        if NUMBERED_ITEM_RE.search(solution):
            numbered_solutions += 1
            if solution_item_count(solution) >= 5:
                long_solutions += 1

        for prefix in REPETITIVE_SOLUTION_PREFIXES:
            if solution.startswith(prefix) or f"1) {prefix}" in solution:
                solution_prefix_hits[prefix] += 1

        for pattern in FAKE_VERIFY_PATTERNS:
            if re.search(pattern, verification, re.IGNORECASE):
                fake_verify_hits += 1
                break

        if failure_log.strip():
            duplicate_failure_log[failure_log.strip()[:120]] += 1
        if lesson.strip():
            duplicate_lesson[lesson.strip()[:100]] += 1

    return {
        "total": len(rows),
        "by_source": by_source,
        "numbered_solutions": numbered_solutions,
        "long_solutions": long_solutions,
        "fake_verify_hits": fake_verify_hits,
        "solution_prefix_hits": solution_prefix_hits,
        "duplicate_failure_log": duplicate_failure_log,
        "duplicate_lesson": duplicate_lesson,
        "avg_user_prompt_len": sum(user_prompt_lengths) / max(len(user_prompt_lengths), 1),
        "avg_solution_len": sum(solution_lengths) / max(len(solution_lengths), 1),
    }


def audit_sft_lengths(path: Path) -> dict | None:
    if not path.exists():
        return None
    lengths = []
    with open(path, "r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            assistant = row["messages"][2]["content"]
            lengths.append(len(assistant))
    if not lengths:
        return None
    return {
        "count": len(lengths),
        "avg": sum(lengths) / len(lengths),
        "max": max(lengths),
        "over_1800": sum(1 for n in lengths if n > 1800),
    }


def top_duplicates(counter: Counter[str], min_count: int = 5, limit: int = 8) -> list[tuple[str, int]]:
    return counter.most_common(limit)


def write_report(train_stats: dict, sft_stats: dict | None) -> None:
    lines = [
        "# SFT Template Audit",
        "",
        f"Train file: `{TRAIN_PATH}`",
        f"SFT file: `{SFT_PATH}`",
        "",
        "## Dataset coverage",
        "",
        f"- Total train rows: **{train_stats['total']}**",
        f"- Avg user prompt length: **{train_stats['avg_user_prompt_len']:.0f}** chars",
        f"- Avg raw solution length: **{train_stats['avg_solution_len']:.0f}** chars",
        "",
        "### Rows by source bucket",
        "",
    ]
    for source, count in train_stats["by_source"].most_common():
        lines.append(f"- `{source}`: {count}")

    lines.extend(
        [
            "",
            "## Template issues (raw train)",
            "",
            f"- Numbered-list solutions: **{train_stats['numbered_solutions']}** ({100 * train_stats['numbered_solutions'] / max(train_stats['total'], 1):.1f}%)",
            f"- Solutions with 5+ numbered steps: **{train_stats['long_solutions']}**",
            f"- Verification rows matching fake/browser boilerplate: **{train_stats['fake_verify_hits']}**",
            "",
            "### Repetitive solution openings (game batch pattern)",
            "",
        ]
    )
    for prefix, count in train_stats["solution_prefix_hits"].most_common():
        lines.append(f"- `{prefix}...`: {count}")

    lines.extend(["", "### Most repeated failure_log prefixes", ""])
    for text, count in top_duplicates(train_stats["duplicate_failure_log"]):
        if count < 3:
            continue
        lines.append(f"- ({count}x) {text}...")

    lines.extend(["", "### Most repeated lesson prefixes", ""])
    for text, count in top_duplicates(train_stats["duplicate_lesson"]):
        if count < 3:
            continue
        lines.append(f"- ({count}x) {text}...")

    lines.extend(
        [
            "",
            "## Build-time mitigations",
            "",
            "- `build_sft_messages.py` now reads from `datasets/mythos_coder_train.jsonl`.",
            "- Assistant responses are compressed: numbered solutions capped to 4 bullets, verification trimmed to 3 checks.",
            "- Diagnosis drops redundant `Initial problem:` prefix and limits investigation steps to 4.",
            "",
        ]
    )

    if sft_stats:
        lines.extend(
            [
                "## SFT output after compression",
                "",
                f"- SFT rows: **{sft_stats['count']}**",
                f"- Avg assistant message: **{sft_stats['avg']:.0f}** chars",
                f"- Max assistant message: **{sft_stats['max']}** chars",
                f"- Rows still over 1800 chars: **{sft_stats['over_1800']}**",
                "",
            ]
        )
    else:
        lines.append("SFT file not found — run `python scripts/build_sft_messages.py` first.\n")

    lines.extend(
        [
            "## Recommendations",
            "",
            "1. Regenerate game-repo raw rows with shorter `solution` prose instead of echoing investigation steps.",
            "2. Replace screenshot/recording verification text with concrete command or browser checks.",
            "3. Keep user prompts messy/vague in eval only; train prompts should stay specific.",
            "4. Retrain LoRA after SFT rebuild and re-run `test_lora_model.py`.",
            "",
        ]
    )

    REPORT_PATH.parent.mkdir(parents=True, exist_ok=True)
    REPORT_PATH.write_text("\n".join(lines), encoding="utf-8")


def main() -> int:
    rows = load_train_rows(TRAIN_PATH)
    train_stats = audit_rows(rows)
    sft_stats = audit_sft_lengths(SFT_PATH)
    write_report(train_stats, sft_stats)
    print(f"Wrote audit report: {REPORT_PATH}")
    print(f"Train rows audited: {train_stats['total']}")
    if sft_stats:
        print(f"SFT avg assistant chars: {sft_stats['avg']:.0f}")
    return 0


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