Add minimal GAP reproduction package

1 files changed, 168 insertions, 0 deletions

diff --git a/src/gap_pipeline/cli.py b/src/gap_pipeline/cli.py
new file mode 100644
index 0000000..105cc97
--- /dev/null
+++ b/src/gap_pipeline/cli.py
@@ -0,0 +1,168 @@
+"""Command-line interface for generation and all offline checks."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+from pathlib import Path
+from typing import Any
+
+from .clients import OpenAIJsonClient
+from .models import CanonicalItem
+from .offline import (
+    align_run_to_release,
+    load_dataset,
+    summarize_run,
+    validate_public_dataset,
+)
+from .pipeline import KernelPipeline, PipelineConfig
+from .release import export_release
+from .store import RunStore
+from .surface import SurfacePipeline
+
+
+def write_or_print(payload: dict[str, Any], output: Path | None) -> None:
+    text = json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
+    if output is None:
+        print(text, end="")
+    else:
+        output.parent.mkdir(parents=True, exist_ok=True)
+        output.write_text(text, encoding="utf-8")
+
+
+async def generate_kernel(args: argparse.Namespace) -> None:
+    records = load_dataset(args.dataset)
+    if args.item_id not in records:
+        raise SystemExit(f"item ID {args.item_id!r} not found in {args.dataset}")
+    item = CanonicalItem.from_public_record(records[args.item_id])
+    config = PipelineConfig(
+        proposer_model=args.proposer_model,
+        judge_model=args.judge_model,
+    )
+    proposer = OpenAIJsonClient(args.proposer_model)
+    judges = [
+        OpenAIJsonClient(
+            args.judge_model,
+            seed=None if args.seed is None else args.seed + judge_id,
+        )
+        for judge_id in range(1, 6)
+    ]
+    pipeline = KernelPipeline(
+        proposer=proposer,
+        judges=judges,
+        store=RunStore(args.run_dir, item.item_id),
+        config=config,
+    )
+    result = await pipeline.run(item)
+    print(result.model_dump_json(indent=2))
+
+
+async def generate_surfaces(args: argparse.Namespace) -> None:
+    records = load_dataset(args.dataset)
+    if args.item_id not in records:
+        raise SystemExit(f"item ID {args.item_id!r} not found in {args.dataset}")
+    item = CanonicalItem.from_public_record(records[args.item_id])
+    store = RunStore(args.run_dir, item.item_id)
+    store.write_input(item)
+    store.write_config(
+        {
+            "protocol_name": "gap-surface-v1",
+            "proposer_model": args.proposer_model,
+        }
+    )
+    pipeline = SurfacePipeline(OpenAIJsonClient(args.proposer_model), store)
+    variants = await pipeline.run_all(item)
+    print(
+        json.dumps(
+            {
+                family: variant.as_release_payload()
+                for family, variant in variants.items()
+            },
+            ensure_ascii=False,
+            indent=2,
+        )
+    )
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="gap-reproduce")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    validate = subparsers.add_parser("validate-data")
+    validate.add_argument("dataset", type=Path)
+    validate.add_argument("--output", type=Path)
+
+    summarize = subparsers.add_parser("summarize-run")
+    summarize.add_argument("run_dir", type=Path)
+    summarize.add_argument("--output", type=Path)
+
+    align = subparsers.add_parser("align-release")
+    align.add_argument("run_dir", type=Path)
+    align.add_argument("dataset", type=Path)
+    align.add_argument("--output", type=Path)
+
+    generate = subparsers.add_parser("generate-kernel")
+    generate.add_argument("--dataset", type=Path, required=True)
+    generate.add_argument("--item-id", required=True)
+    generate.add_argument("--run-dir", type=Path, required=True)
+    generate.add_argument("--proposer-model", default="o3")
+    generate.add_argument("--judge-model", default="o3")
+    generate.add_argument(
+        "--seed",
+        type=int,
+        help="optional provider seed; omitted by default for o3 compatibility",
+    )
+
+    surfaces = subparsers.add_parser("generate-surfaces")
+    surfaces.add_argument("--dataset", type=Path, required=True)
+    surfaces.add_argument("--item-id", required=True)
+    surfaces.add_argument("--run-dir", type=Path, required=True)
+    surfaces.add_argument("--proposer-model", default="o3")
+
+    export = subparsers.add_parser("export-release")
+    export.add_argument("--source-dataset", type=Path, required=True)
+    export.add_argument("--surface-run-dir", type=Path, required=True)
+    export.add_argument("--kernel-run-dir", type=Path, required=True)
+    export.add_argument("--output-root", type=Path, required=True)
+    export.add_argument("--allow-partial", action="store_true")
+    export.add_argument(
+        "--item-id",
+        action="append",
+        dest="item_ids",
+        help="export only this item ID; repeat to select multiple items",
+    )
+    return parser
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    if args.command == "validate-data":
+        write_or_print(validate_public_dataset(args.dataset), args.output)
+    elif args.command == "summarize-run":
+        write_or_print(summarize_run(args.run_dir), args.output)
+    elif args.command == "align-release":
+        write_or_print(
+            align_run_to_release(args.run_dir, args.dataset),
+            args.output,
+        )
+    elif args.command == "generate-kernel":
+        asyncio.run(generate_kernel(args))
+    elif args.command == "generate-surfaces":
+        asyncio.run(generate_surfaces(args))
+    elif args.command == "export-release":
+        manifest = export_release(
+            source_dataset=args.source_dataset,
+            surface_run_root=args.surface_run_dir,
+            kernel_run_root=args.kernel_run_dir,
+            output_root=args.output_root,
+            allow_partial=args.allow_partial,
+            item_ids=set(args.item_ids) if args.item_ids else None,
+        )
+        print(json.dumps(manifest, ensure_ascii=False, indent=2))
+    else:
+        raise AssertionError(args.command)
+
+
+if __name__ == "__main__":
+    main()