from lm_eval import evaluator

MODEL_DIR = "./hf_gpt2_nanogpt"

# GPT-2 技术报告相关数据集任务
TASKS = [
    "lambada_openai",
    "cbt_common",
    "cbt_ne",
    "cbt_v",
    "cbt_p",
    "triviaqa",
    "webqs",
    "arc_easy",
    "arc_challenge",
    "winogrande",
    "race",
    "story_cloze",
]

print("🚀 正在运行完整评测任务集（参考 GPT-2 技术报告）...")

results = evaluator.simple_evaluate(
    model="hf",
    model_args=f"pretrained={MODEL_DIR}",
    tasks=TASKS,
    batch_size=8,
    device="cuda"
)

print("✅ 评测完成，结果为：")
for task, metrics in results["results"].items():
    if isinstance(metrics, dict):
        metric_str = ", ".join([f"{k}: {v:.4f}" for k, v in metrics.items()])
        print(f"{task}: {metric_str}")