pemf/scripts/entrypoint.py

"""
Headless entrypoint for running a single experiment or a sweep.

Usage:
  python scripts/entrypoint.py --task example1 --out_dir experiments
  python scripts/entrypoint.py --task example1 --out_dir experiments --use_wandb

By default Weights & Biases logging is disabled. Use --use_wandb to enable it.
"""
import argparse
import json
import os
import importlib

def main():
    parser = argparse.ArgumentParser(description="Run ARC-AGI experiment (headless).")
    parser.add_argument("--task", type=str, required=True, help="Task name or path to task JSON")
    parser.add_argument("--out_dir", type=str, default="experiments", help="Output directory")
    parser.add_argument("--use_wandb", action="store_true", help="Enable Weights & Biases logging (default: off)")
    parser.add_argument("--params", type=str, default=None, help="Optional JSON string of params")
    args = parser.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)

    # lazy imports to avoid heavy startup cost
    import itt_solver.experiment_driver as ed
    import itt_solver.solver_core as sc

    # load task: if args.task is a JSON file path, load it; otherwise expect a built-in name
    if os.path.exists(args.task):
        with open(args.task) as fh:
            task = json.load(fh)
    else:
        # minimal built-in example if user passed 'example1'
        # Corrected target from real ARC task 007bbfb7 (Kronecker self-similar)
        if args.task == "example1":
            task = {
              'name': 'example1',
              'input': [[0,7,7],[7,7,7],[0,7,7]],
              'target': [
                [0,0,0,0,7,7,0,7,7],
                [0,0,0,7,7,7,7,7,7],
                [0,0,0,0,7,7,0,7,7],
                [0,7,7,0,7,7,0,7,7],
                [7,7,7,7,7,7,7,7,7],
                [0,7,7,0,7,7,0,7,7],
                [0,0,0,0,7,7,0,7,7],
                [0,0,0,7,7,7,7,7,7],
                [0,0,0,0,7,7,0,7,7],
              ],
              'target_shape': (9,9)
            }
        else:
            raise SystemExit(f"Unknown task identifier: {args.task}")

    # parse params if provided
    params = {}
    if args.params:
        try:
            params = json.loads(args.params)
        except Exception:
            print("Warning: could not parse --params JSON; ignoring.")

    # build atomic library using default factory
    atomic_library = ed.default_atomic_factory(params, task)

    # run single experiment
    result = ed.run_single(task, atomic_library, params, out_dir=args.out_dir)

    # optionally run W&B logging externally (only if requested)
    if args.use_wandb:
        try:
            from itt_solver.wandb_runner import run_and_log_wandb
            run_and_log_wandb(task, atomic_library, params, out_dir=args.out_dir,
                              wandb_project=params.get('wandb_project','itt_solver'),
                              wandb_entity=None, resume="allow")
        except Exception as e:
            print("W&B logging failed or not configured:", e)

    print("Run finished. Result summary:")
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()