86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .client import LLMClient, LLMClientConfig
|
|
from .pipeline import run_evaluation
|
|
from .skeletons import generate_skeletons
|
|
|
|
|
|
def _load_task(path: Path) -> dict[str, Any]:
|
|
payload = json.loads(path.read_text(encoding="utf-8"))
|
|
required = {"count", "length", "seed", "output_dir"}
|
|
missing = required - payload.keys()
|
|
if missing:
|
|
raise ValueError(f"task is missing fields: {', '.join(sorted(missing))}")
|
|
return payload
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(prog="logicbench")
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
|
generate = subparsers.add_parser(
|
|
"generate", help="generate hidden systems and answer-bearing behavior skeletons"
|
|
)
|
|
generate.add_argument("--count", type=int, default=3)
|
|
generate.add_argument("--length", default="short")
|
|
generate.add_argument("--seed", type=int, default=0)
|
|
generate.add_argument("--entity-count", type=int, default=128)
|
|
generate.add_argument("--simulation-steps", type=int, default=1000)
|
|
generate.add_argument("--output", type=Path)
|
|
|
|
run = subparsers.add_parser("run", help="run renderer and solver from a task JSON file")
|
|
run.add_argument("task", type=Path)
|
|
run.add_argument(
|
|
"--render-chunk-size",
|
|
type=int,
|
|
help="override the task's records per renderer request (default: 60)",
|
|
)
|
|
return parser
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
args = build_parser().parse_args(argv)
|
|
if args.command == "generate":
|
|
skeletons = generate_skeletons(
|
|
args.count,
|
|
args.length,
|
|
args.seed,
|
|
entity_count=args.entity_count,
|
|
simulation_steps=args.simulation_steps,
|
|
)
|
|
lines = [json.dumps(item.to_dict(), ensure_ascii=False) for item in skeletons]
|
|
text = "\n".join(lines) + "\n"
|
|
if args.output:
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(text, encoding="utf-8")
|
|
else:
|
|
print(text, end="")
|
|
return 0
|
|
|
|
task = _load_task(args.task)
|
|
renderer = LLMClient(LLMClientConfig.from_env("renderer"))
|
|
solver = LLMClient(LLMClientConfig.from_env("solver"))
|
|
render_chunk_size = (
|
|
args.render_chunk_size
|
|
if args.render_chunk_size is not None
|
|
else int(task.get("render_chunk_size", 60))
|
|
)
|
|
summary = run_evaluation(
|
|
renderer=renderer,
|
|
solver=solver,
|
|
count=int(task["count"]),
|
|
length=task["length"],
|
|
seed=int(task["seed"]),
|
|
output_dir=Path(task["output_dir"]),
|
|
entity_count=int(task.get("entity_count", 128)),
|
|
simulation_steps=int(task.get("simulation_steps", 1000)),
|
|
render_chunk_size=render_chunk_size,
|
|
)
|
|
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
|
return 0
|