mlc-ai--mlc-llm/python/mlc_llm/bench/evaluation/mmlu.py

"""Eval MMLU with MLCEngine."""

import argparse
import asyncio
import csv
import json
import string
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional  # noqa: UP035

import numpy as np
import tqdm

from mlc_llm import AsyncMLCEngine

SUBJECTS = [
    "abstract_algebra",
    "anatomy",
    "astronomy",
    "business_ethics",
    "clinical_knowledge",
    "college_biology",
    "college_chemistry",
    "college_computer_science",
    "college_mathematics",
    "college_medicine",
    "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_european_history",
    "high_school_geography",
    "high_school_government_and_politics",
    "high_school_macroeconomics",
    "high_school_mathematics",
    "high_school_microeconomics",
    "high_school_physics",
    "high_school_psychology",
    "high_school_statistics",
    "high_school_us_history",
    "high_school_world_history",
    "human_aging",
    "human_sexuality",
    "international_law",
    "jurisprudence",
    "logical_fallacies",
    "machine_learning",
    "management",
    "marketing",
    "medical_genetics",
    "miscellaneous",
    "moral_disputes",
    "moral_scenarios",
    "nutrition",
    "philosophy",
    "prehistory",
    "professional_accounting",
    "professional_law",
    "professional_medicine",
    "professional_psychology",
    "public_relations",
    "security_studies",
    "sociology",
    "us_foreign_policy",
    "virology",
    "world_religions",
]
PADDING_LEN = max(len(subject) for subject in SUBJECTS)
DEVICES = ["cuda", "rocm", "metal", "vulkan"]
PROMPT_TEMPLATE = string.Template("$Q\nA. $A\nB. $B\nC. $C\nD. $D\nAnswer:")


def parse_args():
    """Parse command line arguments."""

    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, required=True)
    parser.add_argument(
        "--dataset", type=Path, required=True, help="Path to MMLU test dataset home."
    )
    parser.add_argument("--device", type=str, choices=["auto", *DEVICES], default="auto")
    parser.add_argument("--model-lib", type=str, default=None)
    parser.add_argument("-s", "--subject", nargs="+", type=str, choices=SUBJECTS, default=SUBJECTS)
    parser.add_argument("-bs", "--batch-size", type=int, default=16)
    parser.add_argument("--log-dir", type=Path, default=None)
    return parser.parse_args()


async def send_request(
    async_engine: AsyncMLCEngine,
    prompts: List[str],  # noqa: UP006
    semaphore: asyncio.Semaphore,
    subject: str,
):
    """Send the calibration requests to the engine."""
    tasks = []

    async def generate_task(prompt):
        async with semaphore:
            return await async_engine.completions.create(
                prompt=prompt,
                stream=False,
                max_tokens=1,
                temperature=1.0,
                logprobs=True,
                top_logprobs=5,
            )

    for prompt in prompts:
        task = asyncio.create_task(generate_task(prompt))
        tasks.append(task)

    return await tqdm.asyncio.tqdm.gather(
        *tasks,
        desc=f"Running {subject.ljust(PADDING_LEN)}",
        bar_format="{desc} {percentage:3.0f}%|{bar}{r_bar}",
    )


async def evaluate(
    model: str,
    device: str,
    dataset: Path,
    model_lib: Optional[str],
    subjects: List[str],  # noqa: UP006
    semaphore: asyncio.Semaphore,
    log_dir: Optional[Path],
):
    """Evaluate MMLU for the model."""
    async_engine = AsyncMLCEngine(model, device=device, model_lib=model_lib, mode="server")

    results: Dict[str, Any] = {}  # noqa: UP006
    for subject in subjects:
        with open(dataset / "test" / f"{subject}_test.csv", encoding="utf-8") as csvfile:
            tests = list(csv.reader(csvfile, delimiter=",", quotechar='"'))
            assert all(len(test) == 6 for test in tests)

        logs = []
        num_correct = 0
        prompts = [
            PROMPT_TEMPLATE.substitute(Q=test[0], A=test[1], B=test[2], C=test[3], D=test[4])
            for test in tests
        ]
        responses = await send_request(async_engine, prompts, semaphore, subject)

        assert len(responses) == len(tests)
        for response, test in zip(responses, tests):
            token_logprobs = {}
            logprobs = response.choices[0].logprobs.content[0].top_logprobs
            for logprob in logprobs:
                if logprob.token not in token_logprobs:
                    token_logprobs[logprob.token] = logprob.logprob

            abcd_logprobs = {}
            for choice in ["A", "B", "C", "D"]:
                abcd_logprobs[choice] = token_logprobs[choice] if choice in token_logprobs else -100

            pred = {0: "A", 1: "B", 2: "C", 3: "D"}[int(np.argmax(list(abcd_logprobs.values())))]
            num_correct += pred == test[5]

            logs.append(
                {
                    "Question": {
                        "Q": test[0],
                        "A": test[1],
                        "B": test[2],
                        "C": test[3],
                        "D": test[4],
                    },
                    "Answer": test[5],
                    "Response": {
                        "pred": pred,
                        "logprobs": list(abcd_logprobs.values()),
                    },
                }
            )

        results[subject] = {
            "correct": num_correct,
            "total": len(tests),
            "accuracy": num_correct / len(tests),
        }

        if log_dir:
            with open(log_dir / "subjects" / f"{subject}.json", "w", encoding="utf-8") as f:
                json.dump(logs, f, indent=2)

    total_correct, total_tests = 0, 0
    for subject, v in results.items():
        num_correct, num_tests, accuracy = v["correct"], v["total"], v["accuracy"]
        print(f"{subject}: {num_correct} / {num_tests} = {accuracy * 100:.2f}%")
        total_correct += num_correct
        total_tests += num_tests

    total_accuracy = total_correct / total_tests
    results["total"] = {
        "correct": total_correct,
        "total": total_tests,
        "accuracy": total_accuracy,
    }
    print(f"Total accuracy: {total_correct} / {total_tests} = {total_accuracy * 100:.2f}%")

    if log_dir:
        results = {
            "config": {
                "model": model,
                "device": device,
                "model_lib": model_lib,
                "subjects": subjects,
            },
            "results": results,
        }
        with open(log_dir / "summary.json", "w", encoding="utf-8") as f:
            json.dump(results, f, indent=2)


if __name__ == "__main__":
    args = parse_args()
    start_time = datetime.now()
    log_dir: Optional[Path] = None
    if args.log_dir is not None:
        time_dir = start_time.strftime("%Y-%m-%d_%H-%M-%S")
        log_dir = args.log_dir / time_dir
        (log_dir / "subjects").mkdir(parents=True, exist_ok=True)
    asyncio.run(
        evaluate(
            model=args.model,
            device=args.device,
            dataset=args.dataset,
            model_lib=args.model_lib,
            subjects=args.subject,
            semaphore=asyncio.Semaphore(args.batch_size),
            log_dir=log_dir,
        )
    )
    end_time = datetime.now()
    print(f"Time used: {end_time - start_time}")