dataset/scripts/merge_codex_runs.py

#!/usr/bin/env python3
"""Merge isolated per-paper Codex JSON outputs into dataset JSONL files."""

from __future__ import annotations

import json
import re
import subprocess
import sys
import tempfile
from pathlib import Path


ROOT = Path(__file__).resolve().parents[2]
PROCESSED = ROOT / "dataset" / "processed"
RUN_DIR = PROCESSED / "codex_runs"
LOG_DIR = PROCESSED / "logs"
VALIDATE_ONE = ROOT / "dataset" / "scripts" / "validate_one_codex_run.py"
EVIDENCE_TRACE_HEADING_RE = re.compile(r"(?i)^\s*(evidence trace|evidence|citations?|sources?|references?)\s*:?\s*$")
INLINE_EVIDENCE_TRACE_RE = re.compile(r"(?is)\n\s*(evidence trace|citations?|sources?|references?)\s*:.*$")


def iter_json(path: Path):
    with path.open() as f:
        return json.load(f)


def write_jsonl(path: Path, rows: list[dict]) -> None:
    with path.open("w") as f:
        for row in rows:
            f.write(json.dumps(row, sort_keys=True, separators=(",", ":")) + "\n")


def strip_evidence_trace(text: str) -> str:
    """Remove legacy evidence-trace sections from trainable text only."""
    text = INLINE_EVIDENCE_TRACE_RE.sub("", text)
    kept: list[str] = []
    for line in text.splitlines():
        if EVIDENCE_TRACE_HEADING_RE.match(line):
            break
        kept.append(line)
    return "\n".join(kept).rstrip()


def sanitize_trainable_text(obj: dict) -> dict:
    """Keep evidence metadata, but remove evidence traces from SFT/RL text."""
    for ex in obj.get("sft_examples") or []:
        for message in ex.get("messages") or []:
            if isinstance(message, dict) and message.get("role") in {"user", "assistant"}:
                message["content"] = strip_evidence_trace(str(message.get("content", "")))
    for ex in obj.get("rl_examples") or []:
        for field in ["prompt", "chosen_answer", "rejected_answer"]:
            if field in ex:
                ex[field] = strip_evidence_trace(str(ex.get(field, "")))
    return obj


def validate_sanitized_run(obj: dict, original_path: Path, schema_path: Path) -> None:
    with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as tmp:
        tmp_path = Path(tmp.name)
        json.dump(obj, tmp)
        tmp.write("\n")
    try:
        result = subprocess.run(
            [sys.executable, str(VALIDATE_ONE), str(tmp_path), str(schema_path)],
            check=False,
            text=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
    finally:
        tmp_path.unlink(missing_ok=True)

    if result.returncode != 0:
        if result.stdout:
            print(result.stdout, end="")
        if result.stderr:
            print(result.stderr, end="", file=sys.stderr)
        raise SystemExit(f"Refusing to merge invalid run after trainable-text sanitization: {original_path}")


def main() -> int:
    paper_index: list[dict] = []
    sft_examples: list[dict] = []
    rl_examples: list[dict] = []
    skipped: list[dict] = []

    if not RUN_DIR.exists():
        raise SystemExit(f"Missing run directory: {RUN_DIR}")

    for path in sorted(RUN_DIR.glob("*.json")):
        schema_path = LOG_DIR / f"{path.stem}.schema.json"
        if not schema_path.exists():
            print(f"Warning: skipping {path}; missing rendered schema {schema_path}", file=sys.stderr)
            continue
        obj = sanitize_trainable_text(iter_json(path))
        validate_sanitized_run(obj, path, schema_path)
        index = obj.get("paper_index")
        if isinstance(index, dict):
            paper_index.append(index)

        skipped_paper = obj.get("skipped_paper")
        if skipped_paper is not None:
            if isinstance(skipped_paper, dict):
                skipped.append(skipped_paper)
            continue

        sft_examples.extend(obj.get("sft_examples") or [])
        rl_examples.extend(obj.get("rl_examples") or [])

    write_jsonl(PROCESSED / "paper_index.jsonl", paper_index)
    write_jsonl(PROCESSED / "sft.jsonl", sft_examples)
    write_jsonl(PROCESSED / "rl.jsonl", rl_examples)
    write_jsonl(PROCESSED / "skipped_papers.jsonl", skipped)

    print(f"Merged {len(paper_index)} paper_index rows")
    print(f"Merged {len(sft_examples)} SFT examples")
    print(f"Merged {len(rl_examples)} RL examples")
    print(f"Merged {len(skipped)} skipped-paper rows")
    return 0


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