File size: 4,565 Bytes
044f4e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | """
Sinhala MMLU Shared Task - Final inference script
Team: SinhalaNuwana
pip install -r requirements.txt
python inference.py
Downloads the questions from DATA_REPO, the base model and our LoRA adapter from Hugging Face,
runs predict.py six times and merge_votes.py once (both next to this file), and writes
predictions.jsonl and predictions.zip in the current folder. Greedy decoding, no sampling.
About 16 hours for the 2,238 test questions on 2 x T4.
"""
# ======================= ORGANIZERS CHANGE ONLY THIS LINE =======================
DATA_REPO = "SinhalaMMLU-SharedTask/Sinhala-MMLU-SharedTask-Test-Dataset"
# ================================================================================
import argparse
import hashlib
import json
import os
import subprocess
import sys
from pathlib import Path
from huggingface_hub import snapshot_download
MODEL = "google/gemma-4-E4B-it"
REVISION = "ee0ef6023621cff504d758262d4e04895a5af4a2"
ADAPTER = "SinhalaMMLU-SharedTask/team-sinhalanuwana-gemma4-e4b-star3-lora"
ADAPTER_SHA256 = "b82696e2541e0133d24bf14721caa20485dd70ab8ef78a7e40c1d2a77db1ba71"
ROOT = Path(__file__).resolve().parent
PASSES = [
("orders-02", ["--think", "512", "--rotations", "0", "2"]),
("orders-13", ["--think", "512", "--rotations", "1", "3"]),
("hard-2048", ["--level", "hard", "--think", "2048", "--score_batch_size", "8"]),
("easy-2048", ["--level", "easy", "--think", "2048", "--score_batch_size", "8"]),
("medium-rev02", ["--level", "medium", "--reverse", "--think", "512", "--rotations", "0", "2"]),
("medium-rev13", ["--level", "medium", "--reverse", "--think", "512", "--rotations", "1", "3"]),
]
def slice_questions(data_dir, out_dir, per_file):
for path in sorted(data_dir.rglob("*.json")):
if any(part.startswith(".") for part in path.relative_to(data_dir).parts):
continue
target = out_dir / path.relative_to(data_dir)
target.parent.mkdir(parents=True, exist_ok=True)
questions = json.loads(path.read_text(encoding="utf-8"))[:per_file]
target.write_text(json.dumps(questions, ensure_ascii=False), encoding="utf-8")
return out_dir
def main():
parser = argparse.ArgumentParser(
description="Answer the SinhalaMMLU questions with Gemma 4 E4B-it and the round 3 STaR LoRA adapter, "
"and write OUT_DIR/predictions.jsonl and OUT_DIR/predictions.zip"
)
parser.add_argument("--data_dir", type=Path, help=f"local folder of the per-subject .json files; {DATA_REPO} when not given")
parser.add_argument("--out_dir", type=Path, default=Path("."), help="where predictions.jsonl goes; each pass writes to OUT_DIR/runs/")
parser.add_argument("--adapter", help=f"local adapter folder; {ADAPTER} from the Hub when not given")
parser.add_argument("--per_file", type=int, help="answer only the first N questions of each file, for a quick check")
args = parser.parse_args()
snapshot_download(MODEL, revision=REVISION)
adapter = Path(args.adapter or snapshot_download(ADAPTER))
digest = hashlib.sha256((adapter / "adapter_model.safetensors").read_bytes()).hexdigest()
if digest != ADAPTER_SHA256:
raise SystemExit(f"error: {adapter}/adapter_model.safetensors is {digest}, expected {ADAPTER_SHA256}")
data_dir = args.data_dir or Path(snapshot_download(DATA_REPO, repo_type="dataset"))
runs = args.out_dir / "runs"
if args.per_file:
data_dir = slice_questions(data_dir, runs / "questions", args.per_file)
env = {**os.environ, "HF_HUB_OFFLINE": "1"}
long_env = {**env, "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}
common = [
"--model", MODEL, "--revision", REVISION, "--dtype", "float16", "--batch_size", "32",
"--adapter", str(adapter), "--data_dir", str(data_dir),
]
for name, options in PASSES:
print(f"== pass {name}", flush=True)
subprocess.run(
[sys.executable, str(ROOT / "predict.py"), *common, "--out_dir", str(runs / name), *options],
check=True,
env=long_env if "2048" in options else env,
)
run = {name: str(runs / name) for name, _ in PASSES}
subprocess.run(
[
sys.executable, str(ROOT / "merge_votes.py"), run["orders-02"], run["orders-13"],
"--override", run["hard-2048"], "--add", run["easy-2048"],
"--add", run["medium-rev02"], "--add", run["medium-rev13"], "--out_dir", str(args.out_dir),
],
check=True,
env=env,
)
if __name__ == "__main__":
main()
|