feat: add mmlu_bench.py (5-shot MMLU, harness or lightweight)

mmlu_bench.py

@@ -0,0 +1,235 @@
+#!/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.")