File size: 5,610 Bytes
8a5ffa8 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | """
MMLU data construction for Claim 5 (paper Section 5.2 / Appendix B.2).
Builds, deterministically (seed 0):
* a 1,000-example Humanities fine-tuning set,
* a 4x250 = 1,000-example evaluation set (Humanities / Social Sciences /
STEM / Other) that is DISJOINT from the fine-tuning set,
* 7 few-shot exemplars per subject drawn from the MMLU `validation` split
(topped up from `dev` when a subject has fewer than 7 validation rows),
so no exemplar can ever appear in the fine-tuning set or the eval set.
The paper fine-tunes Qwen2.5-3B-Instruct on "the Humanities subset of MMLU,
which contains 1,000 training examples" and evaluates on "the MMLU test set
across four categories ... consisting of 1,000 examples in total". MMLU ships
no dedicated Humanities train split, so we carve the 1,000 fine-tuning rows out
of the Humanities *test* pool and hold the 250 evaluation rows out of it.
"""
import random
HUMANITIES = [
"formal_logic", "high_school_european_history", "high_school_us_history",
"high_school_world_history", "international_law", "jurisprudence",
"logical_fallacies", "moral_disputes", "moral_scenarios", "philosophy",
"prehistory", "professional_law", "world_religions",
]
STEM = [
"abstract_algebra", "astronomy", "college_biology", "college_chemistry",
"college_computer_science", "college_mathematics", "college_physics",
"computer_security", "conceptual_physics", "electrical_engineering",
"elementary_mathematics", "high_school_biology", "high_school_chemistry",
"high_school_computer_science", "high_school_mathematics",
"high_school_physics", "high_school_statistics", "machine_learning",
]
SOCIAL = [
"econometrics", "high_school_geography",
"high_school_government_and_politics", "high_school_macroeconomics",
"high_school_microeconomics", "high_school_psychology", "human_sexuality",
"professional_psychology", "public_relations", "security_studies",
"sociology", "us_foreign_policy",
]
OTHER = [
"anatomy", "business_ethics", "clinical_knowledge", "college_medicine",
"global_facts", "human_aging", "management", "marketing",
"medical_genetics", "miscellaneous", "nutrition", "professional_accounting",
"professional_medicine", "virology",
]
CATEGORIES = {
"Humanities": HUMANITIES,
"Social Sciences": SOCIAL,
"STEM": STEM,
"Other": OTHER,
}
SUBJ2CAT = {s: c for c, ss in CATEGORIES.items() for s in ss}
LETTERS = ["A", "B", "C", "D"]
N_TRAIN = 1000
N_EVAL_PER_CAT = 250
N_SHOT = 7
DATA_SEED = 0
def _pretty(subject):
return subject.replace("_", " ")
def question_text(row):
"""Zero-shot user turn for one MMLU row."""
lines = [
f"The following is a multiple choice question about {_pretty(row['subject'])}.",
"",
row["question"].strip(),
]
for letter, choice in zip(LETTERS, row["choices"]):
lines.append(f"{letter}. {str(choice).strip()}")
lines.append("")
lines.append("Answer with the letter of the correct option (A, B, C, or D).")
return "\n".join(lines)
def answer_text(row):
return LETTERS[int(row["answer"])]
def _rows(ds_split, subjects):
out = []
for r in ds_split:
if r["subject"] in subjects:
out.append({
"subject": r["subject"],
"question": r["question"],
"choices": list(r["choices"]),
"answer": int(r["answer"]),
})
return out
def build(cache_dir=None):
from datasets import load_dataset
dd = load_dataset("cais/mmlu", "all", cache_dir=cache_dir)
test, val, dev = dd["test"], dd["validation"], dd["dev"]
all_subjects = sorted(SUBJ2CAT)
# ---- few-shot exemplars: validation split first, dev as top-up -------
val_rows, dev_rows = {}, {}
for r in _rows(val, set(all_subjects)):
val_rows.setdefault(r["subject"], []).append(r)
for r in _rows(dev, set(all_subjects)):
dev_rows.setdefault(r["subject"], []).append(r)
shots = {}
for s in all_subjects:
pool = list(val_rows.get(s, []))
random.Random(DATA_SEED).shuffle(pool)
picked = pool[:N_SHOT]
if len(picked) < N_SHOT:
picked = picked + dev_rows.get(s, [])[: N_SHOT - len(picked)]
assert len(picked) == N_SHOT, (s, len(picked))
shots[s] = picked
# ---- eval sets (250 per category, from `test`) -----------------------
eval_sets, used_hum = {}, set()
for ci, (cat, subjects) in enumerate(CATEGORIES.items()):
pool = _rows(test, set(subjects))
idx = list(range(len(pool)))
random.Random(DATA_SEED + 100 * (ci + 1)).shuffle(idx)
chosen = idx[:N_EVAL_PER_CAT]
eval_sets[cat] = [pool[i] for i in chosen]
if cat == "Humanities":
hum_pool = pool
used_hum = set(chosen)
# ---- 1,000 Humanities fine-tuning rows, disjoint from eval ----------
remaining = [i for i in range(len(hum_pool)) if i not in used_hum]
random.Random(DATA_SEED + 7).shuffle(remaining)
train_rows = [hum_pool[i] for i in remaining[:N_TRAIN]]
assert len(train_rows) == N_TRAIN
return {"train": train_rows, "eval": eval_sets, "shots": shots}
def chat_messages(row, shots=None):
"""Multi-turn chat prompt. `shots` = list of exemplar rows (same subject)."""
msgs = []
for ex in (shots or []):
msgs.append({"role": "user", "content": question_text(ex)})
msgs.append({"role": "assistant", "content": answer_text(ex)})
msgs.append({"role": "user", "content": question_text(row)})
return msgs
|