| """Self-aggregation module: builds an items file in which each problem is followed by the written-up final parts |
| (the text after </think>) of K independent attempts. The model then reviews the attempts and gives one answer. Ids, |
| tasks and gold answers stay the same, so the file runs through gen.py / posthoc_coc.py / grade.py unchanged. |
| |
| Also writes, per id, the weighted-vote answer of the same K attempts (field "vote"), so the aggregator can be compared |
| with the vote on exactly the same attempts, and used only where the vote is not unanimous ("unanimous"). |
| usage: agg_items.py --recs results/m180_fb4/mini.jsonl --items data/coc/math180.jsonl --out data/agg/m180.jsonl |
| """ |
| import argparse |
| import collections |
| import json |
| import os |
| import random |
| import re |
| import sys |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from grade import last_boxed |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--recs", required=True) |
| ap.add_argument("--items", nargs="+", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--k", type=int, default=4) |
| ap.add_argument("--max-chars", type=int, default=2000) |
| a = ap.parse_args() |
|
|
| items = {} |
| for p in a.items: |
| for l in open(p, encoding="utf-8"): |
| r = json.loads(l) |
| items[r["id"]] = r |
| by = collections.defaultdict(list) |
| for l in open(a.recs, encoding="utf-8"): |
| r = json.loads(l) |
| if r["id"] in items: |
| by[r["id"]].append(r) |
|
|
| MC = re.compile(r"answer is:?\s*\(?\**\s*([A-J])\b") |
|
|
|
|
| def key(task, fin): |
| if task == "mmlupro": |
| m = MC.findall(fin) |
| return m[-1] if m else None |
| b = last_boxed(fin) |
| return b.strip().replace(" ", "") if b else None |
|
|
|
|
| def clip(s): |
| s = s.strip() |
| if len(s) <= a.max_chars: |
| return s |
| return s[: a.max_chars // 3] + "\n[...]\n" + s[-(a.max_chars - a.max_chars // 3):] |
|
|
|
|
| TAIL = {"math": "Put your final answer within \\boxed{}.", |
| "mmlupro": "End your response with 'The answer is (X)' where X is the letter of the correct option."} |
|
|
| n = 0 |
| os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True) |
| with open(a.out, "w", encoding="utf-8") as fo: |
| for pid, g in by.items(): |
| it = items[pid] |
| g = sorted(g, key=lambda r: r["sample"])[: a.k] |
| fins = [r["completion"].split("</think>", 1)[1] if "</think>" in r["completion"] else "" for r in g] |
| votes = collections.defaultdict(float) |
| for r, f in zip(g, fins): |
| k_ = key(it["task"], f) |
| if k_: |
| votes[k_] += (0.25 if r.get("module") == "fallback" else 1.0) / max(r["n_tokens"], 200) ** 0.5 |
| keys = [key(it["task"], f) for f in fins] |
| order = list(range(len(g))) |
| random.Random(pid).shuffle(order) |
| att = "\n\n".join(f"### Attempt {i + 1}\n{clip(fins[j]) or '(no final write-up)'}" for i, j in enumerate(order)) |
| prompt = (f"{it['prompt']}\n\n---\n\n{len(g)} independent attempts at this problem are shown below. They can " |
| f"contain mistakes, and their final answers can disagree.\n\n{att}\n\n---\n\nReview the attempts " |
| "critically. Find the points where they disagree, check the disputed steps yourself, and decide " |
| "which answer is correct; if all of them are wrong, solve the problem yourself. Do not just count " |
| "votes: an answer given by fewer attempts can be the right one. " |
| + TAIL["mmlupro" if it["task"] == "mmlupro" else "math"]) |
| fo.write(json.dumps({**it, "prompt": prompt, "vote": max(votes, key=votes.get) if votes else None, |
| "unanimous": len(set(keys)) == 1 and keys[0] is not None}, ensure_ascii=False) + "\n") |
| n += 1 |
| print(n, "items written") |
|
|
