File size: 2,807 Bytes
ec24fb0 | 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 | """Cross-model pooling: majority vote over the thinking answers of several runs (each <predictions.zip>.ballots.jsonl,
written by inference/predict.py), merged into a base predictions.zip for the questions the ballot files cover.
Ballots are concatenated in the order given; unparsed answers abstain; ties go to the first ballot (the first run's greedy
answer) if it is among the winners, else to the lowest label (the same rule as inference/predict.py).
python scripts/pool_ballots.py --base results/submissions/final_candidate_base_think_gated_sc.zip \
--ballots results/submissions/pool_base_medium.zip.ballots.jsonl results/submissions/pool_rft2_medium.zip.ballots.jsonl \
--out results/submissions/final_candidate_B_plus_pool_medium.zip
"""
from __future__ import annotations
import argparse
import json
import zipfile
from collections import Counter
from pathlib import Path
def vote(b: list):
c = Counter(x for x in b if x)
if not c:
return None
top = max(c.values())
w = [a for a, n in c.items() if n == top]
return b[0] if b[0] in w else min(w)
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--base", required=True, type=Path)
ap.add_argument("--ballots", required=True, nargs="+", type=Path)
ap.add_argument("--out", required=True, type=Path)
args = ap.parse_args()
runs = []
for f in args.ballots:
rows = [json.loads(l) for l in f.read_text(encoding="utf-8").splitlines() if l.strip()]
runs.append({(r["level"], r["subject"], r["q_no"]): r["ballots"] for r in rows})
keys = set(runs[0])
assert all(set(r) == keys for r in runs), "ballot files cover different questions"
pooled = {k: vote([b for r in runs for b in r[k]]) for k in keys}
with zipfile.ZipFile(args.base) as z:
base = [json.loads(l) for l in z.read("predictions.jsonl").decode("utf-8").splitlines() if l.strip()]
changed = matched = 0
for r in base:
k = (r["level"], r["subject"], r["q_no"])
if k in pooled and pooled[k]:
matched += 1
changed += str(pooled[k]) != r["answer"]
r["answer"] = str(pooled[k])
assert matched == sum(1 for v in pooled.values() if v), "some pooled questions are missing from the base file"
args.out.parent.mkdir(parents=True, exist_ok=True)
with zipfile.ZipFile(args.out, "w", zipfile.ZIP_DEFLATED) as z:
z.writestr("predictions.jsonl", "".join(json.dumps(r, ensure_ascii=False) + "\n" for r in base))
n_b = Counter(len([b for r in runs for b in r[k]]) for k in keys)
print(f"{len(keys)} questions pooled over {len(runs)} runs (ballots per question: {dict(n_b)}); "
f"{changed} answers differ from {args.base.name} -> {args.out}")
if __name__ == "__main__":
main()
|