Compute accuracy over completed AIME results.
(
runs_dir: str,
eval_dir: str = "evals",
)
| 106 | |
| 107 | |
| 108 | def evaluate( |
| 109 | runs_dir: str, |
| 110 | eval_dir: str = "evals", |
| 111 | ) -> dict: |
| 112 | """Compute accuracy over completed AIME results.""" |
| 113 | runs_path = Path(runs_dir) |
| 114 | output_dir = Path(eval_dir) |
| 115 | output_dir.mkdir(parents=True, exist_ok=True) |
| 116 | |
| 117 | json_files = sorted(runs_path.glob("*.json")) |
| 118 | if not json_files: |
| 119 | print(f"No JSON files in {runs_path}") |
| 120 | return {} |
| 121 | |
| 122 | results: List[dict] = [] |
| 123 | for json_path in json_files: |
| 124 | with json_path.open("r") as f: |
| 125 | results.append(json.load(f)) |
| 126 | |
| 127 | completed = [r for r in results if r.get("status") == "completed"] |
| 128 | total = len(completed) |
| 129 | |
| 130 | if total == 0: |
| 131 | print("No completed problems to evaluate.") |
| 132 | return {} |
| 133 | |
| 134 | correct = sum(1 for r in completed if r.get("is_correct", False)) |
| 135 | accuracy = round(correct / total, 4) |
| 136 | |
| 137 | # Refusal rate (could not parse an answer) |
| 138 | refused = sum(1 for r in completed if r.get("parsed_answer") is None) |
| 139 | |
| 140 | # Per-dataset breakdown |
| 141 | dataset_stats = defaultdict(lambda: {"total": 0, "correct": 0}) |
| 142 | for r in completed: |
| 143 | ds = r.get("dataset", "unknown") |
| 144 | dataset_stats[ds]["total"] += 1 |
| 145 | if r.get("is_correct", False): |
| 146 | dataset_stats[ds]["correct"] += 1 |
| 147 | |
| 148 | per_dataset = {} |
| 149 | for ds, stats in sorted(dataset_stats.items()): |
| 150 | ds_acc = ( |
| 151 | round(stats["correct"] / stats["total"], 4) if stats["total"] > 0 else 0.0 |
| 152 | ) |
| 153 | per_dataset[ds] = { |
| 154 | "total": stats["total"], |
| 155 | "correct": stats["correct"], |
| 156 | "accuracy": ds_acc, |
| 157 | } |
| 158 | |
| 159 | summary = { |
| 160 | "total_completed": total, |
| 161 | "total_failed": len(results) - total, |
| 162 | "correct": correct, |
| 163 | "accuracy": accuracy, |
| 164 | "refused": refused, |
| 165 | "refusal_rate": round(refused / total, 4) if total > 0 else 0.0, |