team-rexx / inference.py

Document cross-machine re-run (91.6% of answers reproduced); tighten vote-gain wording

90522eb verified 18 days ago

19.2 kB

"""
Sinhala MMLU Shared Task - Final inference script
Team: Rexx
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
Rules
  * Everything (models, adapters, extra data) must be downloaded from Hugging Face inside
    this script. No local paths, no Google Drive / Kaggle paths.
  * No API keys or tokens in the code. If something is gated, the organizers log in themselves.
  * No closed or paid APIs. Only your open-weight system (8B total inference parameters or fewer).
  * Deterministic: no sampling (do_sample=False) and a fixed seed, so a re-run gives the same answers.
  * Do not rely on file names, question counts or answers in the dataset. Load every .json file
    in DATA_REPO, as section 1 does. The final dataset has the same structure as the test dataset.
Our system (our final Codabench submission; README.md gives its ID and score)
  * One model, MODEL_ID below: Gemma 4 E4B-it distilled on Sinhala curriculum reasoning traces and averaged
    50/50 with its base weights. 7,941,100,832 parameters. Nothing else is loaded; no retrieval, no internet
    at inference, no tools.
  * Per question: 8 greedy answers (temperature 0, i.e. do_sample=False) with thinking on, each showing the
    four options in a different fixed order, mapped back to the original numbering; the majority answer wins.
    On our validation sets the vote scores 1.0 point above one greedy answer in the original option order.
  * Deterministic: no sampling (greedy decoding), a fixed seed, and VLLM_ENABLE_V1_MULTIPROCESSING=0 for
    deterministic vLLM scheduling. Re-runs on one machine are identical (tests/test_determinism.py). On another
    machine, even with the same GPU model and packages, the GPU kernels round slightly differently and the long
    reasoning can change: a re-run on a second A100 reproduced 91.6% of our final answers (see README.md).
  * The model and recipe code are those of our earlier sampled submission 08 (macro average 64.6025): section 2
    copies its functions verbatim from bakeoff_run.py (SHA-256 00fb7f9c...). Only the decoding differs: greedy
    with 8 answers here, temperature 1.0 with 16 answers there. See README.md.
  * Needs one NVIDIA GPU with at least 40 GB; tested on an A100 40 GB, where the full test set takes about
    56 minutes, including the model download and start-up.
"""

# ======================= 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-rexx-gemma-4-e4b"  # your model/adapter repo, or a public base model
MODEL_REVISION = "f6936800be930a3251f559616f8b136f70b6cfab"  # same files as sinhala-mmlu/gemma-4-e4b-sinhala-mmlu v2.0.0
SEED = 1234
GREEDY = True         # temperature 0, i.e. do_sample=False (submission 08 sampled at temperature 1.0)
SAMPLES = 8           # answers per question, majority vote
SHUFFLE = True        # each answer sees the four options in a different order
THINK_BUDGET = 7168   # max new tokens per sample
MAX_MODEL_LEN = 8192
GPU_MEM = 0.9
SHARDS = 4            # questions run round-robin in 4 parts, as in submission 08; changes batching only
DETAILS_FILE = "predictions_details.jsonl"  # votes and per-sample answers, for checking a re-run

import os

os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")  # deterministic scheduling; set before vllm loads

import json
import math
import random
import re
import subprocess
import sys
import time
import zlib
from types import SimpleNamespace

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


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. YOUR SYSTEM ===============
# The settings and functions from here to run_think() are copied verbatim from bakeoff_run.py, the script
# that produced submission 08 (MODELS["gemma4_e4b"], SYSTEM, POSITIONAL_RE, ANSWER_RE, grade_of,
# user_message, build_llm, digit_token_variants, option_probs, render, score_prompts, permutations_for,
# run_think). tests/test_recipe_identical.py checks that they are unchanged.

CFG = {
    "hf_id": MODEL_ID,  # replaced by the downloaded snapshot folder in load_system()
    "think_kwargs": {"enable_thinking": True},
    "direct_kwargs": {"enable_thinking": False},
    "think_end": "<channel|>",
    "after_think": "",
    "sampling": {"temperature": 1.0, "top_p": 0.95, "top_k": 64},
    "mm_limits": [{"image": 0, "audio": 0, "video": 0}, {"image": 0, "audio": 0}, {"image": 0}],
}
if GREEDY:  # exactly what bakeoff_run.py --greedy does
    CFG["sampling"] = {"temperature": 0.0}

# bakeoff_run.py's command-line arguments (submission 08: submission08/run_reports/).
ARGS = SimpleNamespace(think_budget=THINK_BUDGET, samples=SAMPLES, max_model_len=MAX_MODEL_LEN, gpu_mem=GPU_MEM,
                       seed=SEED, context=True, shuffle=SHUFFLE, reference=False, save_text=False)

SYSTEM = (
    "You are an expert teacher of the Sri Lankan national school curriculum (grades 6-13), "
    "fluent in Sinhala. You will receive a multiple-choice question written in Sinhala with four "
    "numbered options. Exactly one option is correct. Read the Sinhala carefully, including "
    "negations such as 'නොවන' / 'නොවේ' (NOT). You may reason in English. "
    "Finish with the final answer in the form 'Answer: <number>'."
)

# Options that point at other options must keep their order when --shuffle permutes the options.
POSITIONAL_RE = re.compile(r"ඉහත|above|\b[1-4]\s*(?:හා|සහ|හෝ|,|and|or)\s*[1-4]\b", re.I)
ANSWER_RE = re.compile(r"(?:Answer|පිළිතුර)\s*[:：]?\s*\**\s*\(?([1-4])\)?")


def grade_of(row):
    grade = row.get("grade") or (row.get("meta") or {}).get("grade")
    return "AL" if str(grade) in ("12-13", "12", "13") else grade


def user_message(row, context=False, reference=False):
    """context=True adds the subject and grade (both given in the test file) above the question.
    reference=True (context distillation only, never at test time) shows the row's source passage first."""
    options = "\n".join(f"{i}. {c}" for i, c in enumerate(row["choices"], 1))
    head = ""
    if reference and row.get("reference"):
        head = ("පසුබිම් පාඨය (Reference text from the textbook):\n" + row["reference"].strip() + "\n\n"
                "(Use this material to answer, but reason as if you had already learned it yourself: "
                "do not mention or quote the reference text.)\n\n")
    if context:
        grade = grade_of(row)
        head += f"විෂය (Subject): {row['subject']}" + (f" | ශ්‍රේණිය (Grade): {grade}" if grade else "") + "\n\n"
    return head + (
        f"ප්‍රශ්නය (Question):\n{row['question']}\n\n"
        f"විකල්ප (Options):\n{options}\n\n"
        "Which option (1, 2, 3 or 4) is correct? End your reply with 'Answer: <number>'."
    )


def build_llm(cfg, args):
    from vllm import LLM

    last_error = None
    for limits in cfg["mm_limits"]:
        try:
            return LLM(
                model=cfg["hf_id"],
                dtype="bfloat16",
                max_model_len=args.max_model_len,
                gpu_memory_utilization=args.gpu_mem,
                limit_mm_per_prompt=limits,
                seed=args.seed,
                tensor_parallel_size=cfg.get("tp", 1),
            ), limits
        except (ValueError, TypeError) as e:  # unsupported modality key -> try a smaller dict
            last_error = e
            print(f"[warn] limit_mm_per_prompt={limits} rejected: {e}", flush=True)
    raise last_error


def digit_token_variants(tok):
    """Single-token encodings of '1'..'4' with and without a leading space."""
    variants = {}
    for d in "1234":
        ids = set()
        for text in (d, " " + d):
            enc = tok.encode(text, add_special_tokens=False)
            if len(enc) == 1:
                ids.add(enc[0])
        if not ids:
            raise RuntimeError(f"digit {d} has no single-token encoding")
        variants[int(d)] = sorted(ids)
    return variants


def option_probs(output, variants):
    """4-way distribution from the raw top-k logprobs of the first generated token."""
    top = output.outputs[0].logprobs[0]
    logp = {}
    for d, ids in variants.items():
        vals = [top[i].logprob for i in ids if i in top]
        logp[d] = max(vals) if vals else -30.0
    m = max(logp.values())
    z = sum(math.exp(v - m) for v in logp.values())
    return [math.exp(logp[d] - m) / z for d in (1, 2, 3, 4)]


def render(tok, row, kwargs, context=False, reference=False):
    messages = [{"role": "system", "content": SYSTEM},
                {"role": "user", "content": user_message(row, context, reference)}]
    text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, **kwargs)
    # Avoid a double BOS: drop the template's BOS only if the tokenizer adds one itself
    # (vLLM tokenizes text prompts with add_special_tokens=True).
    if tok.bos_token and text.startswith(tok.bos_token) and tok("x")["input_ids"][:1] == [tok.bos_token_id]:
        text = text[len(tok.bos_token):]
    return text


def score_prompts(llm, prompts, variants, seed):
    from vllm import SamplingParams

    allowed = sorted({i for ids in variants.values() for i in ids})
    sp = SamplingParams(max_tokens=1, temperature=0.0, logprobs=20, allowed_token_ids=allowed, seed=seed)
    outs = llm.generate(prompts, sp, use_tqdm=True)
    return [option_probs(o, variants) for o in outs]


def permutations_for(row, n):
    """n option orders for one question: rotations of a per-question random order (then its reverse), so each
    option appears in each position equally often. perm[j] = original index shown at position j."""
    if POSITIONAL_RE.search(" ".join(row["choices"])):
        return [[0, 1, 2, 3]] * n
    base = list(range(4))
    random.Random(zlib.crc32(row["id"].encode())).shuffle(base)
    orders = [base[i:] + base[:i] for i in range(4)]
    orders += [list(reversed(o)) for o in orders]
    return [orders[k % 8] for k in range(n)]


def run_think(llm, tok, cfg, rows, variants, args):
    from vllm import SamplingParams

    def sampling(seed, n):
        return SamplingParams(max_tokens=args.think_budget, n=n, seed=seed, skip_special_tokens=False,
                              **cfg["sampling"])

    # One request per (question, sample) when shuffling, since each sample sees its own option order;
    # otherwise one request per question with n samples.
    jobs = []  # (row, k, perm, prompt)
    for r in rows:
        if args.shuffle:
            for k, perm in enumerate(permutations_for(r, args.samples)):
                shown = r | {"choices": [r["choices"][i] for i in perm]}
                jobs.append((r, k, perm, render(tok, shown, cfg["think_kwargs"], args.context, args.reference)))
        else:
            jobs.append((r, None, [0, 1, 2, 3], render(tok, r, cfg["think_kwargs"], args.context, args.reference)))
    prompts = [j[3] for j in jobs]
    if args.shuffle:
        sps = [sampling(args.seed * 1000 + k, 1) for _, k, _, _ in jobs]
    else:
        sps = sampling(args.seed, args.samples)
    t0 = time.time()
    outs = llm.generate(prompts, sps, use_tqdm=True)
    gen_secs = time.time() - t0

    end = cfg["think_end"]
    forced_prompts, records = [], []
    for (r, k0, perm, prompt), out in zip(jobs, outs):
        for k, sample in enumerate(out.outputs):
            text = sample.text
            closed = end in text
            thought = text.split(end, 1)[0] if closed else text
            reply = text.split(end, 1)[1] if closed else ""
            matches = ANSWER_RE.findall(reply)
            shown = int(matches[-1]) if matches else None
            forced_prompts.append(prompt + thought + ("" if thought.endswith("\n") else "\n") + end
                                  + cfg["after_think"] + "Answer: ")
            records.append({"row": r, "k": k if k0 is None else k0, "perm": perm,
                            "parsed": None if shown is None else perm[shown - 1] + 1, "closed": closed,
                            "tokens": len(sample.token_ids), "finish": sample.finish_reason,
                            "reply_tail": reply[-200:], "thought": thought if args.save_text else None})
    t1 = time.time()
    forced = score_prompts(llm, forced_prompts, variants, args.seed)
    score_secs = time.time() - t1

    by_row = {}
    for rec, p in zip(records, forced):
        # Map the shown-position probabilities back to the original option order.
        orig = [0.0] * 4
        for j, i in enumerate(rec["perm"]):
            orig[i] = p[j]
        rec["forced_probs"] = orig
        by_row.setdefault(rec["row"]["id"], []).append(rec)
    preds = []
    for r in rows:
        samples = by_row[r["id"]]
        votes = [0.0] * 4
        for s in samples:
            choice = s["parsed"] if s["parsed"] is not None else 1 + max(range(4), key=lambda i: s["forced_probs"][i])
            votes[choice - 1] += 1.0
        mean_p = [sum(s["forced_probs"][i] for s in samples) / len(samples) for i in range(4)]
        pred = 1 + max(range(4), key=lambda i: (votes[i], mean_p[i]))
        preds.append({
            "id": r["id"], "subject": r["subject"], "gold": r.get("answer"), "pred": pred, "probs": mean_p,
            "votes": votes,
            "samples": [{k: s[k] for k in ("k", "perm", "parsed", "closed", "tokens", "finish", "reply_tail")}
                        | {"forced_pred": 1 + max(range(4), key=lambda i: s["forced_probs"][i])}
                        | ({"thought": s["thought"]} if args.save_text else {})
                        for s in samples],
        })
    return preds, {"generation_seconds": gen_secs, "forced_scoring_seconds": score_secs,
                   "example_prompt": prompts[0], "example_output": outs[0].outputs[0].text[:3000]}


# ----- Glue between the template and the recipe above (not part of bakeoff_run.py) -----

LEVEL_ORDER = ("easy", "medium", "hard")


def remove_unused_audio_packages():
    """vLLM 0.32.0 from PyPI brings a torchaudio build for another CUDA version than PyTorch 2.13.0, which stops
    Transformers from importing, and torchcodec's native library fails to load in some images. This text-only
    system uses neither, so they are removed, exactly as in our tested Colab notebook."""
    import importlib.util

    present = [p for p in ("torchaudio", "torchcodec") if importlib.util.find_spec(p)]
    if present:
        print("removing unused packages:", ", ".join(present), flush=True)
        subprocess.run([sys.executable, "-m", "pip", "uninstall", "-y", "-q", *present], check=False)


def task_rows(rows):
    """The template's rows as recipe rows, in the order of submission 08's input file: levels easy, medium,
    hard; files by name; questions in file order. The id "<level>|<file name>|<q_no>" seeds each question's
    option order, exactly as in submission 08."""
    order = sorted(range(len(rows)), key=lambda i: (
        LEVEL_ORDER.index(get_level(rows[i])) if get_level(rows[i]) in LEVEL_ORDER else len(LEVEL_ORDER),
        get_level(rows[i]), rows[i]["_file"].rsplit("/", 1)[-1]))
    out, seen = [], {}
    for i in order:
        r = rows[i]
        if len(r["choices"]) != 4:
            raise ValueError(f"{r['_file']} q_no {r.get('q_no')}: {len(r['choices'])} choices, expected 4")
        qid = f"{get_level(r)}|{r['_file'].rsplit('/', 1)[-1]}|{r['q_no']}"
        seen[qid] = seen.get(qid, 0) + 1
        if seen[qid] > 1:  # keep ids unique: votes are grouped by id
            qid += f"#{seen[qid]}"
        out.append((i, {"id": qid, "subject": r["subject"], "grade": (r.get("metadata") or {}).get("grade"),
                        "question": r["question"], "choices": r["choices"], "answer": None}))
    return out


def load_system():
    remove_unused_audio_packages()
    model_dir = snapshot_download(MODEL_ID, revision=MODEL_REVISION)
    cfg = CFG | {"hf_id": model_dir}
    llm, limits = build_llm(cfg, ARGS)
    tok = llm.get_tokenizer()
    return llm, tok, cfg, digit_token_variants(tok)


def predict_all(system, rows):
    llm, tok, cfg, variants = system
    ordered = task_rows(rows)
    answers, details = [None] * len(rows), []
    for s in range(SHARDS):
        shard = ordered[s::SHARDS]
        t0 = time.time()
        preds, info = run_think(llm, tok, cfg, [r for _, r in shard], variants, ARGS)
        for (i, _), p in zip(shard, preds):
            answers[i] = str(p["pred"])
            details.append(p | {"q_no": rows[i]["q_no"], "level": get_level(rows[i])})
        samples = [x for p in preds for x in p["samples"]]
        print(f"shard {s + 1}/{SHARDS}: {len(shard)} questions in {time.time() - t0:.0f} s; "
              f"answer parsed {sum(x['parsed'] is not None for x in samples) / len(samples):.4f}, "
              f"thinking closed {sum(x['closed'] for x in samples) / len(samples):.4f}", flush=True)
    with open(DETAILS_FILE, "w", encoding="utf-8") as f:
        for p in details:
            f.write(json.dumps(p, ensure_ascii=False) + "\n")
    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)

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

	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

	Rules
	* Everything (models, adapters, extra data) must be downloaded from Hugging Face inside
	this script. No local paths, no Google Drive / Kaggle paths.
	* No API keys or tokens in the code. If something is gated, the organizers log in themselves.
	* No closed or paid APIs. Only your open-weight system (8B total inference parameters or fewer).
	* Deterministic: no sampling (do_sample=False) and a fixed seed, so a re-run gives the same answers.
	* Do not rely on file names, question counts or answers in the dataset. Load every .json file
	in DATA_REPO, as section 1 does. The final dataset has the same structure as the test dataset.

	Our system (our final Codabench submission; README.md gives its ID and score)
	* One model, MODEL_ID below: Gemma 4 E4B-it distilled on Sinhala curriculum reasoning traces and averaged
	50/50 with its base weights. 7,941,100,832 parameters. Nothing else is loaded; no retrieval, no internet
	at inference, no tools.
	* Per question: 8 greedy answers (temperature 0, i.e. do_sample=False) with thinking on, each showing the
	four options in a different fixed order, mapped back to the original numbering; the majority answer wins.
	On our validation sets the vote scores 1.0 point above one greedy answer in the original option order.
	* Deterministic: no sampling (greedy decoding), a fixed seed, and VLLM_ENABLE_V1_MULTIPROCESSING=0 for
	deterministic vLLM scheduling. Re-runs on one machine are identical (tests/test_determinism.py). On another
	machine, even with the same GPU model and packages, the GPU kernels round slightly differently and the long
	reasoning can change: a re-run on a second A100 reproduced 91.6% of our final answers (see README.md).
	* The model and recipe code are those of our earlier sampled submission 08 (macro average 64.6025): section 2
	copies its functions verbatim from bakeoff_run.py (SHA-256 00fb7f9c...). Only the decoding differs: greedy
	with 8 answers here, temperature 1.0 with 16 answers there. See README.md.
	* Needs one NVIDIA GPU with at least 40 GB; tested on an A100 40 GB, where the full test set takes about
	56 minutes, including the model download and start-up.
	"""

	# ======================= 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-rexx-gemma-4-e4b" # your model/adapter repo, or a public base model
	MODEL_REVISION = "f6936800be930a3251f559616f8b136f70b6cfab" # same files as sinhala-mmlu/gemma-4-e4b-sinhala-mmlu v2.0.0
	SEED = 1234
	GREEDY = True # temperature 0, i.e. do_sample=False (submission 08 sampled at temperature 1.0)
	SAMPLES = 8 # answers per question, majority vote
	SHUFFLE = True # each answer sees the four options in a different order
	THINK_BUDGET = 7168 # max new tokens per sample
	MAX_MODEL_LEN = 8192
	GPU_MEM = 0.9
	SHARDS = 4 # questions run round-robin in 4 parts, as in submission 08; changes batching only
	DETAILS_FILE = "predictions_details.jsonl" # votes and per-sample answers, for checking a re-run

	import os

	os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0") # deterministic scheduling; set before vllm loads

	import json
	import math
	import random
	import re
	import subprocess
	import sys
	import time
	import zlib
	from types import SimpleNamespace

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


	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. YOUR SYSTEM ===============
	# The settings and functions from here to run_think() are copied verbatim from bakeoff_run.py, the script
	# that produced submission 08 (MODELS["gemma4_e4b"], SYSTEM, POSITIONAL_RE, ANSWER_RE, grade_of,
	# user_message, build_llm, digit_token_variants, option_probs, render, score_prompts, permutations_for,
	# run_think). tests/test_recipe_identical.py checks that they are unchanged.

	CFG = {
	"hf_id": MODEL_ID, # replaced by the downloaded snapshot folder in load_system()
	"think_kwargs": {"enable_thinking": True},
	"direct_kwargs": {"enable_thinking": False},
	"think_end": "<channel\|>",
	"after_think": "",
	"sampling": {"temperature": 1.0, "top_p": 0.95, "top_k": 64},
	"mm_limits": [{"image": 0, "audio": 0, "video": 0}, {"image": 0, "audio": 0}, {"image": 0}],
	}
	if GREEDY: # exactly what bakeoff_run.py --greedy does
	CFG["sampling"] = {"temperature": 0.0}

	# bakeoff_run.py's command-line arguments (submission 08: submission08/run_reports/).
	ARGS = SimpleNamespace(think_budget=THINK_BUDGET, samples=SAMPLES, max_model_len=MAX_MODEL_LEN, gpu_mem=GPU_MEM,
	seed=SEED, context=True, shuffle=SHUFFLE, reference=False, save_text=False)

	SYSTEM = (
	"You are an expert teacher of the Sri Lankan national school curriculum (grades 6-13), "
	"fluent in Sinhala. You will receive a multiple-choice question written in Sinhala with four "
	"numbered options. Exactly one option is correct. Read the Sinhala carefully, including "
	"negations such as 'නොවන' / 'නොවේ' (NOT). You may reason in English. "
	"Finish with the final answer in the form 'Answer: <number>'."
	)

	# Options that point at other options must keep their order when --shuffle permutes the options.
	POSITIONAL_RE = re.compile(r"ඉහත\|above\|\b[1-4]\s(?:හා\|සහ\|හෝ\|,\|and\|or)\s[1-4]\b", re.I)
	ANSWER_RE = re.compile(r"(?:Answer\|පිළිතුර)\s[:：]?\s\*\s\(?([1-4])\)?")


	def grade_of(row):
	grade = row.get("grade") or (row.get("meta") or {}).get("grade")
	return "AL" if str(grade) in ("12-13", "12", "13") else grade


	def user_message(row, context=False, reference=False):
	"""context=True adds the subject and grade (both given in the test file) above the question.
	reference=True (context distillation only, never at test time) shows the row's source passage first."""
	options = "\n".join(f"{i}. {c}" for i, c in enumerate(row["choices"], 1))
	head = ""
	if reference and row.get("reference"):
	head = ("පසුබිම් පාඨය (Reference text from the textbook):\n" + row["reference"].strip() + "\n\n"
	"(Use this material to answer, but reason as if you had already learned it yourself: "
	"do not mention or quote the reference text.)\n\n")
	if context:
	grade = grade_of(row)
	head += f"විෂය (Subject): {row['subject']}" + (f" \| ශ්‍රේණිය (Grade): {grade}" if grade else "") + "\n\n"
	return head + (
	f"ප්‍රශ්නය (Question):\n{row['question']}\n\n"
	f"විකල්ප (Options):\n{options}\n\n"
	"Which option (1, 2, 3 or 4) is correct? End your reply with 'Answer: <number>'."
	)


	def build_llm(cfg, args):
	from vllm import LLM

	last_error = None
	for limits in cfg["mm_limits"]:
	try:
	return LLM(
	model=cfg["hf_id"],
	dtype="bfloat16",
	max_model_len=args.max_model_len,
	gpu_memory_utilization=args.gpu_mem,
	limit_mm_per_prompt=limits,
	seed=args.seed,
	tensor_parallel_size=cfg.get("tp", 1),
	), limits
	except (ValueError, TypeError) as e: # unsupported modality key -> try a smaller dict
	last_error = e
	print(f"[warn] limit_mm_per_prompt={limits} rejected: {e}", flush=True)
	raise last_error


	def digit_token_variants(tok):
	"""Single-token encodings of '1'..'4' with and without a leading space."""
	variants = {}
	for d in "1234":
	ids = set()
	for text in (d, " " + d):
	enc = tok.encode(text, add_special_tokens=False)
	if len(enc) == 1:
	ids.add(enc[0])
	if not ids:
	raise RuntimeError(f"digit {d} has no single-token encoding")
	variants[int(d)] = sorted(ids)
	return variants


	def option_probs(output, variants):
	"""4-way distribution from the raw top-k logprobs of the first generated token."""
	top = output.outputs[0].logprobs[0]
	logp = {}
	for d, ids in variants.items():
	vals = [top[i].logprob for i in ids if i in top]
	logp[d] = max(vals) if vals else -30.0
	m = max(logp.values())
	z = sum(math.exp(v - m) for v in logp.values())
	return [math.exp(logp[d] - m) / z for d in (1, 2, 3, 4)]


	def render(tok, row, kwargs, context=False, reference=False):
	messages = [{"role": "system", "content": SYSTEM},
	{"role": "user", "content": user_message(row, context, reference)}]
	text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, **kwargs)
	# Avoid a double BOS: drop the template's BOS only if the tokenizer adds one itself
	# (vLLM tokenizes text prompts with add_special_tokens=True).
	if tok.bos_token and text.startswith(tok.bos_token) and tok("x")["input_ids"][:1] == [tok.bos_token_id]:
	text = text[len(tok.bos_token):]
	return text


	def score_prompts(llm, prompts, variants, seed):
	from vllm import SamplingParams

	allowed = sorted({i for ids in variants.values() for i in ids})
	sp = SamplingParams(max_tokens=1, temperature=0.0, logprobs=20, allowed_token_ids=allowed, seed=seed)
	outs = llm.generate(prompts, sp, use_tqdm=True)
	return [option_probs(o, variants) for o in outs]


	def permutations_for(row, n):
	"""n option orders for one question: rotations of a per-question random order (then its reverse), so each
	option appears in each position equally often. perm[j] = original index shown at position j."""
	if POSITIONAL_RE.search(" ".join(row["choices"])):
	return [[0, 1, 2, 3]] * n
	base = list(range(4))
	random.Random(zlib.crc32(row["id"].encode())).shuffle(base)
	orders = [base[i:] + base[:i] for i in range(4)]
	orders += [list(reversed(o)) for o in orders]
	return [orders[k % 8] for k in range(n)]


	def run_think(llm, tok, cfg, rows, variants, args):
	from vllm import SamplingParams

	def sampling(seed, n):
	return SamplingParams(max_tokens=args.think_budget, n=n, seed=seed, skip_special_tokens=False,
	**cfg["sampling"])

	# One request per (question, sample) when shuffling, since each sample sees its own option order;
	# otherwise one request per question with n samples.
	jobs = [] # (row, k, perm, prompt)
	for r in rows:
	if args.shuffle:
	for k, perm in enumerate(permutations_for(r, args.samples)):
	shown = r \| {"choices": [r["choices"][i] for i in perm]}
	jobs.append((r, k, perm, render(tok, shown, cfg["think_kwargs"], args.context, args.reference)))
	else:
	jobs.append((r, None, [0, 1, 2, 3], render(tok, r, cfg["think_kwargs"], args.context, args.reference)))
	prompts = [j[3] for j in jobs]
	if args.shuffle:
	sps = [sampling(args.seed * 1000 + k, 1) for _, k, _, _ in jobs]
	else:
	sps = sampling(args.seed, args.samples)
	t0 = time.time()
	outs = llm.generate(prompts, sps, use_tqdm=True)
	gen_secs = time.time() - t0

	end = cfg["think_end"]
	forced_prompts, records = [], []
	for (r, k0, perm, prompt), out in zip(jobs, outs):
	for k, sample in enumerate(out.outputs):
	text = sample.text
	closed = end in text
	thought = text.split(end, 1)[0] if closed else text
	reply = text.split(end, 1)[1] if closed else ""
	matches = ANSWER_RE.findall(reply)
	shown = int(matches[-1]) if matches else None
	forced_prompts.append(prompt + thought + ("" if thought.endswith("\n") else "\n") + end
	+ cfg["after_think"] + "Answer: ")
	records.append({"row": r, "k": k if k0 is None else k0, "perm": perm,
	"parsed": None if shown is None else perm[shown - 1] + 1, "closed": closed,
	"tokens": len(sample.token_ids), "finish": sample.finish_reason,
	"reply_tail": reply[-200:], "thought": thought if args.save_text else None})
	t1 = time.time()
	forced = score_prompts(llm, forced_prompts, variants, args.seed)
	score_secs = time.time() - t1

	by_row = {}
	for rec, p in zip(records, forced):
	# Map the shown-position probabilities back to the original option order.
	orig = [0.0] * 4
	for j, i in enumerate(rec["perm"]):
	orig[i] = p[j]
	rec["forced_probs"] = orig
	by_row.setdefault(rec["row"]["id"], []).append(rec)
	preds = []
	for r in rows:
	samples = by_row[r["id"]]
	votes = [0.0] * 4
	for s in samples:
	choice = s["parsed"] if s["parsed"] is not None else 1 + max(range(4), key=lambda i: s["forced_probs"][i])
	votes[choice - 1] += 1.0
	mean_p = [sum(s["forced_probs"][i] for s in samples) / len(samples) for i in range(4)]
	pred = 1 + max(range(4), key=lambda i: (votes[i], mean_p[i]))
	preds.append({
	"id": r["id"], "subject": r["subject"], "gold": r.get("answer"), "pred": pred, "probs": mean_p,
	"votes": votes,
	"samples": [{k: s[k] for k in ("k", "perm", "parsed", "closed", "tokens", "finish", "reply_tail")}
	\| {"forced_pred": 1 + max(range(4), key=lambda i: s["forced_probs"][i])}
	\| ({"thought": s["thought"]} if args.save_text else {})
	for s in samples],
	})
	return preds, {"generation_seconds": gen_secs, "forced_scoring_seconds": score_secs,
	"example_prompt": prompts[0], "example_output": outs[0].outputs[0].text[:3000]}


	# ----- Glue between the template and the recipe above (not part of bakeoff_run.py) -----

	LEVEL_ORDER = ("easy", "medium", "hard")


	def remove_unused_audio_packages():
	"""vLLM 0.32.0 from PyPI brings a torchaudio build for another CUDA version than PyTorch 2.13.0, which stops
	Transformers from importing, and torchcodec's native library fails to load in some images. This text-only
	system uses neither, so they are removed, exactly as in our tested Colab notebook."""
	import importlib.util

	present = [p for p in ("torchaudio", "torchcodec") if importlib.util.find_spec(p)]
	if present:
	print("removing unused packages:", ", ".join(present), flush=True)
	subprocess.run([sys.executable, "-m", "pip", "uninstall", "-y", "-q", *present], check=False)


	def task_rows(rows):
	"""The template's rows as recipe rows, in the order of submission 08's input file: levels easy, medium,
	hard; files by name; questions in file order. The id "<level>\|<file name>\|<q_no>" seeds each question's
	option order, exactly as in submission 08."""
	order = sorted(range(len(rows)), key=lambda i: (
	LEVEL_ORDER.index(get_level(rows[i])) if get_level(rows[i]) in LEVEL_ORDER else len(LEVEL_ORDER),
	get_level(rows[i]), rows[i]["_file"].rsplit("/", 1)[-1]))
	out, seen = [], {}
	for i in order:
	r = rows[i]
	if len(r["choices"]) != 4:
	raise ValueError(f"{r['_file']} q_no {r.get('q_no')}: {len(r['choices'])} choices, expected 4")
	qid = f"{get_level(r)}\|{r['_file'].rsplit('/', 1)[-1]}\|{r['q_no']}"
	seen[qid] = seen.get(qid, 0) + 1
	if seen[qid] > 1: # keep ids unique: votes are grouped by id
	qid += f"#{seen[qid]}"
	out.append((i, {"id": qid, "subject": r["subject"], "grade": (r.get("metadata") or {}).get("grade"),
	"question": r["question"], "choices": r["choices"], "answer": None}))
	return out


	def load_system():
	remove_unused_audio_packages()
	model_dir = snapshot_download(MODEL_ID, revision=MODEL_REVISION)
	cfg = CFG \| {"hf_id": model_dir}
	llm, limits = build_llm(cfg, ARGS)
	tok = llm.get_tokenizer()
	return llm, tok, cfg, digit_token_variants(tok)


	def predict_all(system, rows):
	llm, tok, cfg, variants = system
	ordered = task_rows(rows)
	answers, details = [None] * len(rows), []
	for s in range(SHARDS):
	shard = ordered[s::SHARDS]
	t0 = time.time()
	preds, info = run_think(llm, tok, cfg, [r for _, r in shard], variants, ARGS)
	for (i, _), p in zip(shard, preds):
	answers[i] = str(p["pred"])
	details.append(p \| {"q_no": rows[i]["q_no"], "level": get_level(rows[i])})
	samples = [x for p in preds for x in p["samples"]]
	print(f"shard {s + 1}/{SHARDS}: {len(shard)} questions in {time.time() - t0:.0f} s; "
	f"answer parsed {sum(x['parsed'] is not None for x in samples) / len(samples):.4f}, "
	f"thinking closed {sum(x['closed'] for x in samples) / len(samples):.4f}", flush=True)
	with open(DETAILS_FILE, "w", encoding="utf-8") as f:
	for p in details:
	f.write(json.dumps(p, ensure_ascii=False) + "\n")
	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)