Eval cells now record turns/tool_calls/completion+prompt tokens; summary gets efficiency-vs-base ratios

eval_adapter.py

@@ -65,15 +65,25 @@ def run_one(task_index, model):
             kwargs.pop("eval_sampling")
             result = env.run_rollout(**kwargs)
     except Exception as exc:
-        return task_index, model, None, True, f"{type(exc).__name__}: {exc}"
+        return task_index, model, None, True, f"{type(exc).__name__}: {exc}", None
     finally:
         try:
             env.close()
         except Exception:
             pass
     if not result.ok or result.reward is None:
-        return task_index, model, None, True, getattr(result, "error", None) or "not ok"
-    return task_index, model, float(result.reward) == 1.0, False, None
+        return task_index, model, None, True, getattr(result, "error", None) or "not ok", None
+    turns = getattr(result, "turns", []) or []
+    metrics = {
+        "turns": getattr(result, "n_turns", 0) or len(turns),
+        "tool_calls": sum(getattr(t, "n_tools", 0) for t in turns),
+        "completion_tokens": sum(len(getattr(t, "completion_token_ids", []) or []) for t in turns),
+        # Each turn resends the whole conversation, so this grows ~quadratically
+        # with turn count — the quantity the article's input-token savings cut.
+        "prompt_tokens_total": sum(len(getattr(t, "prompt_token_ids", []) or []) for t in turns),
+        "wall_s": getattr(result, "wall_s", 0.0),
+    }
+    return task_index, model, float(result.reward) == 1.0, False, None, metrics
 
 
 def main():
@@ -87,8 +97,8 @@ def main():
         futures = [pool.submit(run_one, i, m) for i, m in work]
         done = 0
         for fut in as_completed(futures):
-            i, m, correct, failed, err = fut.result()
-            cells[(i, m)] = {"correct": correct, "failed": failed, "error": err}
+            i, m, correct, failed, err, metrics = fut.result()
+            cells[(i, m)] = {"correct": correct, "failed": failed, "error": err, "metrics": metrics}
             done += 1
             if done % 20 == 0 or done == len(work):
                 print(f"[eval] {done}/{len(work)} cells done", flush=True)

eval_watch.py

@@ -62,23 +62,42 @@ def lora_rpc(route, name, path=None):
     urllib.request.urlopen(req, timeout=900)
 
 
-def pass_at_1(cells_path):
+def agg(cells_path):
+    """pass@1 plus efficiency means over solved tasks (the article's second axis)."""
     cells = json.loads(Path(cells_path).read_text())["cells"]
     scored = [c["correct"] for c in cells.values() if not c["failed"]]
-    return (sum(scored) / len(scored)) if scored else None, len(scored)
+    solved = [c["metrics"] for c in cells.values() if c["correct"] and c.get("metrics")]
+
+    def mean(key):
+        vals = [m[key] for m in solved if key in m]
+        return round(sum(vals) / len(vals), 1) if vals else None
+
+    return {
+        "pass_at_1": (sum(scored) / len(scored)) if scored else None,
+        "scorable": len(scored),
+        "solved": len(solved),
+        "tool_calls_solved_mean": mean("tool_calls"),
+        "turns_solved_mean": mean("turns"),
+        "completion_tokens_solved_mean": mean("completion_tokens"),
+        "prompt_tokens_solved_mean": mean("prompt_tokens_total"),
+    }
 
 
 def main():
     fs = repo_files()
     # 1. Base-model cells, once. Every checkpoint compares against these.
-    if "eval/base_cells.json" not in fs:
+    #    Recomputed if the cached file predates per-cell metrics.
+    if "eval/base_cells.json" in fs:
+        local = api.hf_hub_download(REPO, "eval/base_cells.json")
+        Path("/tmp/base_cells.json").write_text(Path(local).read_text())
+    if "eval/base_cells.json" not in fs or any(
+        c.get("metrics") is None
+        for c in json.loads(Path("/tmp/base_cells.json").read_text())["cells"].values()
+        if c["correct"]
+    ):
         print("[watch] evaluating base model first", flush=True)
         run_eval(BASE_MODEL, "/tmp/base_cells.json")
         upload("/tmp/base_cells.json", "eval/base_cells.json")
-    elif not Path("/tmp/base_cells.json").exists():
-        # A previous worker already computed base cells: reuse them locally.
-        local = api.hf_hub_download(REPO, "eval/base_cells.json")
-        Path("/tmp/base_cells.json").write_text(Path(local).read_text())
 
     summary = {}
     if "eval/summary.json" in fs:
@@ -105,14 +124,21 @@ def main():
                 print(f"[watch] {ckpt} eval failed ({exc}); will retry next poll", flush=True)
                 continue
             upload(f"/tmp/{ckpt}_cells.json", f"eval/{ckpt}_results.json")
-            p1, n = pass_at_1(f"/tmp/{ckpt}_cells.json")
-            pb, nb = pass_at_1("/tmp/base_cells.json")
-            summary[ckpt] = {"pass_at_1": p1, "scorable": n,
-                             "base_pass_at_1": pb,
-                             "delta": (p1 - pb) if p1 is not None and pb is not None else None}
+            a, b = agg(f"/tmp/{ckpt}_cells.json"), agg("/tmp/base_cells.json")
+            entry = dict(a)
+            entry["base_pass_at_1"] = b["pass_at_1"]
+            if a["pass_at_1"] is not None and b["pass_at_1"] is not None:
+                entry["delta"] = a["pass_at_1"] - b["pass_at_1"]
+            for key in ("tool_calls_solved_mean", "turns_solved_mean",
+                        "completion_tokens_solved_mean", "prompt_tokens_solved_mean"):
+                if a.get(key) is not None and b.get(key):
+                    entry[key.replace("_mean", "_vs_base")] = round(a[key] / b[key], 3)
+            summary[ckpt] = entry
             Path("/tmp/summary.json").write_text(json.dumps(summary, indent=2))
             upload("/tmp/summary.json", "eval/summary.json")
-            print(f"[watch] {ckpt}: pass@1={p1} (base {pb})", flush=True)
+            print(f"[watch] {ckpt}: pass@1={a['pass_at_1']} (base {b['pass_at_1']}), "
+                  f"tool_calls={a['tool_calls_solved_mean']} (base {b['tool_calls_solved_mean']})",
+                  flush=True)
         time.sleep(POLL_S)