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