#!/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.")