Something went wrong. Try again.
zig langref cli nate.tngl.io/zigman
Something went wrong. Try again.
3.2 kB · 92 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293"""the ablation loop: run each task bare vs. with-zigman, summarize the lift."""
from __future__ import annotations
import jsonfrom dataclasses import dataclassfrom pathlib import Pathfrom typing import Callable
from .config import Configfrom .pi import ArmResult, run_taskfrom .tasks import TASKS, Task
@dataclassclass Report: model: str arms: list[str] tasks: list[Task] results: dict[str, dict[str, ArmResult]] # task.id -> arm -> result
def total(self, arm: str) -> int: return sum(self.results[t.id][arm].oracle_ok for t in self.tasks)
def lift(self) -> int | None: if "bare" in self.arms and "zigman" in self.arms: return self.total("zigman") - self.total("bare") return None
def zigman_use_count(self) -> int: if "zigman" not in self.arms: return 0 return sum(self.results[t.id]["zigman"].used_zigman for t in self.tasks)
def ablate( cfg: Config, task_filter: str = "", arms: tuple[str, ...] = ("bare", "zigman"), on_result: Callable[[ArmResult], None] | None = None,) -> Report: tasks = [t for t in TASKS if not task_filter or task_filter in t.id] results: dict[str, dict[str, ArmResult]] = {} for t in tasks: results[t.id] = {} for arm in arms: r = run_task(cfg, t, with_zigman=(arm == "zigman")) results[t.id][arm] = r if on_result: on_result(r) return Report(model=cfg.model, arms=list(arms), tasks=tasks, results=results)
def render(report: Report) -> str: lines = [f"\nagent: Pi model: {report.model}"] header = "task".ljust(16) + "".join(a.ljust(8) for a in report.arms) lines.append(header) for t in report.tasks: row = t.id.ljust(16) for a in report.arms: row += ("pass" if report.results[t.id][a].oracle_ok else "·").ljust(8) lines.append(row) lines.append("-" * len(header)) lines.append("total".ljust(16) + "".join(f"{report.total(a)}/{len(report.tasks)}".ljust(8) for a in report.arms)) if "zigman" in report.arms: lines.append(f"zigman used in {report.zigman_use_count()}/{len(report.tasks)} zigman-arm runs") if (lift := report.lift()) is not None: lines.append(f"zigman lift: {lift:+d}") return "\n".join(lines)
def result_line(r: ArmResult) -> str: to = " TIMEOUT" if r.timed_out else "" uz = (" zigman:" + (str(r.zigman_queries) if r.used_zigman else "none")) if r.arm == "zigman" else "" return (f" {r.arm:<7} {r.task:<14} {'PASS' if r.oracle_ok else 'fail'} " f"[{','.join(r.tools) or 'no-tools'}{to}]{uz}")
def save(report: Report, out_dir: Path) -> Path: out_dir.mkdir(parents=True, exist_ok=True) payload = { "model": report.model, "arms": report.arms, "lift": report.lift(), "results": { tid: {arm: vars(r) | {"used_zigman": r.used_zigman} for arm, r in by_arm.items()} for tid, by_arm in report.results.items() }, } out = out_dir / f"pi_{Path(report.model).name}.json" out.write_text(json.dumps(payload, indent=2)) return out