fix: add option permutation to MMLU eval
- Few-shot examples now include subject preamble (consistent format) - Add --seed flag for option permutation (default 0, -1 to disable) - Shuffles A/B/C/D positions per-question to neutralise positional bias
This commit is contained in:
@@ -4,6 +4,7 @@ import argparse
|
|||||||
import csv
|
import csv
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
|
import random
|
||||||
from collections import defaultdict
|
from collections import defaultdict
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
@@ -154,19 +155,22 @@ def build_prompt(question: str, choices: dict, subject: str) -> str:
|
|||||||
|
|
||||||
|
|
||||||
def apply_chat(
|
def apply_chat(
|
||||||
tokenizer, raw_prompt: str, n_shot: int, dev_data: list[dict] | None
|
tokenizer,
|
||||||
|
raw_prompt: str,
|
||||||
|
n_shot: int,
|
||||||
|
dev_data: list[dict] | None,
|
||||||
|
subject: str = "",
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Wrap raw MMLU prompt in the model's chat template format.
|
"""Wrap raw MMLU prompt in the model's chat template format.
|
||||||
|
|
||||||
For few-shot, prepend example Q&A pairs as a second user/assistant exchange.
|
For few-shot, prepend example Q&A pairs as user/assistant exchanges.
|
||||||
|
Few-shot examples use the same subject preamble as the test question to
|
||||||
|
keep the format consistent.
|
||||||
"""
|
"""
|
||||||
messages = []
|
messages = []
|
||||||
if n_shot > 0 and dev_data:
|
if n_shot > 0 and dev_data:
|
||||||
for item in dev_data[:n_shot]:
|
for item in dev_data[:n_shot]:
|
||||||
q = f"Question: {item['question']}\n"
|
q = build_prompt(item["question"], item, subject)
|
||||||
for k in ("A", "B", "C", "D"):
|
|
||||||
q += f"{k}. {item[k]}\n"
|
|
||||||
q += "Answer:"
|
|
||||||
messages.append({"role": "user", "content": q})
|
messages.append({"role": "user", "content": q})
|
||||||
messages.append({"role": "assistant", "content": item["answer"]})
|
messages.append({"role": "assistant", "content": item["answer"]})
|
||||||
messages.append({"role": "user", "content": raw_prompt})
|
messages.append({"role": "user", "content": raw_prompt})
|
||||||
@@ -202,6 +206,25 @@ def choice_logprob(
|
|||||||
return score
|
return score
|
||||||
|
|
||||||
|
|
||||||
|
def _permute_choices(item: dict, rng: random.Random) -> tuple[dict, str]:
|
||||||
|
"""Shuffle the option order of a question.
|
||||||
|
|
||||||
|
Returns ``(permuted_item, new_answer_letter)``. The question text and
|
||||||
|
the *content* of each choice are unchanged; only which letter (A/B/C/D)
|
||||||
|
maps to which content is shuffled. This neutralises the model's
|
||||||
|
positional bias (e.g. always picking B).
|
||||||
|
"""
|
||||||
|
letters = ("A", "B", "C", "D")
|
||||||
|
contents = [item[k] for k in letters]
|
||||||
|
perm = list(letters)
|
||||||
|
rng.shuffle(perm)
|
||||||
|
permuted = {"question": item["question"]}
|
||||||
|
for new_letter, orig_letter in zip(letters, perm):
|
||||||
|
permuted[new_letter] = item[orig_letter]
|
||||||
|
new_answer = letters[perm.index(item["answer"])]
|
||||||
|
return permuted, new_answer
|
||||||
|
|
||||||
|
|
||||||
def evaluate_subject(
|
def evaluate_subject(
|
||||||
model,
|
model,
|
||||||
tokenizer,
|
tokenizer,
|
||||||
@@ -210,18 +233,24 @@ def evaluate_subject(
|
|||||||
dev_data: list[dict] | None,
|
dev_data: list[dict] | None,
|
||||||
device: str,
|
device: str,
|
||||||
n_shot: int,
|
n_shot: int,
|
||||||
|
seed: int = 0,
|
||||||
) -> tuple[float, int, int]:
|
) -> tuple[float, int, int]:
|
||||||
|
rng = random.Random(seed) if seed >= 0 else None
|
||||||
correct = 0
|
correct = 0
|
||||||
total = 0
|
total = 0
|
||||||
for item in tqdm.tqdm(test_data, desc=f"{subject:40s}", leave=False):
|
for item in tqdm.tqdm(test_data, desc=f"{subject:40s}", leave=False):
|
||||||
raw_prompt = build_prompt(item["question"], item, subject)
|
if rng is not None:
|
||||||
context = apply_chat(tokenizer, raw_prompt, n_shot, dev_data or [])
|
permuted, answer = _permute_choices(item, rng)
|
||||||
|
else:
|
||||||
|
permuted, answer = item, item["answer"]
|
||||||
|
raw_prompt = build_prompt(permuted["question"], permuted, subject)
|
||||||
|
context = apply_chat(tokenizer, raw_prompt, n_shot, dev_data or [], subject)
|
||||||
context_ids = tokenizer.encode(context)
|
context_ids = tokenizer.encode(context)
|
||||||
scores = {
|
scores = {
|
||||||
c: choice_logprob(model, tokenizer, context_ids, c, device)
|
c: choice_logprob(model, tokenizer, context_ids, c, device)
|
||||||
for c in ("A", "B", "C", "D")
|
for c in ("A", "B", "C", "D")
|
||||||
}
|
}
|
||||||
if max(scores, key=scores.get) == item["answer"]:
|
if max(scores, key=scores.get) == answer:
|
||||||
correct += 1
|
correct += 1
|
||||||
total += 1
|
total += 1
|
||||||
return correct / total, correct, total
|
return correct / total, correct, total
|
||||||
@@ -256,6 +285,12 @@ def main():
|
|||||||
default="bfloat16" if torch.cuda.is_available() else "float32",
|
default="bfloat16" if torch.cuda.is_available() else "float32",
|
||||||
help="Torch dtype",
|
help="Torch dtype",
|
||||||
)
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--seed",
|
||||||
|
type=int,
|
||||||
|
default=0,
|
||||||
|
help="Seed for option permutation (0 to enable, -1 to disable)",
|
||||||
|
)
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
if args.download or not os.path.exists(args.data_dir):
|
if args.download or not os.path.exists(args.data_dir):
|
||||||
@@ -287,7 +322,14 @@ def main():
|
|||||||
test_data = load_csv(test_path)
|
test_data = load_csv(test_path)
|
||||||
|
|
||||||
acc, corr, tot = evaluate_subject(
|
acc, corr, tot = evaluate_subject(
|
||||||
model, tokenizer, subject, test_data, dev_data, device, args.n_shot
|
model,
|
||||||
|
tokenizer,
|
||||||
|
subject,
|
||||||
|
test_data,
|
||||||
|
dev_data,
|
||||||
|
device,
|
||||||
|
args.n_shot,
|
||||||
|
seed=args.seed,
|
||||||
)
|
)
|
||||||
results[subject] = {"accuracy": round(acc, 4), "correct": corr, "total": tot}
|
results[subject] = {"accuracy": round(acc, 4), "correct": corr, "total": tot}
|
||||||
total_correct += corr
|
total_correct += corr
|
||||||
|
|||||||
Reference in New Issue
Block a user