| """ |
| 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. |
| """ |
|
|
| |
| DATA_REPO = "SinhalaMMLU-SharedTask/Sinhala-MMLU-SharedTask-Test-Dataset" |
| |
|
|
| OUTPUT_FILE = "predictions.jsonl" |
|
|
| |
| MODEL_ID = "SinhalaMMLU-SharedTask/team-qbits-cpt-v3" |
| SEED = 42 |
| BATCH = 16 |
| MAX_SEQ_LEN = 1024 |
|
|
| |
| |
| |
| 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) |
|
|
|
|
| |
| 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] |
|
|
|
|
| |
| |
| |
|
|
| |
| 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: |
| 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") |
| 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 |
|
|
|
|
| |
| 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) |
|
|
