training/make_predictions.py

#!/usr/bin/env python
"""
Produce the shared-task submission: predictions.zip containing predictions.jsonl.

Required line format (one JSON object per question):
    {"q_no":1,"subject":"Arts","category":"humanities","level":"easy","answer":"2"}

`answer` is a STRING holding the 1-based option number.

    python make_predictions.py \
        --data /workspace/data/TEST \
        --model /workspace/outputs/mcq_sft_base_v2/best_adapter \
        --out predictions.jsonl --zip predictions.zip

--data may be either a single folder of *.json files, or a root containing easy/ medium/ hard/
subfolders; the level comes from metadata.difficulty, falling back to the folder name.

Scoring: for every option, score the joint tokenisation of prompt+" N" and take the highest total
log-probability ("loglik"). This is deterministic, always yields a valid 1..N answer, and measured
highest on TEST/easy. Because the task is submitted as PREDICTIONS, we are free to choose the
decoding method - we are not restricted to the organisers' bare-digit logit comparison.

Reproducibility: the organisers re-run submitted code on their own hardware and only verified runs
enter the final ranking. Nothing here samples; seeds are fixed and generation is not used. Keep this
file, the adapter/model reference and the prompt builder together in the submitted code.
"""
import argparse
import glob
import json
import os
import re
import random
import sys
import time
import zipfile

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

# --- the organisers' prompt, verbatim from their 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",
}
ENGLISH_INSTRUCTIONS = {
    "easy": "Choose the correct or most appropriate answer from 1, 2, 3, 4 for the question below.\n\n",
    "medium": "Choose the correct or most appropriate answer from 1, 2, 3, 4 for the question below.\n\n",
    "hard": "Select the correct answer from 1, 2, 3, 4, 5 for the given question.\n\n",
}
LEVELS = ("easy", "medium", "hard")

# Prompt variants to ensemble over. Each is a real configuration the organisers' harness supports
# (sinhala/english x intro true/false), plus option rotations to cancel answer-position bias.
# "rot" rotates the DISPLAYED option order; predictions are mapped back to the original indices.
VARIANTS = {
    "si_intro":    {"lang": "si", "intro": True,  "rot": 0},
    "si_nointro":  {"lang": "si", "intro": False, "rot": 0},
    "en_intro":    {"lang": "en", "intro": True,  "rot": 0},
    "si_rot1":     {"lang": "si", "intro": True,  "rot": 1},
    "si_rot2":     {"lang": "si", "intro": True,  "rot": 2},
    "si_rot3":     {"lang": "si", "intro": True,  "rot": 3},
    "en_nointro":  {"lang": "en", "intro": False, "rot": 0},
    "en_rot1":     {"lang": "en", "intro": True,  "rot": 1},
    "si_nointro_r2": {"lang": "si", "intro": False, "rot": 2},
}
DEFAULT_ENSEMBLE = ["si_intro", "si_nointro", "en_intro", "si_rot1", "si_rot2"]


def fit_numbers(instruction, n):
    """The organisers hard-code "1, 2, 3, 4, 5" in the HARD instruction because the original
    SinhalaMMLU hard split had five options. SinhalaMMLU 2.0 dropped one incorrect option, so hard
    questions now have four. Telling the model to choose from five options when four are shown is
    strictly worse, so the number run is rewritten to match the actual 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):
    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], len(choices))
    return intro + instr + f"ප්‍රශ්නය: {question}\n{body}\nපිළිතුර:"


def english_prompt(level, subject, question, choices, intro_type=True):
    """The organisers' prompt with language="english" (their english_intro uses the ENGLISH subject)."""
    intro = (f"This is a multiple choice question related to the subject {subject}. "
             "Given in Sinhala Language.\n") if intro_type else ""
    body = "\n".join(f"{i + 1}. {str(c).strip()}" for i, c in enumerate(choices))
    instr = fit_numbers(ENGLISH_INSTRUCTIONS[level], len(choices))
    return intro + instr + f"Question: {question}\n{body}\nAnswer:"


def rotate(choices, rot):
    """Displayed order; displayed position j holds original option (j + rot) % n."""
    n = len(choices)
    r = rot % n if n else 0
    return choices[r:] + choices[:r], r


def variant_prompt(row, spec):
    shown, r = rotate(row["choices"], spec["rot"])
    if spec["lang"] == "en":
        p = english_prompt(row["level"], row["subject"], row["question"], shown, spec["intro"])
    else:
        p = official_prompt(row["level"], row["subject_original"], row["question"], shown, spec["intro"])
    return p, r, len(shown)


def content_free_prompt(row, spec, filler="N/A"):
    """Same template and option count, but the question and every option replaced by a neutral
    filler. Scoring the digits here measures the model's a-priori preference among option POSITIONS,
    independent of content (contextual calibration / PMI-DC). Depends only on
    (variant, level, subject, n_choices), so it is cached across questions."""
    stub = dict(row)
    stub["question"] = filler
    stub["choices"] = [filler] * len(row["choices"])
    return variant_prompt(stub, spec)


def level_of(item, path, override=""):
    """metadata.difficulty, else a top-level 'level'/'difficulty' field, else the file name or any
    path component containing easy/medium/hard, else 'easy'."""
    if override:
        return override
    md = item.get("metadata") or {}
    for cand in (md.get("difficulty"), md.get("level"), item.get("level"), item.get("difficulty")):
        d = str(cand or "").strip().lower()
        if d in LEVELS:
            return d
    low = path.replace("\\", "/").lower()
    # substring match so "hard.json", "test_hard.jsonl" and ".../hard/..." all work
    hits = [lv for lv in LEVELS if lv in low]
    if len(hits) == 1:
        return hits[0]
    if hits:                                  # e.g. ".../easy/medium_extra.json": trust the last one
        return max(hits, key=low.rfind)
    return "easy"


def read_any(path):
    """A JSON array, a JSONL file, or a dict wrapping the list under a common key."""
    raw = open(path, encoding="utf-8-sig").read().strip()
    if not raw:
        return []
    if raw[0] == "[":
        return json.loads(raw)
    if raw[0] == "{":
        try:
            obj = json.loads(raw)
        except json.JSONDecodeError:
            obj = None
        if isinstance(obj, dict):
            for k in ("data", "questions", "items", "rows", "test"):
                if isinstance(obj.get(k), list):
                    return obj[k]
            return [obj]
        # not one object -> treat as JSONL
    out = []
    for line in raw.splitlines():
        line = line.strip()
        if line:
            out.append(json.loads(line))
    return out


def load_items(data_dir, level_override=""):
    if os.path.isfile(data_dir):
        files = [data_dir]
    else:
        files = sorted(set(glob.glob(os.path.join(data_dir, "**", "*.json"), recursive=True))
                       | set(glob.glob(os.path.join(data_dir, "**", "*.jsonl"), recursive=True)))
    if not files:
        sys.exit(f"no .json/.jsonl files found under {data_dir}")
    rows, skipped, per_file = [], 0, []
    for path in files:
        try:
            data = read_any(path)
        except Exception as e:                                           # noqa: BLE001
            print(f"[warn] {os.path.basename(path)}: {type(e).__name__}: {e}")
            continue
        if not isinstance(data, list) or not data:
            print(f"[warn] {os.path.basename(path)}: no question list found")
            continue
        before = len(rows)
        for it in data:
            md = it.get("metadata") or {}
            ch = [c for c in (it.get("choices") or []) if str(c).strip()]
            if len(ch) < 2:
                skipped += 1          # never drop silently: these still need a prediction line
            ans = it.get("answer")
            rows.append({
                "file": os.path.basename(path),
                "gold": ans if isinstance(ans, int) and 1 <= ans <= max(2, len(ch)) else None,
                "q_no": it.get("q_no"),
                "subject": it.get("subject"),
                "category": it.get("category"),
                "level": level_of(it, path, level_override),
                "question": str(it.get("question") or "").strip(),
                "choices": ch,
                "subject_original": md.get("subject_original") or it.get("subject"),
            })
        lv = {r["level"] for r in rows[before:]}
        per_file.append((os.path.basename(path), len(rows) - before, sorted(lv)))
    by_level = {}
    for r in rows:
        by_level[r["level"]] = by_level.get(r["level"], 0) + 1
    print(f"[data] {len(rows)} questions from {len(files)} file(s); by level: {by_level}")
    for name, n, lv in per_file:
        print(f"    {name:<44} {n:>6} questions   level -> {','.join(lv)}")
    missing = [lv for lv in LEVELS if lv not in by_level]
    if missing:
        print(f"[warn] no questions detected for: {missing}. If that is wrong, the level could not be "
              "read from metadata.difficulty or the file name - check the table above.")
    if skipped:
        print(f"[warn] {skipped} question(s) have fewer than 2 usable options; they still get a line")
    return rows


def common_prefix_len(a, b):
    n = 0
    for x, y in zip(a, b):
        if x != y:
            break
        n += 1
    return n


def predict(model, tok, rows, max_len, batch_seqs, variants, calibrate=False):
    """Deterministic prompt ensemble. For every (question, variant) the options are scored, converted
    to a log-softmax over the DISPLAYED positions (so variants with different prompt lengths are on a
    common scale), mapped back to the original option order, and summed across variants.
    No sampling, no generation: the same inputs always give the same predictions."""
    import math
    import torch
    model.eval()
    dev = next(model.parameters()).device
    pad = tok.pad_token_id if tok.pad_token_id is not None else 0

    items, meta = [], {}                 # (ri, vi, shown_pos, ids, answer_start) ; (ri,vi)->(rot,n)
    for ri, r in enumerate(rows):
        if len(r["choices"]) < 2:
            continue
        for vi, vname in enumerate(variants):
            prompt, rot, n = variant_prompt(r, VARIANTS[vname])
            meta[(ri, vi)] = (rot, n)
            p_ids = tok(prompt)["input_ids"]
            for j in range(n):
                f = tok(prompt + f" {j + 1}")["input_ids"]
                if len(f) > max_len:
                    f = f[-max_len:]
                items.append((ri, vi, j, f, max(1, common_prefix_len(p_ids, f))))

    # ---- calibration priors: one short prompt per (variant, level, subject, n_choices)
    prior_items, prior_key = [], {}
    if calibrate:
        for ri, r in enumerate(rows):
            if len(r["choices"]) < 2:
                continue
            for vi, vname in enumerate(variants):
                key = (vi, r["level"], r["subject_original"], len(r["choices"]))
                prior_key[(ri, vi)] = key
                if key in [k for _, k, _, _, _ in prior_items]:
                    continue
                cf, rot, n = content_free_prompt(r, VARIANTS[vname])
                p_ids = tok(cf)["input_ids"]
                for j in range(n):
                    f = tok(cf + f" {j + 1}")["input_ids"]
                    prior_items.append((len(prior_items), key, j, f, max(1, common_prefix_len(p_ids, f))))
        # de-duplicate: keep the first occurrence of each (key, j)
        seen, uniq = set(), []
        for _, key, j, f, k in prior_items:
            if (key, j) in seen:
                continue
            seen.add((key, j))
            uniq.append((key, j, f, k))
        prior_items = uniq
        print(f"[calibrate] {len(prior_items)} content-free candidates "
              f"({len({k for k, _, _, _ in prior_items})} distinct prompts)")

    print(f"[predict] {len(rows)} questions x {len(variants)} variant(s) = {len(items)} candidates")
    scores, t0 = {}, time.time()
    with torch.no_grad():
        for i in range(0, len(items), batch_seqs):
            chunk = items[i:i + batch_seqs]
            L = max(len(x[3]) 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, (_, _, _, sq, _) in enumerate(chunk):
                ids[b, :len(sq)] = torch.tensor(sq)
                mask[b, :len(sq)] = 1
            logits = model(input_ids=ids.to(dev), attention_mask=mask.to(dev)).logits
            for b, (ri, vi, j, sq, k) in enumerate(chunk):
                lp = logits[b, k - 1:len(sq) - 1].float().log_softmax(-1)
                tgt = ids[b, k:len(sq)].to(dev)
                scores[(ri, vi, j)] = lp.gather(-1, tgt[:, None]).sum().item()
            del logits
            if (i // max(1, batch_seqs)) % 100 == 0:
                done = min(i + batch_seqs, len(items))
                print(f"  scored {done}/{len(items)} ({time.time() - t0:.0f}s)", flush=True)

    priors = {}
    if calibrate and prior_items:
        with torch.no_grad():
            for i in range(0, len(prior_items), batch_seqs):
                chunk = prior_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, (_, _, sq, _) in enumerate(chunk):
                    ids[b, :len(sq)] = torch.tensor(sq)
                    mask[b, :len(sq)] = 1
                logits = model(input_ids=ids.to(dev), attention_mask=mask.to(dev)).logits
                for b, (key, j, sq, k) in enumerate(chunk):
                    lp = logits[b, k - 1:len(sq) - 1].float().log_softmax(-1)
                    tgt = ids[b, k:len(sq)].to(dev)
                    priors[(key, j)] = lp.gather(-1, tgt[:, None]).sum().item()
                del logits
        print(f"[calibrate] priors computed in {time.time() - t0:.0f}s")

    preds, fallback, agree = [], 0, 0
    for ri, r in enumerate(rows):
        n = len(r["choices"])
        if n < 2:
            preds.append(1)
            fallback += 1
            continue
        total, used, per_variant = [0.0] * n, 0, []
        for vi in range(len(variants)):
            if (ri, vi) not in meta:
                continue
            rot, nv = meta[(ri, vi)]
            raw = [scores.get((ri, vi, j)) for j in range(nv)]
            if any(x is None for x in raw):
                continue
            m = max(raw)                                     # log-softmax over displayed positions
            z = m + math.log(sum(math.exp(x - m) for x in raw))
            adj = [raw[j] - z for j in range(nv)]
            if calibrate:
                key = prior_key.get((ri, vi))
                pr = [priors.get((key, j)) for j in range(nv)]
                if all(x is not None for x in pr):
                    pm = max(pr)
                    pz = pm + math.log(sum(math.exp(x - pm) for x in pr))
                    adj = [adj[j] - (pr[j] - pz) for j in range(nv)]   # subtract the position prior
            best_j = max(range(nv), key=lambda j: adj[j])
            per_variant.append(1 + (best_j + rot) % n)
            for j in range(nv):
                total[(j + rot) % n] += adj[j]                 # displayed j == original (j+rot)%n
            used += 1
        if used == 0:
            preds.append(1)
            fallback += 1
            continue
        preds.append(1 + max(range(n), key=lambda i: total[i]))
        if len(set(per_variant)) == 1:
            agree += 1

    if fallback:
        print(f"[warn] {fallback} question(s) fell back to answer 1")
    if len(variants) > 1:
        print(f"[predict] all variants agreed on {agree}/{len(rows)} ({agree / max(1, len(rows)):.1%}) "
              "- the rest are where the ensemble changes the answer")
    print(f"[predict] done in {time.time() - t0:.0f}s")
    return preds, None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True, help="folder of test *.json files, or a root with easy/medium/hard")
    ap.add_argument("--model", required=True, help="adapter dir, merged dir, or HF repo")
    ap.add_argument("--out", default="predictions.jsonl")
    ap.add_argument("--zip", default="predictions.zip", help="'' to skip zipping")
    ap.add_argument("--level", default="", choices=["", *LEVELS], help="force the level for every item")
    ap.add_argument("--max_seq_length", type=int, default=1024)
    ap.add_argument("--eval_batch", type=int, default=16)
    ap.add_argument("--calibrate", action="store_true",
                    help="contextual calibration: measure the model's a-priori preference among "
                         "option POSITIONS on a content-free prompt and subtract it. ~50 extra "
                         "forward passes in total; standard fix for log-likelihood MCQ position bias.")
    ap.add_argument("--score", action="store_true",
                    help="if the input carries gold 'answer' values (e.g. your val3.jsonl), report "
                         "accuracy as well. Use this to compare variant sets on HELD-OUT data "
                         "instead of spending Codabench submissions.")
    ap.add_argument("--ensemble", nargs="*", default=None,
                    help=f"prompt variants to average over (default: {' '.join(DEFAULT_ENSEMBLE)}). "
                         f"Choices: {' '.join(VARIANTS)}. Pass 'si_intro' alone for single-prompt scoring.")
    ap.add_argument("--seed", type=int, default=3407)
    ap.add_argument("--dry_run", action="store_true", help="show the prompt and exit, no model load")
    args = ap.parse_args()
    for stream in (sys.stdout, sys.stderr):
        try:
            stream.reconfigure(encoding="utf-8")
        except Exception:                                                # noqa: BLE001
            pass

    random.seed(args.seed)
    variants = args.ensemble if args.ensemble else list(DEFAULT_ENSEMBLE)
    unknown = [v for v in variants if v not in VARIANTS]
    if unknown:
        sys.exit(f"unknown variant(s) {unknown}; choose from {list(VARIANTS)}")
    print(f"[ensemble] {len(variants)} variant(s): {variants}")
    rows = load_items(args.data, args.level)

    if args.dry_run:
        r = next((x for x in rows if len(x["choices"]) >= 2), None)
        if r:
            for v in variants:
                prompt, rot, _ = variant_prompt(r, VARIANTS[v])
                print(f"\n--- variant {v} (rot={rot}) ---")
                print(prompt)
            print("--- end ---")
        return

    import torch
    torch.manual_seed(args.seed)
    from unsloth import FastModel
    print(f"[model] loading {args.model}")
    model, processor = FastModel.from_pretrained(model_name=args.model, dtype=None,
                                                 max_seq_length=args.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

    preds, _ = predict(model, tok, rows, args.max_seq_length, args.eval_batch, variants,
                       calibrate=args.calibrate)

    with open(args.out, "w", encoding="utf-8") as f:
        for r, p in zip(rows, preds):
            f.write(json.dumps({"q_no": r["q_no"], "subject": r["subject"],
                                "category": r["category"], "level": r["level"],
                                "answer": str(p)}, ensure_ascii=False) + "\n")
    print(f"[out] {len(rows)} lines -> {args.out}")

    # self-checks: the submission is rejected or mis-scored if any of these are wrong
    lines = [json.loads(l) for l in open(args.out, encoding="utf-8") if l.strip()]
    assert len(lines) == len(rows), "line count does not match the question count"
    bad_keys = [i for i, d in enumerate(lines)
                if set(d) != {"q_no", "subject", "category", "level", "answer"}]
    bad_ans = [i for i, d in enumerate(lines)
               if not (isinstance(d["answer"], str) and d["answer"] in {"1", "2", "3", "4", "5"})]
    bad_lvl = [i for i, d in enumerate(lines) if d["level"] not in LEVELS]
    for name, bad in [("wrong key set", bad_keys), ("answer not a '1'..'5' string", bad_ans),
                      ("level not easy/medium/hard", bad_lvl)]:
        print(f"[check] {name}: {len(bad)}" + (f"  e.g. line {bad[0] + 1}" if bad else "  OK"))
    dist = {}
    for d in lines:
        dist[d["answer"]] = dist.get(d["answer"], 0) + 1
    print(f"[check] answer distribution: {dict(sorted(dist.items()))}")
    dup = len(lines) - len({(d["level"], d["subject"], d["q_no"]) for d in lines})
    print(f"[check] duplicate (level, subject, q_no) keys: {dup}")
    if bad_keys or bad_ans or bad_lvl:
        sys.exit("fix the problems above before submitting")

    if args.score:
        import collections
        have = [(r, p) for r, p in zip(rows, preds) if r["gold"] is not None]
        if not have:
            print("[score] the input has no gold 'answer' field, nothing to score")
        else:
            n = sum(1 for _ in have)
            c = sum(int(p == r["gold"]) for r, p in have)
            print(f"\n[score] accuracy {c / n:.4f}  ({c}/{n})   variants: {variants}")
            for key, label in (("level", "level"), ("subject", "subject")):
                agg = collections.defaultdict(lambda: [0, 0])
                for r, p in have:
                    k = r.get(key) or "?"
                    agg[k][0] += int(p == r["gold"])
                    agg[k][1] += 1
                print(f"  by {label}:")
                for k, (cc, tt) in sorted(agg.items(), key=lambda x: -x[1][1]):
                    if tt >= 20 or label == "level":
                        print(f"    {str(k):<34} {cc / tt:.4f}  ({cc}/{tt})")

    if args.zip:
        with zipfile.ZipFile(args.zip, "w", zipfile.ZIP_DEFLATED) as z:
            z.write(args.out, arcname="predictions.jsonl")   # must sit at the archive root
        with zipfile.ZipFile(args.zip) as z:
            print(f"[zip] {args.zip} contains: {z.namelist()}")
    print("\nFirst two lines:")
    for l in open(args.out, encoding="utf-8").readlines()[:2]:
        print("  " + l.rstrip())


if __name__ == "__main__":
    main()