training/eval_sinhalammlu.py

#!/usr/bin/env python
"""
SinhalaMMLU evaluation with the ORGANISERS' own prompt (copied from their util_prompt.py).

Two subcommands:

  blocklist  - convert the official TEST json files into the mcq_dataset.py schema, so they can be
               used for DECONTAMINATION:
                   python eval_sinhalammlu.py blocklist --data /workspace/data/TEST/easy \
                          --out test_easy_block.jsonl
                   python mcq_dataset.py dedupe data/cleaned.jsonl data/deduped.jsonl \
                          --blocklist test_easy_block.jsonl --conflicts data/conflicts.jsonl
               Any non-zero "overlap with blocklist" means your training data contained test
               questions and they have now been removed. Re-run split/build-sft after this.

  eval       - score a model on those files. Reports BOTH:
                 loglik   = pick the option whose answer-token has the highest log-probability
                            (this is how the model was trained and how we measured validation)
                 generate = greedy-decode after "පිළිතුර:" and read the first digit
                            (this is what the organisers' evaluate.py does)
               If generate is much worse than loglik, the model knows the answer but is not
               emitting a clean digit, which would cost real marks.

                   python eval_sinhalammlu.py eval --data /workspace/data/TEST/easy \
                          --model /workspace/outputs/mcq_sft_base/best_adapter \
                          --out_csv results_easy.csv

NOTE ON THE RULES: the shared task forbids "training, tuning or prompt design on test questions in
any form". Use this as a one-off read-only sanity check. Do NOT use it to choose a checkpoint, a
prompt or hyperparameters: use your own held-out val.jsonl for that.
"""
import argparse
import csv
import glob
import json
import os
import re
import sys
import time
from collections import defaultdict

os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")

# --- the organisers' prompt text, verbatim from util_prompt.py -------------------------------
SINHALA_INSTRUCTIONS = {
    "easy": "පහත ප්‍රශ්නයට 1, 2, 3, 4 යන පිළිතුරුවලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරන්න.\n\n",
    "medium": "පහත ප්‍රශ්නයට 1, 2, 3, 4  යන පිළිතුරුවලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරාගන්න.\n\n",
    "hard": "පහත ප්‍රශ්නයට 1, 2, 3, 4, 5 යන පිළිතුරුවලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරාගන්න.\n\n",
}


def fit_numbers(instruction, n):
    """The organisers hard-code "1, 2, 3, 4, 5" in the HARD instruction (the original SinhalaMMLU
    hard split had five options). SinhalaMMLU 2.0 dropped one, so hard questions now have four;
    the number run is rewritten to match the actual option count."""
    return re.sub(r"1,\s*2,\s*3,\s*4(?:,\s*5)?",
                  ", ".join(str(i + 1) for i in range(n)), instruction, count=1)


def official_prompt(level, subject_original, question, choices, intro_type=True):
    """Zero-shot Sinhala prompt, identical to generate_instruction_prompt(..., language='sinhala')."""
    intro = f"මෙය {subject_original} විෂයයට අදාළ බහුවරණ ප්‍රශ්නයකි.\n" if intro_type else ""
    body = "\n".join(f"{i + 1}. {str(c).strip()}" for i, c in enumerate(choices))
    instr = fit_numbers(SINHALA_INSTRUCTIONS[level.lower()], len(choices))
    return intro + instr + f"ප්‍රශ්නය: {question}\n{body}\nපිළිතුර:"


def load_items(data_dir, level_override=""):
    """Read every *.json in data_dir (each a list of question dicts)."""
    files = sorted(glob.glob(os.path.join(data_dir, "**", "*.json"), recursive=True))
    if not files:
        sys.exit(f"no .json files found under {data_dir}")
    rows = []
    for path in files:
        try:
            data = json.load(open(path, encoding="utf-8"))
        except Exception as e:                                           # noqa: BLE001
            print(f"[warn] skipping {os.path.basename(path)}: {type(e).__name__}: {e}")
            continue
        if not isinstance(data, list):
            print(f"[warn] skipping {os.path.basename(path)}: not a JSON list")
            continue
        for it in data:
            md = it.get("metadata") or {}
            level = level_override or md.get("difficulty") or "easy"
            if level.lower() not in SINHALA_INSTRUCTIONS:
                level = "easy"
            ch = [c for c in (it.get("choices") or []) if str(c).strip()]
            ans = it.get("answer")
            if len(ch) < 2 or not isinstance(ans, int) or not (1 <= ans <= len(ch)):
                continue
            rows.append({"file": os.path.basename(path), "q_no": it.get("q_no"),
                         "subject": it.get("subject"), "category": it.get("category"),
                         "question": str(it.get("question") or "").strip(), "choices": ch,
                         "answer": ans, "level": level.lower(),
                         "subject_original": md.get("subject_original") or it.get("subject"),
                         "source": md.get("source")})
    print(f"[data] {len(rows)} questions from {len(files)} file(s) in {data_dir}")
    return rows


# ------------------------------------------------------------------ blocklist (no GPU needed)
def cmd_blocklist(a):
    rows = load_items(a.data, a.level)
    with open(a.out, "w", encoding="utf-8") as f:
        for i, r in enumerate(rows):
            f.write(json.dumps({"id": f"mmlu_{r['level']}_{r['file']}_{r['q_no']}",
                                "subject": r["subject"], "question": r["question"],
                                "choices": r["choices"], "answer": r["answer"],
                                "metadata": {"source": r["source"], "split": "official_test"}},
                               ensure_ascii=False) + "\n")
    print(f"[out] {len(rows)} questions -> {a.out}")
    print("Now decontaminate your TRAINING data (this is the compliant use of these files):")
    print(f"  python mcq_dataset.py dedupe data/cleaned.jsonl data/deduped.jsonl --blocklist {a.out}")
    print("Any non-zero 'overlap with blocklist' = test questions were in your training set.")


# ------------------------------------------------------------------ eval (needs the GPU)
def common_prefix_len(x, y):
    n = 0
    for p, q in zip(x, y):
        if p != q:
            break
        n += 1
    return n


def score_loglik(model, tok, rows, prompts, max_len, batch_seqs):
    """Total log-prob of each candidate answer token; returns a predicted 1-based option per row."""
    import torch
    dev = next(model.parameters()).device
    pad = tok.pad_token_id if tok.pad_token_id is not None else 0
    items, skipped = [], set()
    for ri, (r, p) in enumerate(zip(rows, prompts)):
        p_ids = tok(p)["input_ids"]
        seqs = []
        for ci in range(len(r["choices"])):
            f = tok(p + f" {ci + 1}")["input_ids"]
            seqs.append((f, max(1, common_prefix_len(p_ids, f))))
        if max(len(s) for s, _ in seqs) > max_len:
            skipped.add(ri)
            continue
        for ci, (s, k) in enumerate(seqs):
            items.append((ri, ci, s, k))

    scores = {}
    with torch.no_grad():
        for i in range(0, len(items), batch_seqs):
            chunk = items[i:i + batch_seqs]
            L = max(len(x[2]) for x in chunk)
            ids = torch.full((len(chunk), L), pad, dtype=torch.long)
            mask = torch.zeros((len(chunk), L), dtype=torch.long)
            for b, (_, _, s, _) in enumerate(chunk):
                ids[b, :len(s)] = torch.tensor(s)
                mask[b, :len(s)] = 1
            logits = model(input_ids=ids.to(dev), attention_mask=mask.to(dev)).logits
            for b, (ri, ci, s, k) in enumerate(chunk):
                lp = logits[b, k - 1:len(s) - 1].float().log_softmax(-1)
                tgt = ids[b, k:len(s)].to(dev)
                scores[(ri, ci)] = lp.gather(-1, tgt[:, None]).sum().item()
            del logits

    preds = []
    for ri, r in enumerate(rows):
        if ri in skipped:
            preds.append(None)
            continue
        cs = [scores.get((ri, ci)) for ci in range(len(r["choices"]))]
        preds.append(None if any(c is None for c in cs)
                     else 1 + max(range(len(cs)), key=lambda i: cs[i]))
    return preds, len(skipped)


def score_generate(model, tok, prompts, max_new_tokens, batch_rows):
    """Greedy-decode and read the first digit, the way the organisers' evaluate.py does."""
    import torch
    dev = next(model.parameters()).device
    pad = tok.pad_token_id if tok.pad_token_id is not None else 0
    preds, raw = [], []
    with torch.no_grad():
        for i in range(0, len(prompts), batch_rows):
            batch = prompts[i:i + batch_rows]
            enc = [tok(p)["input_ids"] for p in batch]
            L = max(len(e) for e in enc)
            ids = torch.full((len(enc), L), pad, dtype=torch.long)
            mask = torch.zeros((len(enc), L), dtype=torch.long)
            for b, e in enumerate(enc):                      # LEFT pad so generation continues correctly
                ids[b, L - len(e):] = torch.tensor(e)
                mask[b, L - len(e):] = 1
            out = model.generate(input_ids=ids.to(dev), attention_mask=mask.to(dev),
                                 max_new_tokens=max_new_tokens, do_sample=False,
                                 pad_token_id=pad)
            for b in range(len(enc)):
                text = tok.decode(out[b, L:], skip_special_tokens=True)
                raw.append(text)
                m = re.search(r"[1-5]", text)
                preds.append(int(m.group()) if m else None)
    return preds, raw


def score_official(model, tok, prompts, n_choices, system="You must output the Answer",
                   batch_rows=8, no_chat=False):
    """Exact replica of the organisers' predict_classification_causal_by_number (utils.py):
    wrap in the chat template with a system message, ONE forward pass, take the logits at the final
    position, and argmax over the token ids of the BARE digits '1'..'N' (no leading space).
    This is what actually decides their score - there is no generation involved."""
    import torch
    dev = next(model.parameters()).device
    pad = tok.pad_token_id if tok.pad_token_id is not None else 0

    mode = {"how": "raw prompt (no chat template requested)"} if no_chat else {}

    def wrap(instruction):
        if no_chat:
            return instruction
        msgs = [{"role": "system", "content": system}, {"role": "user", "content": instruction}]
        try:
            out = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
            mode.setdefault("how", "chat template, system + user roles")
            return out
        except Exception:
            pass
        try:
            # many Gemma templates reject a system role: fold it into the first user turn
            msgs = [{"role": "user", "content": system + "\n\n" + instruction}]
            out = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
            mode.setdefault("how", "chat template, system folded into the user turn")
            return out
        except Exception:
            # NO chat template on this tokenizer at all (base Gemma 4 has none).
            # The organisers' evaluate.py calls apply_chat_template unconditionally, so it would
            # CRASH here too -> scoring the raw prompt is the closest thing we can measure.
            mode.setdefault("how", "NO CHAT TEMPLATE on this tokenizer: scored the raw prompt. "
                                   "The organisers' evaluate.py would fail on this model as-is.")
            return instruction

    texts = [wrap(p) for p in prompts]
    # digit token ids, exactly as they compute them
    dig = {k: tok.encode(str(k), add_special_tokens=False)[-1] for k in range(1, 6)}
    preds = []
    with torch.no_grad():
        for i in range(0, len(texts), batch_rows):
            batch = texts[i:i + batch_rows]
            enc = [tok(t)["input_ids"] for t in batch]          # add_special_tokens default, as they do
            L = max(len(e) for e in enc)
            ids = torch.full((len(enc), L), pad, dtype=torch.long)
            mask = torch.zeros((len(enc), L), dtype=torch.long)
            for b, e in enumerate(enc):                          # LEFT pad: last column is the last real token
                ids[b, L - len(e):] = torch.tensor(e)
                mask[b, L - len(e):] = 1
            logits = model(input_ids=ids.to(dev), attention_mask=mask.to(dev)).logits[:, -1, :].float()
            for b in range(len(enc)):
                nc = n_choices[i + b]
                cand = [dig[k] for k in range(1, nc + 1)]
                row = logits[b, torch.tensor(cand, device=logits.device)]
                preds.append(1 + int(row.argmax().item()))
            del logits
    print(f"[official] prompt construction: {mode.get('how', 'unknown')}")
    return preds, texts[0]


def report(rows, preds, label):
    n = correct = unparsed = 0
    by_cat, by_subj = defaultdict(lambda: [0, 0]), defaultdict(lambda: [0, 0])
    for r, p in zip(rows, preds):
        if p is None:
            unparsed += 1
            continue
        n += 1
        ok = int(p == r["answer"])
        correct += ok
        by_cat[r["category"] or "?"][0] += ok
        by_cat[r["category"] or "?"][1] += 1
        by_subj[r["subject"] or "?"][0] += ok
        by_subj[r["subject"] or "?"][1] += 1
    print(f"\n=== {label} ===")
    print(f"accuracy {correct / max(1, n):.4f} ({correct}/{n}); no prediction: {unparsed}")
    print("  by category:")
    for k, (c, t) in sorted(by_cat.items()):
        print(f"    {k:<18} {c / t:.4f}  ({c}/{t})")
    print("  by subject:")
    for k, (c, t) in sorted(by_subj.items()):
        if t >= 20:
            print(f"    {k:<18} {c / t:.4f}  ({c}/{t})")
    return {"accuracy": correct / max(1, n), "n": n, "no_prediction": unparsed,
            "by_category": {k: v[0] / v[1] for k, v in by_cat.items()},
            "by_subject": {k: v[0] / v[1] for k, v in by_subj.items()}}


def cmd_eval(a):
    rows = load_items(a.data, a.level)
    if a.limit:
        rows = rows[:a.limit]
    prompts = [official_prompt(r["level"], r["subject_original"], r["question"], r["choices"],
                               intro_type=not a.no_intro) for r in rows]
    print("\n--- prompt actually sent to the model (first question) ---")
    print(prompts[0])
    print("--- end of prompt ---\n")
    if a.dry_run:
        return

    from unsloth import FastModel
    print(f"[model] loading {a.model}")
    model, processor = FastModel.from_pretrained(model_name=a.model, dtype=None,
                                                 max_seq_length=a.max_seq_length,
                                                 load_in_4bit=False, full_finetuning=False)
    tok = getattr(processor, "tokenizer", processor)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    model.eval()

    out = {"model": a.model, "data": a.data, "n_questions": len(rows)}
    t0 = time.time()
    ll_preds, skipped = score_loglik(model, tok, rows, prompts, a.max_seq_length, a.eval_batch)
    print(f"[loglik] {time.time() - t0:.0f}s; skipped (too long) {skipped}")
    out["loglik"] = report(rows, ll_preds, "loglik (how the model was trained/validated)")

    if not a.skip_official:
        t0 = time.time()
        off_preds, shown = score_official(model, tok, prompts, [len(r["choices"]) for r in rows],
                                          system=a.system, batch_rows=a.gen_batch, no_chat=a.no_chat)
        print(f"[official] {time.time() - t0:.0f}s")
        print("--- prompt as the organisers' harness builds it (chat template applied) ---")
        print(shown[:900] + ("..." if len(shown) > 900 else ""))
        print("--- end ---")
        out["official"] = report(rows, off_preds, "OFFICIAL (bare-digit logits at last position)")

    gen_preds, gen_raw = [None] * len(rows), [""] * len(rows)
    if not a.skip_generate:
        t0 = time.time()
        gen_preds, gen_raw = score_generate(model, tok, prompts, a.max_new_tokens, a.gen_batch)
        print(f"[generate] {time.time() - t0:.0f}s")
        out["generate"] = report(rows, gen_preds, "generate (how the organisers score)")
        bad = [t for t, p in zip(gen_raw, gen_preds) if p is None][:5]
        if bad:
            print(f"  examples with no digit in the output: {bad}")

    if a.out_csv:
        # columns match the organisers' get_accuracy_all.py so their aggregator can read this
        with open(a.out_csv, "w", encoding="utf-8", newline="") as f:
            w = csv.writer(f)
            w.writerow(["q_no", "subject", "category", "question", "choices", "answer",
                        "generated_text", "model_predicted_answer", "loglik_pred", "raw_output"])
            for r, lp, gp, g in zip(rows, ll_preds, gen_preds, gen_raw):
                w.writerow([r["q_no"], r["subject"], r["category"], r["question"],
                            json.dumps(r["choices"], ensure_ascii=False), r["answer"],
                            gp if gp is not None else "", gp if gp is not None else "", lp, g])
        print(f"[out] per-question results -> {a.out_csv}")
    if a.out_json:
        json.dump(out, open(a.out_json, "w", encoding="utf-8"), ensure_ascii=False, indent=1)
        print(f"[out] summary -> {a.out_json}")
    print("\nREMINDER: these are official test questions. Do not use this score to choose a "
          "checkpoint, prompt or hyperparameter; the task rules forbid tuning on test questions.")


def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("blocklist", help="convert official TEST json -> jsonl for decontamination")
    p.add_argument("--data", required=True)
    p.add_argument("--out", required=True)
    p.add_argument("--level", default="")
    p.set_defaults(f=cmd_blocklist)

    p = sub.add_parser("eval", help="score a model with the organisers' prompt")
    p.add_argument("--data", required=True)
    p.add_argument("--model", required=True, help="adapter dir, merged dir or HF repo")
    p.add_argument("--level", default="", help="force easy/medium/hard (default: read from metadata)")
    p.add_argument("--no_intro", action="store_true", help="drop the subject intro line (intro_type=False)")
    p.add_argument("--max_seq_length", type=int, default=1024)
    p.add_argument("--eval_batch", type=int, default=16)
    p.add_argument("--gen_batch", type=int, default=8)
    p.add_argument("--max_new_tokens", type=int, default=5)
    p.add_argument("--skip_generate", action="store_true")
    p.add_argument("--skip_official", action="store_true")
    p.add_argument("--no_chat", action="store_true", help="official scoring WITHOUT the chat template")
    p.add_argument("--system", default="You must output the Answer",
                   help="system message the organisers' evaluate.py uses")
    p.add_argument("--limit", type=int, default=0)
    p.add_argument("--dry_run", action="store_true", help="print the prompt and exit (no model load)")
    p.add_argument("--out_csv", default="")
    p.add_argument("--out_json", default="")
    p.set_defaults(f=cmd_eval)

    a = ap.parse_args()
    for stream in (sys.stdout, sys.stderr):
        try:
            stream.reconfigure(encoding="utf-8")
        except Exception:                                                # noqa: BLE001
            pass
    a.f(a)


if __name__ == "__main__":
    main()