MMLU

from datasets import load_dataset
from openai import OpenAI
import re

# MMLU 데이터셋 로드 (HuggingFace)
dataset = load_dataset("cais/mmlu", "all", split="test")

# 특정 과목 필터링 (예: computer_science)
cs_questions = [q for q in dataset if q["subject"] == "high_school_computer_science"]

client = OpenAI()

def evaluate_mmlu_question(question, choices, answer_idx):
    """MMLU 문제 평가 함수"""
    # 5-shot 프롬프트 구성 (실제로는 few-shot 예시 추가)
    prompt = f"""Answer the following multiple choice question. Only output A, B, C, or D.

Question: {question}
A) {choices[0]}
B) {choices[1]}
C) {choices[2]}
D) {choices[3]}

Answer:"""

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=1,
        temperature=0
    )

    predicted = response.choices[0].message.content.strip().upper()
    correct_answer = ["A", "B", "C", "D"][answer_idx]

    return predicted == correct_answer

# 샘플 평가 실행
correct = 0
total = min(50, len(cs_questions))  # 샘플 50문제

for i, q in enumerate(cs_questions[:total]):
    is_correct = evaluate_mmlu_question(
        q["question"],
        q["choices"],
        q["answer"]
    )
    correct += int(is_correct)
    if (i + 1) % 10 == 0:
        print(f"Progress: {i+1}/{total}, Accuracy: {correct/(i+1)*100:.1f}%")

print(f"\nFinal Accuracy: {correct/total*100:.1f}%")
print(f"Subject: high_school_computer_science")

# lm-evaluation-harness 사용 (더 정확한 평가)
# pip install lm-eval
# lm_eval --model openai --model_args model=gpt-4o-mini --tasks mmlu --batch_size 8