sweep_program_synth_evo_smartbatch.py

"""Sweep program_synth_evo_smartbatch.py across Evo 2 checkpoints.

This driver enforces per-trial randomized encodings for every run by passing
`--per-trial-random` to the smartbatch harness.
"""

from __future__ import annotations

import argparse
import json
import subprocess
from pathlib import Path
from typing import List, Sequence, Tuple


DEFAULT_SHOTS: Sequence[int] = (1, 2, 4, 8, 16, 32, 64, 128)
DEFAULT_MODELS: Tuple[str, ...] = ("evo2_1b_base", "evo2_7b", "evo2_40b")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Sweep program_synth_evo_smartbatch.py across Evo models with per-trial randomized encodings."
        )
    )
    parser.add_argument(
        "--models",
        type=str,
        nargs="*",
        default=list(DEFAULT_MODELS),
        help=(
            "Evo model identifiers to evaluate (default: evo2_1b_base evo2_7b). "
            "Each value is forwarded to --model."
        ),
    )
    parser.add_argument(
        "--programs",
        type=Path,
        default=Path("curated_transformations.jsonl"),
        help="Path to the JSONL file of programs to evaluate (default: curated_transformations.jsonl).",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("sweep_results_evo_smartbatch"),
        help="Directory to store per-run results and the summary JSON (default: sweep_results_evo_smartbatch).",
    )
    parser.add_argument(
        "--shots",
        type=int,
        nargs="*",
        default=list(DEFAULT_SHOTS),
        help="Explicit list of in-context example counts to sweep (default: 1 2 4 8 16 32 64 128).",
    )
    parser.add_argument(
        "--trials-per-program",
        type=int,
        default=8,
        help="Number of trials per program passed to program_synth_evo_smartbatch.py (default: 8).",
    )
    parser.add_argument(
        "--bit-length",
        type=int,
        default=8,
        help="Bit length argument forwarded to program_synth_evo_smartbatch.py (default: 8).",
    )
    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=None,
        help="Optional max_new_tokens forwarded to program_synth_evo_smartbatch.py (default: bit length).",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Seed value forwarded to program_synth_evo_smartbatch.py (default: 42).",
    )
    parser.add_argument(
        "--python-exe",
        default="python",
        help="Python executable used to invoke program_synth_evo_smartbatch.py (default: python).",
    )
    parser.add_argument(
        "--program-synth",
        type=Path,
        default=Path("program_synth_evo_smartbatch.py"),
        help="Path to the program_synth_evo_smartbatch.py script (default: program_synth_evo_smartbatch.py).",
    )
    return parser.parse_args()


def run_command(command: Sequence[str]) -> None:
    subprocess.run(command, check=True)


def slugify_model(identifier: str) -> str:
    return identifier.replace("/", "_")


def main() -> None:
    args = parse_args()

    backend = "evo"

    model_entries = [(model_id, slugify_model(model_id)) for model_id in args.models]
    if not model_entries:
        raise SystemExit("No models provided for sweeping.")

    output_dir = args.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)

    summary_runs: List[dict] = []

    for model_identifier, model_slug in model_entries:
        for shots in args.shots:
            output_file = output_dir / f"{model_slug}_ic{shots}.json"
            command = [
                args.python_exe,
                str(args.program_synth),
                "--model",
                model_identifier,
                "--programs",
                str(args.programs),
                "--output",
                str(output_file),
                "--in-context-examples",
                str(shots),
                "--trials-per-program",
                str(args.trials_per_program),
                "--bit-length",
                str(args.bit_length),
                "--seed",
                str(args.seed),
                "--backend",
                backend,
                "--per-trial-random",
            ]
            if args.max_new_tokens is not None:
                command.extend(["--max-new-tokens", str(args.max_new_tokens)])

            print("Running:", " ".join(command))
            run_command(command)

            with output_file.open("r", encoding="utf-8") as handle:
                result_data = json.load(handle)

            overall = result_data.get("overall", {})
            config = result_data.get("config", {})
            summary_runs.append(
                {
                    "model": model_identifier,
                    "model_slug": model_slug,
                    "backend": backend,
                    "in_context_examples": shots,
                    "trials_per_program": args.trials_per_program,
                    "bit_length": args.bit_length,
                    "max_new_tokens": config.get(
                        "max_new_tokens", args.max_new_tokens or args.bit_length
                    ),
                    "per_trial_random": config.get("per_trial_random", True),
                    "mean_accuracy": overall.get("mean_accuracy"),
                    "accuracy_stderr": overall.get("stderr"),
                    "mean_edit_distance": overall.get("mean_edit_distance"),
                    "edit_distance_stderr": overall.get("edit_distance_stderr"),
                    "output_file": str(output_file),
                }
            )

    summary = {
        "models": list(args.models),
        "programs_file": str(args.programs),
        "shots": list(args.shots),
        "trials_per_program": args.trials_per_program,
        "bit_length": args.bit_length,
        "max_new_tokens": args.max_new_tokens,
        "seed": args.seed,
        "backend": backend,
        "per_trial_random": True,
        "runs": summary_runs,
    }

    summary_path = output_dir / "sweep_summary.json"
    with summary_path.open("w", encoding="utf-8") as handle:
        json.dump(summary, handle, indent=2)

    print(f"Wrote summary to {summary_path}")


if __name__ == "__main__":
    main()