378 lines
10 KiB
Python
Executable File
378 lines
10 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import argparse
|
|
import json
|
|
import os
|
|
import random
|
|
from pathlib import Path
|
|
|
|
from datasets import get_dataset_config_names, load_dataset
|
|
|
|
|
|
MMLU_DATASET = "cais/mmlu"
|
|
GPQA_DATASET = "Idavidrein/gpqa"
|
|
CEVAL_DATASET = "ceval/ceval-exam"
|
|
CMMLU_DATASET = "haonan-li/cmmlu"
|
|
|
|
MMLU_CONFIGS = [
|
|
"abstract_algebra",
|
|
"anatomy",
|
|
"astronomy",
|
|
"business_ethics",
|
|
"clinical_knowledge",
|
|
"college_biology",
|
|
"college_chemistry",
|
|
"college_computer_science",
|
|
"college_mathematics",
|
|
"college_physics",
|
|
"computer_security",
|
|
"conceptual_physics",
|
|
"econometrics",
|
|
"electrical_engineering",
|
|
"elementary_mathematics",
|
|
"formal_logic",
|
|
"global_facts",
|
|
"high_school_biology",
|
|
"high_school_chemistry",
|
|
"high_school_computer_science",
|
|
"high_school_government_and_politics",
|
|
"high_school_macroeconomics",
|
|
"high_school_mathematics",
|
|
"high_school_physics",
|
|
"high_school_statistics",
|
|
"international_law",
|
|
"jurisprudence",
|
|
"machine_learning",
|
|
"management",
|
|
"marketing",
|
|
"medical_genetics",
|
|
"moral_disputes",
|
|
"nutrition",
|
|
"philosophy",
|
|
"professional_law",
|
|
"professional_medicine",
|
|
"professional_psychology",
|
|
"public_relations",
|
|
"security_studies",
|
|
"sociology",
|
|
"us_foreign_policy",
|
|
"world_religions",
|
|
]
|
|
|
|
CEVAL_CONFIGS = [
|
|
"computer_network",
|
|
"operating_system",
|
|
"computer_architecture",
|
|
"college_programming",
|
|
"college_physics",
|
|
"college_chemistry",
|
|
"advanced_mathematics",
|
|
"probability_and_statistics",
|
|
"discrete_mathematics",
|
|
"electrical_engineer",
|
|
"metrology_engineer",
|
|
"high_school_mathematics",
|
|
"high_school_physics",
|
|
"high_school_chemistry",
|
|
"high_school_biology",
|
|
"legal_professional",
|
|
"business_administration",
|
|
"marxism",
|
|
"mao_zedong_thought",
|
|
"education_science",
|
|
"teacher_qualification",
|
|
"modern_chinese_history",
|
|
"chinese_language_and_literature",
|
|
"logic",
|
|
]
|
|
|
|
CMMLU_CONFIGS = [
|
|
"agronomy",
|
|
"anatomy",
|
|
"ancient_chinese",
|
|
"arts",
|
|
"astronomy",
|
|
"business_ethics",
|
|
"chinese_civil_service_exam",
|
|
"chinese_driving_rule",
|
|
"chinese_food_culture",
|
|
"chinese_foreign_policy",
|
|
"chinese_history",
|
|
"college_actuarial_science",
|
|
"college_education",
|
|
"college_engineering_hydrology",
|
|
"college_law",
|
|
"college_mathematics",
|
|
"college_medical_statistics",
|
|
"college_medicine",
|
|
"computer_science",
|
|
"conceptual_physics",
|
|
"econometrics",
|
|
"education",
|
|
"electrical_engineering",
|
|
"elementary_chinese",
|
|
"elementary_commonsense",
|
|
"elementary_information_and_technology",
|
|
"elementary_mathematics",
|
|
"ethnology",
|
|
"food_science",
|
|
"genetics",
|
|
"global_facts",
|
|
"high_school_biology",
|
|
"high_school_chemistry",
|
|
"high_school_geography",
|
|
"high_school_mathematics",
|
|
"high_school_physics",
|
|
"human_sexuality",
|
|
"international_law",
|
|
"journalism",
|
|
"jurisprudence",
|
|
"legal_and_moral_basis",
|
|
"logical",
|
|
"machine_learning",
|
|
"management",
|
|
"marketing",
|
|
"marxist_theory",
|
|
"modern_chinese",
|
|
"nutrition",
|
|
"philosophy",
|
|
"professional_accounting",
|
|
"professional_law",
|
|
"professional_medicine",
|
|
"professional_psychology",
|
|
"public_relations",
|
|
"security_study",
|
|
"sociology",
|
|
"sports_science",
|
|
"traditional_chinese_medicine",
|
|
"virology",
|
|
"world_history",
|
|
"world_religions",
|
|
]
|
|
|
|
|
|
def norm(text):
|
|
return " ".join(str(text or "").replace("\x00", " ").split()).strip()
|
|
|
|
|
|
def stable_sample(rng, rows, n):
|
|
if len(rows) <= n:
|
|
return list(rows)
|
|
return rng.sample(rows, n)
|
|
|
|
|
|
def add_item(items, prefix, idx, source, subset, question, choices, answer_idx, metadata=None):
|
|
question = norm(question)
|
|
choices = [norm(x) for x in choices]
|
|
if not question or len(choices) < 2:
|
|
return
|
|
if answer_idx < 0 or answer_idx >= len(choices):
|
|
return
|
|
answer_text = choices[answer_idx]
|
|
if not answer_text:
|
|
return
|
|
mcq_prompt = (
|
|
"Answer the following multiple-choice question. Choose the single best option.\n\n"
|
|
f"Question:\n{question}\n\nAnswer:"
|
|
)
|
|
ppl_text = f"Question: {question}\nCorrect answer: {answer_text}"
|
|
items.append(
|
|
{
|
|
"id": f"{prefix}_{idx:05d}",
|
|
"category": prefix,
|
|
"source": source,
|
|
"subset": subset,
|
|
"ppl_text": ppl_text,
|
|
"mcq_prompt": mcq_prompt,
|
|
"choices": choices,
|
|
"answer_idx": int(answer_idx),
|
|
"answer_text": answer_text,
|
|
"metadata": metadata or {},
|
|
}
|
|
)
|
|
|
|
|
|
def build_mmlu(rng, quota):
|
|
configs = list(MMLU_CONFIGS)
|
|
rng.shuffle(configs)
|
|
rows = []
|
|
per_config = max(20, quota // 24)
|
|
for cfg in configs:
|
|
print(f"[info] load mmlu {cfg}", flush=True)
|
|
try:
|
|
ds = load_dataset(MMLU_DATASET, cfg, split="test")
|
|
except Exception as exc:
|
|
print(f"[warn] skip MMLU {cfg}: {type(exc).__name__}: {str(exc)[:160]}")
|
|
continue
|
|
sampled = stable_sample(rng, list(ds), min(per_config, len(ds)))
|
|
for r in sampled:
|
|
rows.append((cfg, r))
|
|
if len(rows) >= quota * 2:
|
|
break
|
|
rng.shuffle(rows)
|
|
items = []
|
|
for i, (cfg, r) in enumerate(rows[:quota]):
|
|
add_item(
|
|
items,
|
|
"mmlu",
|
|
i,
|
|
MMLU_DATASET,
|
|
cfg,
|
|
r.get("question"),
|
|
r.get("choices", []),
|
|
r.get("answer"),
|
|
{"subject": r.get("subject", cfg)},
|
|
)
|
|
return items
|
|
|
|
|
|
def build_gpqa(rng, quota, config="gpqa_main"):
|
|
ds = load_dataset(GPQA_DATASET, config, split="train", token=True)
|
|
rows = stable_sample(rng, list(ds), min(quota, len(ds)))
|
|
items = []
|
|
for i, r in enumerate(rows):
|
|
choices = [
|
|
r.get("Correct Answer"),
|
|
r.get("Incorrect Answer 1"),
|
|
r.get("Incorrect Answer 2"),
|
|
r.get("Incorrect Answer 3"),
|
|
]
|
|
order = list(range(4))
|
|
rng.shuffle(order)
|
|
shuffled = [choices[j] for j in order]
|
|
answer_idx = order.index(0)
|
|
add_item(
|
|
items,
|
|
"gpqa",
|
|
i,
|
|
GPQA_DATASET,
|
|
config,
|
|
r.get("Question"),
|
|
shuffled,
|
|
answer_idx,
|
|
{
|
|
"high_level_domain": r.get("High-level domain"),
|
|
"subdomain": r.get("Subdomain"),
|
|
"record_id": r.get("Record ID"),
|
|
},
|
|
)
|
|
return items
|
|
|
|
|
|
def letter_answer_idx(ans):
|
|
if ans is None:
|
|
return -1
|
|
s = str(ans).strip()
|
|
if s in {"0", "1", "2", "3"}:
|
|
return int(s)
|
|
if s:
|
|
ch = s[0].upper()
|
|
if ch in "ABCD":
|
|
return ord(ch) - ord("A")
|
|
return -1
|
|
|
|
|
|
def build_abcd_dataset(rng, dataset_name, prefix, quota):
|
|
if prefix == "ceval":
|
|
configs = list(CEVAL_CONFIGS)
|
|
elif prefix == "cmmlu":
|
|
configs = list(CMMLU_CONFIGS)
|
|
else:
|
|
configs = get_dataset_config_names(dataset_name)
|
|
rng.shuffle(configs)
|
|
rows = []
|
|
per_config = max(24, quota // 18)
|
|
for cfg in configs:
|
|
print(f"[info] load {prefix} {cfg}", flush=True)
|
|
split = "test"
|
|
try:
|
|
ds = load_dataset(dataset_name, cfg, split=split)
|
|
except Exception:
|
|
try:
|
|
split = "val"
|
|
ds = load_dataset(dataset_name, cfg, split=split)
|
|
except Exception as exc:
|
|
print(f"[warn] skip {prefix} {cfg}: {type(exc).__name__}: {str(exc)[:160]}")
|
|
continue
|
|
sampled = stable_sample(rng, list(ds), min(per_config, len(ds)))
|
|
for r in sampled:
|
|
rows.append((cfg, split, r))
|
|
if len(rows) >= quota * 2:
|
|
break
|
|
rng.shuffle(rows)
|
|
items = []
|
|
for i, (cfg, split, r) in enumerate(rows[:quota]):
|
|
choices = [r.get("A"), r.get("B"), r.get("C"), r.get("D")]
|
|
add_item(
|
|
items,
|
|
prefix,
|
|
i,
|
|
dataset_name,
|
|
cfg,
|
|
r.get("question"),
|
|
choices,
|
|
letter_answer_idx(r.get("answer")),
|
|
{"split": split},
|
|
)
|
|
return items
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--out-dir", required=True)
|
|
ap.add_argument("--seed", type=int, default=20260607)
|
|
ap.add_argument("--mmlu", type=int, default=700)
|
|
ap.add_argument("--gpqa", type=int, default=400)
|
|
ap.add_argument("--ceval", type=int, default=450)
|
|
ap.add_argument("--cmmlu", type=int, default=450)
|
|
args = ap.parse_args()
|
|
|
|
rng = random.Random(args.seed)
|
|
out_dir = Path(args.out_dir)
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
print("[info] HF_ENDPOINT", os.environ.get("HF_ENDPOINT", ""))
|
|
builders = [
|
|
("mmlu", lambda: build_mmlu(rng, args.mmlu)),
|
|
("gpqa", lambda: build_gpqa(rng, args.gpqa)),
|
|
("ceval", lambda: build_abcd_dataset(rng, CEVAL_DATASET, "ceval", args.ceval)),
|
|
("cmmlu", lambda: build_abcd_dataset(rng, CMMLU_DATASET, "cmmlu", args.cmmlu)),
|
|
]
|
|
all_items = []
|
|
errors = {}
|
|
for name, fn in builders:
|
|
if getattr(args, name) <= 0:
|
|
print(f"[info] skip {name}: quota=0")
|
|
continue
|
|
try:
|
|
items = fn()
|
|
print(f"[info] built {name}: {len(items)}")
|
|
all_items.extend(items)
|
|
except Exception as exc:
|
|
errors[name] = f"{type(exc).__name__}: {str(exc)[:500]}"
|
|
print(f"[error] {name}: {errors[name]}")
|
|
|
|
rng.shuffle(all_items)
|
|
bench_path = out_dir / "heldout_public_mcq_2k.jsonl"
|
|
with bench_path.open("w", encoding="utf-8") as f:
|
|
for item in all_items:
|
|
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
|
|
|
counts = {}
|
|
for item in all_items:
|
|
counts[item["category"]] = counts.get(item["category"], 0) + 1
|
|
stats = {
|
|
"seed": args.seed,
|
|
"requested": {"mmlu": args.mmlu, "gpqa": args.gpqa, "ceval": args.ceval, "cmmlu": args.cmmlu},
|
|
"total_items": len(all_items),
|
|
"counts": counts,
|
|
"errors": errors,
|
|
"output": str(bench_path),
|
|
}
|
|
with (out_dir / "heldout_public_mcq_2k_stats.json").open("w", encoding="utf-8") as f:
|
|
json.dump(stats, f, ensure_ascii=False, indent=2)
|
|
print(json.dumps(stats, ensure_ascii=False, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|