""" Sinhala MMLU Shared Task - Final inference script Team: hakuru labs 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. """ # ======================= ORGANIZERS CHANGE ONLY THIS LINE ======================= DATA_REPO = "SinhalaMMLU-SharedTask/Sinhala-MMLU-SharedTask-Test-Dataset" # ================================================================================ OUTPUT_FILE = "predictions.jsonl" # ----------------------- TEAM SETTINGS ----------------------- MODEL_RQ_BASE = "iCIIT/redqueenprotocol-sin-llama3.2-3B-model" MODEL_RQ_ADAPTER = "SinhalaMMLU-SharedTask/team-hakuru-labs-redqueen" MODEL_HL_BASE = "Jainamshahhh/helalm-4b" MODEL_HL_ADAPTER = "SinhalaMMLU-SharedTask/team-hakuru-labs-helalm" SEED = 3407 BATCH = 1 MAX_NEW_TOKENS = 16 # ------------------------------------------------------------- import json import os import random import re import numpy as np import torch from huggingface_hub import HfFileSystem, list_repo_files from collections import Counter LETTERS = ["A", "B", "C", "D"] ABSOLUTE_WORDS = ["සියල්ල", "කිසිවිටෙක", "සෑම", "හැමවිටම", "කිසිදා"] PENALTY = 2.0 INSTRUCTION = "පහත බහුවරණ ප්‍රශ්නයට නිවැරදි පිළිතුර තෝරන්න. පළමුව සෑම විකල්පයක් ගැන පියවරෙන් පියවර සිතන්න, පසුව පිළිතුර ලෙස A, B, C, හෝ D පමණක් ලියන්න." def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def check_parameters(model, name): total_params = sum(p.numel() for p in model.parameters()) print(f"[{name}] Logical Parameters: {total_params:,} ({total_params/1e9:.3f} B)") return total_params def get_lexical_overlap(question, option): q_words = {w for w in re.findall(r'\b\w+\b', question) if len(w) > 3} o_words = {w for w in re.findall(r'\b\w+\b', option) if len(w) > 3} return len(q_words.intersection(o_words)) # =============== 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 =============== def load_system(): from unsloth import FastLanguageModel from peft import PeftModel print("Loading RedQueen v7...") model_rq, tokenizer_rq = FastLanguageModel.from_pretrained( model_name=MODEL_RQ_BASE, max_seq_length=2048, load_in_4bit=True ) model_rq = PeftModel.from_pretrained(model_rq, MODEL_RQ_ADAPTER) model_rq = model_rq.to(torch.float16) FastLanguageModel.for_inference(model_rq) rq_params = check_parameters(model_rq, "RedQueen v7") print("Loading HelaLM v2...") model_hl, tokenizer_hl = FastLanguageModel.from_pretrained( model_name=MODEL_HL_BASE, max_seq_length=2048, load_in_4bit=True ) model_hl = PeftModel.from_pretrained(model_hl, MODEL_HL_ADAPTER) model_hl = model_hl.to(torch.float16) FastLanguageModel.for_inference(model_hl) hl_params = check_parameters(model_hl, "HelaLM v2") total_ensemble_params = rq_params + hl_params print(f"\n=======================================================") print(f"TOTAL ENSEMBLE PARAMETERS: {total_ensemble_params:,} ({total_ensemble_params/1e9:.3f} B)") if total_ensemble_params <= 8_000_000_000: print("✅ COMPLIANT: Total parameters are under the 8B limit.") else: print("❌ WARNING: Total parameters exceed the 8B limit!") print(f"=======================================================\n") return (model_rq, tokenizer_rq, model_hl, tokenizer_hl) @torch.no_grad() def predict_all(system, rows): model_rq, tokenizer_rq, model_hl, tokenizer_hl = system def build_rq_prompt(q, opts): block = "\n".join(f"{LETTERS[i]}) {opts[i].strip()}" for i in range(4)) return f"### Instruction:\n{INSTRUCTION}\n\n### Input:\n{q.strip()}\n{block}\n\n### Response:\nපිළිතුර:" def build_hl_prompt(q, opts): block = "\n".join(f"{LETTERS[i]}) {opts[i].strip()}" for i in range(4)) return (f"user\n{q.strip()}\n{block}\n" f"model\nපිළිතුර:") ids_rq = {L: list(set(tokenizer_rq(f" {L}", add_special_tokens=False).input_ids + tokenizer_rq(L, add_special_tokens=False).input_ids)) for L in LETTERS} ids_hl = {L: list(set(tokenizer_hl(f" {L}", add_special_tokens=False).input_ids + tokenizer_hl(L, add_special_tokens=False).input_ids)) for L in LETTERS} @torch.no_grad() def get_consistent_logits(model, tok, q, opts, ids_map, prompt_builder, n_runs=5): rng = random.Random(SEED) avg_logits = {L: 0.0 for L in LETTERS} for _ in range(n_runs): order = list(range(4)) rng.shuffle(order) shuffled_opts = [opts[i] for i in order] prompt = prompt_builder(q, shuffled_opts) inputs = tok(prompt, return_tensors="pt").to("cuda") logits = model(**inputs).logits[0, -1, :] run_scores = {L: max(logits[t].item() for t in ids_map[L]) for L in LETTERS} for i, original_idx in enumerate(order): avg_logits[LETTERS[original_idx]] += run_scores[LETTERS[i]] return {L: avg_logits[L] / n_runs for L in LETTERS} answers = [] for i, rec in enumerate(rows): q = rec["question"] opts = rec["choices"] level = str(get_level(rec)).lower() if level == "hard": # Route ONLY to HelaLM for Hard questions final_scores = get_consistent_logits(model_hl, tokenizer_hl, q, opts, ids_hl, build_hl_prompt) else: # Use Ensemble for Easy/Medium rq_avg = get_consistent_logits(model_rq, tokenizer_rq, q, opts, ids_rq, build_rq_prompt) hl_avg = get_consistent_logits(model_hl, tokenizer_hl, q, opts, ids_hl, build_hl_prompt) final_scores = {L: (rq_avg[L] + hl_avg[L]) / 2.0 for L in LETTERS} for j, opt in enumerate(opts): if any(word in opt for word in ABSOLUTE_WORDS): final_scores[LETTERS[j]] -= PENALTY sorted_scores = sorted(final_scores.items(), key=lambda x: x[1], reverse=True) top_letter, top_score = sorted_scores[0] second_letter, second_score = sorted_scores[1] pred_idx = LETTERS.index(top_letter) if (top_score - second_score) < 1.5: top_idx = LETTERS.index(top_letter) second_idx = LETTERS.index(second_letter) if get_lexical_overlap(q, opts[second_idx]) > get_lexical_overlap(q, opts[top_idx]): pred_idx = second_idx answers.append(str(pred_idx + 1)) if (i+1) % 100 == 0: print(f"Processed {i+1}/{len(rows)}") 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)