| """
|
| 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.
|
| """
|
|
|
|
|
| DATA_REPO = "SinhalaMMLU-SharedTask/Sinhala-MMLU-SharedTask-Test-Dataset"
|
|
|
|
|
| OUTPUT_FILE = "predictions.jsonl"
|
|
|
|
|
| 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))
|
|
|
|
|
| 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]
|
|
|
|
|
| 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"<start_of_turn>user\n{q.strip()}\n{block}<end_of_turn>\n"
|
| f"<start_of_turn>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":
|
|
|
| final_scores = get_consistent_logits(model_hl, tokenizer_hl, q, opts, ids_hl, build_hl_prompt)
|
| else:
|
|
|
| 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
|
|
|
|
|
|
|
| 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) |
