File size: 3,741 Bytes
67ebcef | 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 | """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 # noqa: E402
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")
|