MicroLLM2/mmlu_bench.py

#!/usr/bin/env python3
"""
MicroLLM2 — MMLU Benchmark
Evaluates MLVXN/MicroLLM2 (or local checkpoint) on MMLU (5-shot by default)
Uses lm-evaluation-harness if available, else lightweight HF implementation.

Usage:
  python mmlu_bench.py                        # 5-shot MMLU on MLVXN/MicroLLM2
  python mmlu_bench.py --model local          # use /home/zeus/microllm2/microllm2-checkpoints/final_merged
  python mmlu_bench.py --shots 0              # zero-shot
  python mmlu_bench.py --subset abstract_algebra,philosophy  # only those subjects
  python mmlu_bench.py --limit 20             # 20 samples per subject for quick smoke test

Outputs: prints per-subject + average accuracy, saves mmlu_results.json
"""
import os, sys, json, argparse, re
from pathlib import Path

os.environ["HF_HUB_DISABLE_XET"]="1"
os.environ["TOKENIZERS_PARALLELISM"]="false"

parser = argparse.ArgumentParser(description="MicroLLM2 MMLU benchmark")
parser.add_argument("--model", default="auto", help="HF id or 'local' or path; default auto -> local if exists else MLVXN/MicroLLM2")
parser.add_argument("--shots", type=int, default=5, help="few-shot examples (0-5)")
parser.add_argument("--limit", type=int, default=None, help="max samples per subject (None=all)")
parser.add_argument("--subset", type=str, default=None, help="comma-separated MMLU subjects to run")
parser.add_argument("--batch", type=int, default=8)
parser.add_argument("--output", default="mmlu_results.json")
args = parser.parse_args()

LOCAL = Path("/home/zeus/microllm2/microllm2-checkpoints/final_merged")
HF_ID = "MLVXN/MicroLLM2"
if args.model == "auto":
    MODEL_ID = str(LOCAL) if LOCAL.exists() else HF_ID
elif args.model == "local":
    MODEL_ID = str(LOCAL)
else:
    MODEL_ID = args.model

print(f"[*] MicroLLM2 MMLU — model: {MODEL_ID}  shots={args.shots}  limit={args.limit}")
print(f"[*] GPT2-XL 1.5B 1024ctx  vocab=50259 (ChatML) — MMLU via direct eval (no harness needed)")

# Try harness first — if installed use it (more accurate), else fallback
USE_HARNESS = False
try:
    import lm_eval  # noqa
    USE_HARNESS = True
except ImportError:
    USE_HARNESS = False

if USE_HARNESS:
    print("[*] Detected lm-evaluation-harness — using official MMLU task")
    # harness expects HF model type; gpt2 works
    import lm_eval
    from lm_eval.models.huggingface import HFLM
    from lm_eval.tasks.mmlu import MMLUTask  # if available
    print("[*] Running: lm_eval --model hf --model_args pretrained={} --tasks mmlu --num_fewshot {} --batch_size {} {}".format(
        MODEL_ID, args.shots, args.batch, f"--limit {args.limit}" if args.limit else ""))
    # delegate to CLI so output is standard
    import subprocess
    cmd = [
        sys.executable, "-m", "lm_eval",
        "--model", "hf",
        "--model_args", f"pretrained={MODEL_ID},dtype=bfloat16,trust_remote_code=False",
        "--tasks", "mmlu",
        "--num_fewshot", str(args.shots),
        "--batch_size", str(args.batch),
        "--output_path", args.output,
    ]
    if args.limit:
        cmd += ["--limit", str(args.limit)]
    print(" ".join(cmd))
    subprocess.run(cmd, check=False)
    if Path(args.output).exists():
        print(f"[+] Saved to {args.output}")
        # also print summary if harness wrote it
        try:
            data = json.loads(open(args.output).read())
            # harness output is nested; try to find results
            print(json.dumps(data.get("results", data), indent=2)[:4000])
        except: pass
    sys.exit(0)

# --- Lightweight fallback: direct HF evaluation (no harness) ---
print("[*] lm-eval not installed — using lightweight direct MMLU eval (same logic, no harness)")
print("[*] Install harness for official numbers: pip install lm-eval==0.4.4")

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset
from tqdm import tqdm

# MMLU subjects (57) — full list from hendrycks/mmlu or cais/mmlu
MMLU_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"
]
if args.subset:
    wanted = [s.strip() for s in args.subset.split(",") if s.strip()]
    MMLU_SUBJECTS = [s for s in MMLU_SUBJECTS if s in wanted]
    print(f"[*] Subset: {MMLU_SUBJECTS}")

# Load model
print(f"[*] Loading tokenizer + model {MODEL_ID} ...")
tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=False)
if tok.pad_token is None: tok.pad_token = tok.eos_token
if "<|im_start|>" not in tok.get_vocab():
    try: tok.add_special_tokens({"additional_special_tokens":["<|im_start|>","<|im_end|>"]})
    except: pass

dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
device_map = "auto" if torch.cuda.is_available() else None
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=dtype, device_map=device_map, trust_remote_code=False)
model.eval()
device = next(model.parameters()).device
print(f"[+] Loaded on {device} dtype={dtype} — starting MMLU")

CHOICES = ["A","B","C","D"]

def format_example(question, choices, answer=None, include_answer=False):
    # Standard MMLU 5-shot format (Hendrycks)
    prompt = question.strip() + "\n"
    for i, c in enumerate(choices):
        prompt += f"{CHOICES[i]}. {c}\n"
    prompt += "Answer:"
    if include_answer and answer is not None:
        # answer is 0-3 int or letter
        if isinstance(answer, int): ans = CHOICES[answer]
        else: ans = str(answer).strip().upper()[0]
        prompt += f" {ans}"
    return prompt

def get_answer_letter(example):
    a = example["answer"]
    if isinstance(a, int): return CHOICES[a]
    return str(a).strip().upper()[0]

# cache datasets by subject
results = {}
overall_correct = 0
overall_total = 0

# Load MMLU from cais/mmlu (canonical) with fallback to hendrycks
def load_mmlu_subject(subject):
    for name in ["cais/mmlu", "hendrycks/mmlu"]:
        try:
            ds = load_dataset(name, subject)
            return ds
        except Exception as e:
            continue
    raise RuntimeError(f"Could not load MMLU subject {subject}")

for subject in tqdm(MMLU_SUBJECTS, desc="MMLU subjects"):
    print(f"\n{'='*60}\n[>] {subject}  (shots={args.shots})")
    try:
        ds = load_mmlu_subject(subject)
    except Exception as e:
        print(f"[!] Skip {subject}: {e}")
        continue
    # hendrycks/mmlu has test split; cais/mmlu has test
    dev = ds.get("dev") or ds.get("validation") or ds["train"]
    test = ds.get("test") or ds.get("validation") or ds["train"]
    if args.limit:
        test = test.select(range(min(args.limit, len(test))))
    # build few-shot prefix from dev (5 examples)
    few_shot_prefix = ""
    if args.shots > 0:
        shots = min(args.shots, len(dev))
        for i in range(shots):
            ex = dev[i]
            q, ch, ans = ex["question"], ex["choices"], ex["answer"]
            few_shot_prefix += format_example(q, ch, ans, include_answer=True) + "\n\n"

    correct = 0
    total = 0
    # Evaluate — score by logprob of A/B/C/D next token (proper MMLU method)
    # For GPT2 we compute which choice token has highest logit after "Answer:"
    for ex in tqdm(test, desc=subject, leave=False):
        q, ch, ans = ex["question"], ex["choices"], ex["answer"]
        true_letter = get_answer_letter(ex)
        prompt = few_shot_prefix + format_example(q, ch, include_answer=False)
        # Tokenize prompt
        inputs = tok(prompt, return_tensors="pt", truncation=True, max_length=900).to(device)
        with torch.no_grad():
            logits = model(**inputs).logits[0, -1]  # last token logits
        # Get logits for " A", " B", etc. (with leading space)
        # GPT2 BPE: " A" is single token 32 etc. — check both with and without space
        scores = {}
        for letter in CHOICES:
            for variant in [f" {letter}", letter, f" {letter}.", f"\n{letter}"]:
                tid = tok.encode(variant, add_special_tokens=False)
                if len(tid)==1:
                    scores[letter] = logits[tid[0]].item()
                    break
            if letter not in scores:
                scores[letter] = float("-inf")
        pred = max(scores, key=scores.get)
        if pred == true_letter:
            correct += 1
        total += 1
        overall_total += 1
        if pred == true_letter:
            overall_correct += 1

    acc = correct/total if total else 0
    results[subject] = {"correct": correct, "total": total, "accuracy": acc}
    print(f"[=] {subject}: {correct}/{total} = {acc*100:.1f}%  (running avg {(overall_correct/overall_total*100):.1f}%)")

avg = overall_correct/overall_total if overall_total else 0
print("\n" + "="*60)
print(f"MMLU RESULT — {MODEL_ID}")
print(f"Shots: {args.shots}  Subjects: {len(results)}/{len(MMLU_SUBJECTS)}")
for subj, r in sorted(results.items()):
    print(f"  {subj:35s} {r['accuracy']*100:5.1f}%  ({r['correct']}/{r['total']})")
print(f"\n  OVERALL: {overall_correct}/{overall_total} = {avg*100:.2f}%")
print("="*60)

out = {"model": MODEL_ID, "shots": args.shots, "limit": args.limit,
       "overall": {"correct": overall_correct, "total": overall_total, "accuracy": avg},
       "subjects": results}
Path(args.output).write_text(json.dumps(out, indent=2))
print(f"[+] Saved {args.output}")

# Also compare note
print("\nNote: GPT2-XL base ~24-26% MMLU (random 25%). MicroLLM2 distilled should be 25-30% —")
print("MMLU is knowledge-heavy; GPT2 1.5B 1024ctx cannot match 7B+ models. Use as sanity check, not SOTA claim.")