inference.py

"""
Sinhala MMLU Shared Task - Final inference script
Team: QBitsLK

The organizers will run this file exactly as it is, changing ONLY the DATA_REPO line
below to point to the final Sinhala MMLU dataset:

    pip install -r requirements.txt
    python inference.py

System: one merged 16-bit model (7,996,157,418 parameters, <= 8B). Gemma 4 E4B with Sinhala
continued pretraining, LoRA-fine-tuned on past-paper MCQs and merged back into the base weights.

Answers are chosen by per-option log-likelihood under a 5-prompt ensemble. No text is generated,
nothing is sampled, and the seed is fixed, so a re-run gives identical answers.
"""

# ======================= ORGANIZERS CHANGE ONLY THIS LINE =======================
DATA_REPO = "SinhalaMMLU-SharedTask/Sinhala-MMLU-SharedTask-Test-Dataset"
# ================================================================================

OUTPUT_FILE = "predictions.jsonl"

# ----------------------- TEAM SETTINGS (keep exactly as used for your final submission)
MODEL_ID = "SinhalaMMLU-SharedTask/team-qbits-cpt-v3"   # merged model, self-contained
SEED = 42
BATCH = 16                 # candidate sequences scored per forward pass
MAX_SEQ_LEN = 1024

# Prompt variants the ensemble averages over. Each is a configuration the organizers' own
# reference harness supports (sinhala/english x subject-intro on/off), plus two rotations of the
# option order to cancel answer-position bias. CHANGING THIS LIST CHANGES THE PREDICTIONS.
VARIANTS = [
    {"lang": "si", "intro": True,  "rot": 0},
    {"lang": "si", "intro": False, "rot": 0},
    {"lang": "en", "intro": True,  "rot": 0},
    {"lang": "si", "intro": True,  "rot": 1},
    {"lang": "si", "intro": True,  "rot": 2},
]

import json
import math
import os
import random
import re

import numpy as np
import torch
from huggingface_hub import HfFileSystem, list_repo_files

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


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


# =============== 1. LOAD THE DATASET FROM HUGGING FACE (do not change) ===============
def load_questions(data_repo):
    fs = HfFileSystem()
    files = sorted(f for f in list_repo_files(data_repo, repo_type="dataset") if f.endswith(".json"))
    rows = []
    for f in files:
        with fs.open(f"datasets/{data_repo}/{f}", "r", encoding="utf-8") as fh:
            raw = json.load(fh)
        if isinstance(raw, dict):
            raw = next(v for v in raw.values() if isinstance(v, list))
        for r in raw:
            r["_file"] = f
            rows.append(r)
    return rows


def get_level(r):
    return (r.get("metadata") or {}).get("difficulty") or r["_file"].split("/")[0]


# =============== 2. OUR SYSTEM ===============
# The two function names are kept. predict_all() returns one answer ("1".."4") per question,
# in the same order as `rows`.

# Instruction text taken verbatim from the organizers' 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",
}


def fit_numbers(instruction, n):
    """The reference prompt hard-codes "1, 2, 3, 4, 5" for the HARD level because the original
    SinhalaMMLU hard split had five options. SinhalaMMLU 2.0 drops one incorrect option, so hard
    questions have four. Telling the model to choose from five when four are shown is strictly
    worse, so 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 build_prompt(level, subject, subject_original, question, choices, spec):
    """Returns (prompt, rot). Displayed position j holds original option (j + rot) % n."""
    n = len(choices)
    rot = spec["rot"] % n if n else 0
    shown = choices[rot:] + choices[:rot]
    body = "\n".join(f"{i + 1}. {str(c).strip()}" for i, c in enumerate(shown))
    if spec["lang"] == "en":
        intro = (f"This is a multiple choice question related to the subject {subject}. "
                 "Given in Sinhala Language.\n") if spec["intro"] else ""
        instr = fit_numbers(ENGLISH_INSTRUCTIONS[level], n)
        return intro + instr + f"Question: {question}\n{body}\nAnswer:", rot
    intro = f"මෙය {subject_original} විෂයයට අදාළ බහුවරණ ප්‍රශ්නයකි.\n" if spec["intro"] else ""
    instr = fit_numbers(SINHALA_INSTRUCTIONS[level], n)
    return intro + instr + f"ප්‍රශ්නය: {question}\n{body}\nපිළිතුර:", rot


def load_system():
    """Loads the merged model. Gemma 4 is a multimodal architecture, so the right AutoClass
    depends on the transformers version; we try the plausible ones in order."""
    try:
        from unsloth import FastModel
        model, proc = FastModel.from_pretrained(model_name=MODEL_ID, dtype=None,
                                                max_seq_length=MAX_SEQ_LEN,
                                                load_in_4bit=False, full_finetuning=False)
        tok = getattr(proc, "tokenizer", proc)
        if getattr(tok, "pad_token", None) is None:
            tok.pad_token = tok.eos_token
        model.eval()
        print("[model] loaded with unsloth.FastModel (the loader used for our submission)")
        return tok, model
    except Exception as e:
        print(f"[model] unsloth unavailable or failed ({type(e).__name__}: {e});"
              " falling back to plain transformers - expect ~1% of answers to differ")

    from transformers import AutoConfig, AutoProcessor, AutoTokenizer

    try:
        proc = AutoProcessor.from_pretrained(MODEL_ID)
        tok = getattr(proc, "tokenizer", proc)
    except Exception:
        tok = AutoTokenizer.from_pretrained(MODEL_ID)
    if getattr(tok, "pad_token", None) is None:
        tok.pad_token = tok.eos_token

    cfg = AutoConfig.from_pretrained(MODEL_ID)
    print("[model] architectures:", getattr(cfg, "architectures", None))

    candidates = []
    import transformers
    for name in ("AutoModelForCausalLM", "AutoModelForImageTextToText",
                 "AutoModelForVision2Seq", "AutoModel"):
        if hasattr(transformers, name):
            candidates.append(getattr(transformers, name))
    last = None
    for cls in candidates:
        try:
            model = cls.from_pretrained(MODEL_ID, dtype=torch.bfloat16, device_map="auto")
            print(f"[model] loaded with {cls.__name__}")
            break
        except Exception as e:                                           # noqa: BLE001
            last = e
            print(f"[model] {cls.__name__} failed: {type(e).__name__}")
    else:
        raise RuntimeError(f"could not load {MODEL_ID}: {last}")
    model.eval()
    return tok, model


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_all(system, rows):
    """For every (question, variant, option) the answer token is scored; within each variant the
    option scores become a log-softmax over displayed positions (putting variants of different
    prompt lengths on one scale), are mapped back to the original option order, and summed."""
    tok, model = system
    dev = next(model.parameters()).device
    pad = tok.pad_token_id if tok.pad_token_id is not None else 0

    items, meta = [], {}
    for ri, r in enumerate(rows):
        choices = [c for c in (r.get("choices") or []) if str(c).strip()]
        if len(choices) < 2:
            continue
        level = get_level(r)
        if level not in SINHALA_INSTRUCTIONS:
            level = "easy"
        md = r.get("metadata") or {}
        for vi, spec in enumerate(VARIANTS):
            prompt, rot = build_prompt(level, r.get("subject"),
                                       md.get("subject_original") or r.get("subject"),
                                       str(r.get("question") or "").strip(), choices, spec)
            meta[(ri, vi)] = (rot, len(choices))
            p_ids = tok(prompt)["input_ids"]
            for j in range(len(choices)):
                f = tok(prompt + f" {j + 1}")["input_ids"]
                if len(f) > MAX_SEQ_LEN:
                    f = f[-MAX_SEQ_LEN:]
                items.append((ri, vi, j, f, max(1, _common_prefix_len(p_ids, f))))

    print(f"{len(rows)} questions x {len(VARIANTS)} variants = {len(items)} candidates")
    scores = {}
    with torch.no_grad():
        for i in range(0, len(items), BATCH):
            chunk = items[i:i + BATCH]
            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
            out = model(input_ids=ids.to(dev), attention_mask=mask.to(dev))
            logits = out.logits if hasattr(out, "logits") else out[0]
            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, out
            print(f"{min(i + BATCH, len(items))}/{len(items)}", end="\r")

    answers = []
    for ri, r in enumerate(rows):
        n = len([c for c in (r.get("choices") or []) if str(c).strip()])
        if n < 2:
            answers.append("1")                 # a prediction is still required
            continue
        total, used = [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)
            z = m + math.log(sum(math.exp(x - m) for x in raw))
            for j in range(nv):
                total[(j + rot) % n] += raw[j] - z
            used += 1
        answers.append("1" if used == 0 else str(1 + max(range(n), key=lambda i: total[i])))
    return answers


# =============== 3. WRITE PREDICTIONS (do not change) ===============
if __name__ == "__main__":
    set_seed(SEED)
    rows = load_questions(DATA_REPO)
    print(len(rows), "questions loaded from", DATA_REPO)

    system = load_system()
    answers = predict_all(system, rows)

    assert len(answers) == len(rows), "predict_all() must return one answer per question"
    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        for r, a in zip(rows, answers):
            a = str(a).strip()
            assert a in {"1", "2", "3", "4"}, f"invalid answer {a!r} for q_no={r.get('q_no')}"
            f.write(json.dumps({
                "q_no": r["q_no"],
                "subject": r["subject"],
                "category": r["category"],
                "level": get_level(r),
                "answer": a,
            }, ensure_ascii=False, separators=(",", ":")) + "\n")
    print("\nSaved", OUTPUT_FILE)