team-nena / inference.py

Team Nena: Backup submission package

ca24bf9 verified 12 days ago

7.74 kB

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

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