Blame for inference.py

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

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 (keep exactly as used for your final submission)
MODEL_ID = "SinhalaMMLU-SharedTask/team-nena-sinllama-8b-lora"
SEED = 42
BATCH = 8
MAX_NEW_TOKENS = 160

import json
import os
import random
import re
import sys

if hasattr(sys.stdout, "reconfigure"):
    try:
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
        sys.stderr.reconfigure(encoding="utf-8", errors="replace")
    except Exception:
        pass

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


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 (replace this section with your own logic) ===============
# Keep the two function names. predict_all() must return one answer ("1", "2", "3" or "4")
# per question, in the same order as `rows`.

VIRAMA = "\u0DCA"
ZWJ    = "\u200D"
ZWNJ   = "\u200C"
RA     = "\u0DBB"
YA     = "\u0DBA"

def normalize(text):
    if not text or not isinstance(text, str):
        return ""
    t = text.replace("\ufeff", "").replace("\u200b", "")
    t = t.replace(VIRAMA + ZWNJ, VIRAMA)
    t = t.replace(ZWNJ, "")
    t = re.sub(f"{VIRAMA}(?!{ZWJ}){RA}", f"{VIRAMA}{ZWJ}{RA}", t)
    t = re.sub(f"{VIRAMA}(?!{ZWJ}){YA}", f"{VIRAMA}{ZWJ}{YA}", t)
    t = re.sub(f"{ZWJ}+", ZWJ, t)
    return re.sub(r"[ \t]+", " ", t).strip()


SYSTEM_PROMPT = (
    "ඔබ ශ්‍රී ලංකා ජාතික විෂය නිර්දේශයේ ප්‍රවීණ ගුරුවරයෙකි. "
    "පහත බහුවරණ ප්‍රශ්නයට කෙටි හේතු දැක්වීමක් සහිතව නිවැරදි පිළිතුර තෝරන්න."
)


def load_system():
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import PeftModel

    BASE_MODEL = "SAWithanage/SinLlama-Llama-3-8B-Merged"
    tok = AutoTokenizer.from_pretrained(BASE_MODEL)
    tok.padding_side = "left"
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token

    # Load in 4-bit NF4 for universal compatibility across GPUs (T4 / A100 / RTX)
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True
    )
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        quantization_config=bnb_config,
        device_map="auto"
    )
    model = PeftModel.from_pretrained(base_model, MODEL_ID)
    model.eval()
    return tok, model


def make_prompt(r):
    subj = r.get("subject", "General")
    lvl = get_level(r) or "General"
    q_text = normalize(r["question"])
    choices = [normalize(str(c)) for c in r["choices"]]
    choices_block = "\n".join(f"{i + 1}. {c}" for i, c in enumerate(choices))

    user_content = (
        f"<|subject: {subj}|> <|grade: {lvl}|>\n\n"
        f"ප්‍රශ්නය: {q_text}\n\n"
        f"වරණ:\n"
        f"{choices_block}\n\n"
        f"පිළිතුර:"
    )
    return (
        f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"
        f"{SYSTEM_PROMPT}<|eot_id|>"
        f"<|start_header_id|>user<|end_header_id|>\n\n"
        f"{user_content}<|eot_id|>"
        f"<|start_header_id|>assistant<|end_header_id|>\n\n"
    )


def extract_answer(gen_raw, choices=None):
    # Tier 1: Anchored format
    m_ans = re.findall(r"(?:නිවැරදි\s*)?පිළිතුර[:\s]+([1-4])", gen_raw)
    if m_ans:
        return m_ans[-1]

    # Tier 2: Option text matching
    if choices:
        for c_i, c in enumerate(choices):
            c_clean = normalize(str(c))
            if c_clean and len(c_clean) >= 3 and f"{c_i+1}. {c_clean}" in gen_raw:
                return str(c_i + 1)

    # Tier 3: Clean digit fallback
    digits = re.findall(r"\b([1-4])\b", gen_raw)
    return digits[-1] if digits else "1"


def predict_all(system, rows):
    tok, model = system
    answers = [None] * len(rows)

    for b in range(0, len(rows), BATCH):
        batch = rows[b : b + BATCH]
        prompts = [make_prompt(r) for r in batch]
        enc = tok(prompts, return_tensors="pt", padding=True, add_special_tokens=False).to(model.device)

        with torch.no_grad():
            out = model.generate(
                **enc,
                max_new_tokens=MAX_NEW_TOKENS,
                do_sample=False,
                pad_token_id=tok.pad_token_id,
                eos_token_id=tok.eos_token_id
            )

        decoded = tok.batch_decode(out[:, enc["input_ids"].shape[1]:], skip_special_tokens=True)
        for j, d in enumerate(decoded):
            answers[b + j] = extract_answer(d, batch[j]["choices"])
        print(f"{min(b + BATCH, len(rows))}/{len(rows)}", end="\r")

    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)