chore: import upstream snapshot with attribution

python/mlc_llm/bench/evaluation/gsm8k.py

@@ -0,0 +1,315 @@
+"""Eval GSM8K with MLCEngine."""
+
+import argparse
+import asyncio
+import json
+import random
+import re
+from datetime import datetime
+from pathlib import Path
+from typing import List, Literal, Optional  # noqa: UP035
+
+import tqdm
+
+from mlc_llm import AsyncMLCEngine
+
+DEVICES = ["cuda", "rocm", "metal", "vulkan"]
+ANSWER_TRIGGER = "The answer is"
+INVALID_ANS = "[invalid]"
+
+
+def extract_answer(text: str, regex: re.Pattern, select_index: int) -> str:
+    """Extract the answer from the text."""
+    match_all = regex.findall(text)
+    if len(match_all) == 0:
+        return INVALID_ANS
+    match = match_all[select_index]
+    if isinstance(match, tuple):
+        match = next(m for m in match if m)
+    match_str: str = match.strip()
+    match_str = match_str.lstrip("$").rstrip(".").replace(",", "")
+    return match_str
+
+
+def extract_ground_truth(text: str) -> str:
+    """Extract the ground truth from the text."""
+    return extract_answer(text, re.compile(r"#### (\-?[0-9\.\,]+)"), 0)
+
+
+def strict_extract_answer(text: str) -> str:
+    """Strictly extract the answer from the text."""
+    return extract_answer(text, re.compile(r"The answer is \$?(\-?[0-9\.\,]+)."), 0)
+
+
+def flexible_extract_answer(text: str) -> str:
+    """Extract the last number from the text."""
+    return extract_answer(text, re.compile(r"(-?[$0-9.,]{2,})|(-?[0-9]+)"), -1)
+
+
+def create_few_shot_prompt(n_shot: int, use_cot: bool, random_order=False) -> str:
+    """
+    Create a prompt for the few-shot learning task.
+
+    Note
+    ----
+    The examples are taken from the paper https://arxiv.org/pdf/2201.11903.pdf page 35.
+    """
+    question, chain, answer = [], [], []
+
+    question.append(
+        "There are 15 trees in the grove. "
+        "Grove workers will plant trees in the grove today. "
+        "After they are done, there will be 21 trees. "
+        "How many trees did the grove workers plant today?"
+    )
+    chain.append(
+        "There are 15 trees originally. "
+        "Then there were 21 trees after some more were planted. "
+        "So there must have been 21 - 15 = 6."
+    )
+    answer.append("6")
+
+    question.append(
+        "If there are 3 cars in the parking lot and 2 more cars arrive, "
+        "how many cars are in the parking lot?"
+    )
+    chain.append("There are originally 3 cars. 2 more cars arrive. 3 + 2 = 5.")
+    answer.append("5")
+
+    question.append(
+        "Leah had 32 chocolates and her sister had 42. If they ate 35, "
+        "how many pieces do they have left in total?"
+    )
+    chain.append(
+        "Originally, Leah had 32 chocolates. "
+        "Her sister had 42. So in total they had 32 + 42 = 74. "
+        "After eating 35, they had 74 - 35 = 39."
+    )
+    answer.append("39")
+
+    question.append(
+        "Jason had 20 lollipops. He gave Denny some lollipops. Now Jason "
+        "has 12 lollipops. How many lollipops did Jason give to Denny?"
+    )
+    chain.append(
+        "Jason started with 20 lollipops. Then he had 12 after giving some "
+        "to Denny. So he gave Denny 20 - 12 = 8."
+    )
+    answer.append("8")
+
+    question.append(
+        "Shawn has five toys. For Christmas, he got two toys each from his "
+        "mom and dad. How many toys does he have now?"
+    )
+    chain.append(
+        "Shawn started with 5 toys. If he got 2 toys each from his mom and "
+        "dad, then that is 4 more toys. 5 + 4 = 9."
+    )
+    answer.append("9")
+
+    question.append(
+        "There were nine computers in the server room. Five more computers "
+        "were installed each day, from monday to thursday. "
+        "How many computers are now in the server room?"
+    )
+    chain.append(
+        "There were originally 9 computers. For each of 4 days, 5 more "
+        "computers were added. So 5 * 4 = 20 computers were added. "
+        "9 + 20 is 29."
+    )
+    answer.append("29")
+
+    question.append(
+        "Michael had 58 golf balls. On tuesday, he lost 23 golf balls. On "
+        "wednesday, he lost 2 more. "
+        "How many golf balls did he have at the end of wednesday?"
+    )
+    chain.append(
+        "Michael started with 58 golf balls. After losing 23 on tuesday, "
+        "he had 58 - 23 = 35. After losing 2 more, "
+        "he had 35 - 2 = 33 golf balls."
+    )
+    answer.append("33")
+
+    question.append(
+        "Olivia has $23. She bought five bagels for $3 each. How much money does she have left?"
+    )
+    chain.append(
+        "Olivia had 23 dollars. "
+        "5 bagels for 3 dollars each will be 5 x 3 = 15 dollars. "
+        "So she has 23 - 15 dollars left. 23 - 15 is 8."
+    )
+    answer.append("8")
+
+    index_list = list(range(len(question)))
+    if random_order:
+        random.shuffle(index_list)
+
+    prompt = ""
+    for i in index_list[:n_shot]:
+        if use_cot:
+            prompt += f"Q: {question[i]}\nA: {chain[i]} {ANSWER_TRIGGER} {answer[i]}.\n\n"
+        else:
+            prompt += f"Question: {question[i]}\nAnswer: {ANSWER_TRIGGER} {answer[i]}.\n\n"
+    return prompt
+
+
+def create_prompt(question: str, n_shot: int, use_cot: bool, random_order: bool = False) -> str:
+    """Create a prompt for the few-shot learning task."""
+    prompt = create_few_shot_prompt(n_shot, use_cot, random_order)
+    if use_cot:
+        prompt += f"Q: {question}\nA:"
+    else:
+        prompt += f"Question: {question}\nAnswer:"
+    return prompt
+
+
+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 GSM8K 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("--n-shot", type=int, default=8)
+    parser.add_argument("--disable_cot", action="store_true", default=False)
+    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,
+):
+    """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=512,
+                stop=["Q:", "Question:"],
+                temperature=0.0,
+            )
+
+    for prompt in prompts:
+        task = asyncio.create_task(generate_task(prompt))
+        tasks.append(task)
+
+    return await tqdm.asyncio.tqdm.gather(*tasks)
+
+
+async def evaluate(
+    model: str,
+    device: str,
+    dataset: Path,
+    model_lib: Optional[str],
+    n_shot: int,
+    use_cot: bool,
+    batch_size: int,
+    log_dir: Optional[Path],
+):
+    """Evaluate GSM8K for the model."""
+    mode: Literal["local", "interactive", "server"] = (
+        "server" if batch_size > 4 else "interactive" if batch_size == 1 else "local"
+    )
+    async_engine = AsyncMLCEngine(model, device=device, model_lib=model_lib, mode=mode)
+
+    with open(dataset / "test.jsonl", encoding="utf-8") as file:
+        tests = [json.loads(line) for line in file]
+
+    prompts = [create_prompt(test["question"], n_shot, use_cot) for test in tests]
+    responses = await send_request(async_engine, prompts, asyncio.Semaphore(batch_size))
+    assert len(responses) == len(tests)
+
+    num_strict_correct, num_flexible_correct = 0, 0
+    num_tests = len(tests)
+    logs = []
+
+    for response, test in zip(responses, tests):
+        response_text = response.choices[0].text.strip()
+        gt_answer = extract_ground_truth(test["answer"])
+        assert gt_answer != INVALID_ANS
+        strict_answer = strict_extract_answer(response_text)
+        flexible_answer = flexible_extract_answer(response_text)
+
+        if gt_answer == strict_extract_answer(response_text):
+            # If the answer is exactly the same as the response, then it is correct
+            num_strict_correct += 1
+            num_flexible_correct += 1
+
+        elif gt_answer == flexible_extract_answer(response_text):
+            # Try flexible extract if the strict match fails
+            num_flexible_correct += 1
+
+        logs.append(
+            {
+                "question": test["question"],
+                "response": response_text,
+                "ground_truth": gt_answer,
+                "strict_answer": strict_answer,
+                "flexible_answer": flexible_answer,
+                "strict_match": gt_answer == strict_answer,
+                "flexible_match": gt_answer == flexible_answer,
+            }
+        )
+
+    results = {
+        "config": {
+            "model": model,
+            "device": device,
+            "model_lib": model_lib,
+            "n_shot": n_shot,
+            "use_cot": use_cot,
+        },
+        "results": {
+            "strict_match": num_strict_correct,
+            "flexible_match": num_flexible_correct,
+            "total": num_tests,
+        },
+    }
+    print(
+        f"Strict Matching Accuracy: {num_strict_correct} / {num_tests} = "
+        f"{num_strict_correct / num_tests * 100:.2f}%"
+    )
+    print(
+        f"Flexible Matching Accuracy: {num_flexible_correct} / {num_tests} = "
+        f"{num_flexible_correct / num_tests * 100:.2f}%"
+    )
+
+    if log_dir:
+        with open(log_dir / "summary.json", "w", encoding="utf-8") as f:
+            json.dump(results, f, indent=2)
+        with open(log_dir / "logs.json", "w", encoding="utf-8") as f:
+            json.dump(logs, 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.mkdir(parents=True, exist_ok=True)
+    asyncio.run(
+        evaluate(
+            model=args.model,
+            device=args.device,
+            dataset=args.dataset,
+            model_lib=args.model_lib,
+            n_shot=args.n_shot,
+            use_cot=not args.disable_cot,
+            batch_size=args.batch_size,
+            log_dir=log_dir,
+        )
+    )
+    end_time = datetime.now()
+    print(f"Time used: {end_time - start_time}")

python/mlc_llm/bench/evaluation/mmlu.py

@@ -0,0 +1,245 @@
+"""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}")