File size: 5,212 Bytes
de46078 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """Command line interface."""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
from .runner import format_report, get_spark, run_suite, write_report
from .schema import TaskValidationError, load_tasks
DEFAULT_TASKS = Path(__file__).resolve().parent.parent / "tasks"
def _add_selection_args(p: argparse.ArgumentParser) -> None:
p.add_argument("--tasks", type=Path, default=DEFAULT_TASKS, help="task directory")
p.add_argument("--category", action="append", dest="categories", help="filter (repeatable)")
p.add_argument("--id", action="append", dest="ids", help="run specific task ids")
def cmd_run(args: argparse.Namespace) -> int:
from .models import build_model # local import: keeps `validate` dependency-light
tasks = load_tasks(args.tasks, args.categories, args.ids)
if not tasks:
print("no tasks matched", file=sys.stderr)
return 1
model = build_model(args.model, timeout=args.request_timeout)
ks = tuple(sorted({1, *(args.k or [])}))
if args.n < max(ks):
print(
f"error: --n {args.n} is too small for pass@{max(ks)}; "
f"the estimator needs n >= k",
file=sys.stderr,
)
return 2
report = run_suite(
tasks,
model,
n_samples=args.n,
ks=ks,
temperature=args.temperature,
max_tokens=args.max_tokens,
timeout=args.timeout,
keep_responses=not args.no_responses,
)
print(format_report(report))
if args.out:
write_report(report, args.out)
print(f"wrote {args.out}")
return 0
def cmd_validate(args: argparse.Namespace) -> int:
"""Structural checks only. `selfcheck` is the one that executes anything."""
try:
tasks = load_tasks(args.tasks, args.categories, args.ids)
except TaskValidationError as exc:
print(f"INVALID: {exc}", file=sys.stderr)
return 1
from collections import Counter
counts = Counter(t.category for t in tasks)
print(f"{len(tasks)} tasks, all structurally valid\n")
for cat, n in sorted(counts.items()):
print(f" {cat:<22} {n:>4}")
missing = [t.id for t in tasks if not t.probes]
if missing:
print(f"\nwarning: {len(missing)} tasks have no 'probes' note: {missing[:5]}")
return 0
def cmd_selfcheck(args: argparse.Namespace) -> int:
"""Execute every reference solution.
This is the check that matters. If a gold solution does not run, every
model scored against that task gets a meaningless result.
"""
from .harness import evaluate_candidate
tasks = load_tasks(args.tasks, args.categories, args.ids)
spark = get_spark("spark-eval-selfcheck")
spark.sparkContext.setLogLevel("ERROR")
failures = []
for i, task in enumerate(tasks, 1):
result = evaluate_candidate(spark, task, task.solution, timeout=args.timeout)
status = "ok" if result.ok else f"BROKEN ({result.status})"
print(f"[{i:>3}/{len(tasks)}] {task.id:<40} {status}", flush=True)
if not result.ok:
failures.append((task.id, result.detail))
spark.stop()
if failures:
print(f"\n{len(failures)} reference solution(s) failed:\n", file=sys.stderr)
for tid, detail in failures:
print(f" {tid}: {detail}", file=sys.stderr)
return 1
print(f"\nall {len(tasks)} reference solutions execute and self-compare")
return 0
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(
prog="spark-eval",
description="Execution-based benchmark for PySpark code generation.",
)
sub = parser.add_subparsers(dest="command", required=True)
p_run = sub.add_parser("run", help="score a model against the suite")
_add_selection_args(p_run)
p_run.add_argument(
"--model",
required=True,
help="ollama:<tag> | openai:<model> | dummy:reference",
)
p_run.add_argument("--n", type=int, default=1, help="samples per task")
p_run.add_argument(
"--k", type=int, action="append", help="report pass@k (repeatable, needs n>=k)"
)
p_run.add_argument("--temperature", type=float, default=0.2)
p_run.add_argument("--max-tokens", type=int, default=1024)
p_run.add_argument("--timeout", type=int, default=60, help="per-task exec seconds")
p_run.add_argument("--request-timeout", type=int, default=300)
p_run.add_argument("--out", type=Path, help="write full JSON report here")
p_run.add_argument(
"--no-responses",
action="store_true",
help="omit raw generations from the report (smaller files)",
)
p_run.set_defaults(func=cmd_run)
p_val = sub.add_parser("validate", help="structural check on task files")
_add_selection_args(p_val)
p_val.set_defaults(func=cmd_validate)
p_self = sub.add_parser("selfcheck", help="execute every reference solution")
_add_selection_args(p_self)
p_self.add_argument("--timeout", type=int, default=60)
p_self.set_defaults(func=cmd_selfcheck)
args = parser.parse_args(argv)
return args.func(args)
if __name__ == "__main__":
raise SystemExit(main())
|