""" 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)