Upload folder using huggingface_hub

.ipynb_checkpoints/README-checkpoint.md

@@ -0,0 +1,86 @@
+# Team QBitsLK – Sinhala MMLU Shared Task System
+
+## System summary
+- **Base model(s):** `google/gemma-4-E4B` (text-only), with Sinhala continued pretraining, then LoRA supervised fine-tuning on examination-paper MCQs.
+- **What we did:** We built a 7,872-question Sinhala multiple-choice dataset by extracting from publicly available Sri Lankan O/L and Grade 6–11 past papers, cleaning OCR artefacts and decontaminating against the released SinhalaMMLU test split. We continued pretraining Gemma 4 E4B on Sinhala text, then fine-tuned with LoRA (r=32) using answer-only loss, and merged the adapter back into the base weights. At inference each option is scored by per-option log-likelihood under a 5-prompt ensemble (Sinhala and English instructions, subject intro on/off, two option-order rotations) rather than by generation. Nothing is sampled.
+- **Total inference parameters:** **7,996,157,418** (one merged model; ≤ 8B)
+- **Extra data used for training or retrieval:** our own past-paper MCQ dataset, published at `SinhalaMMLU-SharedTask/team-qbits-mcq-data`. No retrieval, no external APIs, no closed models.
+
+## Models and resources in this submission
+| Repo | What it is |
+|---|---|
+| `SinhalaMMLU-SharedTask/team-qbits` | code, inference script, final predictions (this repo) |
+| `SinhalaMMLU-SharedTask/team-qbits-cpt-v3` | **the submitted model** — merged 16-bit, 7,996,157,418 params |
+| `SinhalaMMLU-SharedTask/team-qbits-cpt-v3-lora` | the LoRA adapter for the above (r=32), for inspection |
+| `SinhalaMMLU-SharedTask/team-qbits-base-v3b` | ablation — same recipe on the plain Gemma 4 E4B base, merged |
+| `SinhalaMMLU-SharedTask/team-qbits-base-v3b-lora` | the LoRA adapter for the ablation |
+| `SinhalaMMLU-SharedTask/team-qbits-mcq-data` | the SFT dataset (prompt/completion files) |
+
+`inference.py` uses **only** `team-qbits-cpt-v3`. The other repos are provided for transparency and
+are not part of the inference pipeline.
+
+## How to run
+```bash
+pip install -r requirements.txt
+python inference.py
+```
+- **GPU used:** 1x NVIDIA A40 48 GB (any GPU with ≥ 20 GB should work; the model is 16 GB in bf16)
+- **Approximate run time on the test set:** ~18 minutes for 2,238 questions (5 prompt variants × 4 options = ~44,800 scored candidates); peak VRAM ~27 GB
+- **Anything else the organizers need to know:**
+  - The Gemma license must be accepted on the Hugging Face account used to run this. Our merged model is self-contained, so no separate base-model download is required.
+  - `inference.py` scores options by log-likelihood and never calls `generate`, so there is no sampling. With the fixed seed, repeated runs produce byte-identical `predictions.jsonl`.
+  - `VARIANTS` at the top of `inference.py` defines the prompt ensemble. **Changing that list changes the predictions**; it is set exactly as used for our final Codabench submission.
+  - The reference `util_prompt.py` hard-codes the *hard* instruction as "1, 2, 3, 4, 5", which was correct for the original five-option hard split. SinhalaMMLU 2.0 reduced hard questions to four options, so `fit_numbers()` derives the number list from the actual option count. Verified: 0 of 2,238 prompts disagree with their option count.
+
+## Results (Codabench)
+
+| system | easy | medium | hard | **macro** |
+|---|---|---|---|---|
+| **submitted** — `team-qbits-cpt-v3`, 5-prompt ensemble | 58.10 | 62.56 | 52.79 | **57.81** |
+| the same model with the adapter applied at load time instead of merged | 58.26 | 62.02 | 52.58 | 57.62 |
+| ablation — `team-qbits-base-v3b` (plain Gemma 4 base), 5-prompt ensemble | 57.01 | 63.10 | 51.34 | 57.15 |
+| ablation — plain base, single prompt, no ensemble | — | — | — | 55.66 |
+
+`inference.py` reproduces the submitted row exactly: it loads the merged model, which is the system
+we ship. Row 2 is the same fine-tune evaluated with the LoRA applied at load time (base + adapter).
+That form is 8,069,557,738 parameters and therefore **exceeds the 8B cap**, so it is not the
+submitted system; it is listed only because merging is mathematically equivalent yet not
+bit-identical in bf16 (`W + BA` rounded once, versus `Wx + B(Ax)`), which resolves about 1% of
+near-tied questions differently.
+
+## Files
+- `inference.py` – the exact script that produced our final Codabench submission
+- `predictions.jsonl` – the exact file we submitted to the final leaderboard (macro 57.81),
+  byte-for-byte reproducible by running `inference.py`
+- `predictions_ablation_base_v3b.jsonl` – an earlier submission from the ablation model
+  (`team-qbits-base-v3b`, macro 57.15), included for transparency. To regenerate it, set
+  `MODEL_ID` to `SinhalaMMLU-SharedTask/team-qbits-base-v3b` and change nothing else.
+- `requirements.txt` – package versions we used
+- `training/` – training, fine-tuning and data-preparation code
+  - `sft_gemma4_mcq.py` – LoRA SFT with answer-only loss and candidate-scoring evaluation
+  - `mcq_dataset.py` – extraction cleanup, validation, decontamination, splitting, prompt building
+  - `add_new_data.py`, `normalize_subjects.py` – dataset assembly and subject-name alignment
+  - `merge_adapter.py`, `push_merged_to_hf.py` – adapter merging and publishing
+  - `eval_sinhalammlu.py` – read-only analysis; also builds the decontamination blocklist
+  - `check_env.py` – dependency verification
+  - `run_data_pipeline.sh`, `run_training.sh`, `run_inference.sh` – end-to-end drivers
+  - `REPRODUCE.md` – exact commands, expected counts and the training curve
+
+## Data and decontamination
+Source: publicly available Sri Lankan examination papers (`pastpapers.wiki`,
+`e-thaksalawa.moe.gov.lk`), O/L and Grades 6–11. We do not redistribute the source PDFs; the derived
+question set is published at `team-qbits-mcq-data`.
+
+The released `TEST/easy` split was used **only** as a decontamination blocklist. It matched and
+removed **26** questions from our training data. It was never used for training, checkpoint
+selection, prompt design or hyperparameter choice. Model selection used a held-out split of our own
+data, grouped so that no source paper appears in both training and validation.
+
+## Citation
+This work evaluates on the SinhalaMMLU benchmark:
+
+> Ashmari Pramodya, Nirasha Nelki, Heshan Shalinda, Chamila Liyanage, Yusuke Sakai, Randil
+> Pushpananda, Ruvan Weerasinghe, Hidetaka Kamigaito, and Taro Watanabe. 2025. **SinhalaMMLU: A
+> Comprehensive Benchmark for Evaluating Multitask Language Understanding in Sinhala.** In
+> *Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing*, pages
+> 32943–32961. https://aclanthology.org/2025.emnlp-main.1673/

README.md

@@ -0,0 +1,86 @@
+# Team QBitsLK – Sinhala MMLU Shared Task System
+
+## System summary
+- **Base model(s):** `google/gemma-4-E4B` (text-only), with Sinhala continued pretraining, then LoRA supervised fine-tuning on examination-paper MCQs.
+- **What we did:** We built a 7,872-question Sinhala multiple-choice dataset by extracting from publicly available Sri Lankan O/L and Grade 6–11 past papers, cleaning OCR artefacts and decontaminating against the released SinhalaMMLU test split. We continued pretraining Gemma 4 E4B on Sinhala text, then fine-tuned with LoRA (r=32) using answer-only loss, and merged the adapter back into the base weights. At inference each option is scored by per-option log-likelihood under a 5-prompt ensemble (Sinhala and English instructions, subject intro on/off, two option-order rotations) rather than by generation. Nothing is sampled.
+- **Total inference parameters:** **7,996,157,418** (one merged model; ≤ 8B)
+- **Extra data used for training or retrieval:** our own past-paper MCQ dataset, published at `SinhalaMMLU-SharedTask/team-qbits-mcq-data`. No retrieval, no external APIs, no closed models.
+
+## Models and resources in this submission
+| Repo | What it is |
+|---|---|
+| `SinhalaMMLU-SharedTask/team-qbits` | code, inference script, final predictions (this repo) |
+| `SinhalaMMLU-SharedTask/team-qbits-cpt-v3` | **the submitted model** — merged 16-bit, 7,996,157,418 params |
+| `SinhalaMMLU-SharedTask/team-qbits-cpt-v3-lora` | the LoRA adapter for the above (r=32), for inspection |
+| `SinhalaMMLU-SharedTask/team-qbits-base-v3b` | ablation — same recipe on the plain Gemma 4 E4B base, merged |
+| `SinhalaMMLU-SharedTask/team-qbits-base-v3b-lora` | the LoRA adapter for the ablation |
+| `SinhalaMMLU-SharedTask/team-qbits-mcq-data` | the SFT dataset (prompt/completion files) |
+
+`inference.py` uses **only** `team-qbits-cpt-v3`. The other repos are provided for transparency and
+are not part of the inference pipeline.
+
+## How to run
+```bash
+pip install -r requirements.txt
+python inference.py
+```
+- **GPU used:** 1x NVIDIA A40 48 GB (any GPU with ≥ 20 GB should work; the model is 16 GB in bf16)
+- **Approximate run time on the test set:** ~18 minutes for 2,238 questions (5 prompt variants × 4 options = ~44,800 scored candidates); peak VRAM ~27 GB
+- **Anything else the organizers need to know:**
+  - The Gemma license must be accepted on the Hugging Face account used to run this. Our merged model is self-contained, so no separate base-model download is required.
+  - `inference.py` scores options by log-likelihood and never calls `generate`, so there is no sampling. With the fixed seed, repeated runs produce byte-identical `predictions.jsonl`.
+  - `VARIANTS` at the top of `inference.py` defines the prompt ensemble. **Changing that list changes the predictions**; it is set exactly as used for our final Codabench submission.
+  - The reference `util_prompt.py` hard-codes the *hard* instruction as "1, 2, 3, 4, 5", which was correct for the original five-option hard split. SinhalaMMLU 2.0 reduced hard questions to four options, so `fit_numbers()` derives the number list from the actual option count. Verified: 0 of 2,238 prompts disagree with their option count.
+
+## Results (Codabench)
+
+| system | easy | medium | hard | **macro** |
+|---|---|---|---|---|
+| **submitted** — `team-qbits-cpt-v3`, 5-prompt ensemble | 58.10 | 62.56 | 52.79 | **57.81** |
+| the same model with the adapter applied at load time instead of merged | 58.26 | 62.02 | 52.58 | 57.62 |
+| ablation — `team-qbits-base-v3b` (plain Gemma 4 base), 5-prompt ensemble | 57.01 | 63.10 | 51.34 | 57.15 |
+| ablation — plain base, single prompt, no ensemble | — | — | — | 55.66 |
+
+`inference.py` reproduces the submitted row exactly: it loads the merged model, which is the system
+we ship. Row 2 is the same fine-tune evaluated with the LoRA applied at load time (base + adapter).
+That form is 8,069,557,738 parameters and therefore **exceeds the 8B cap**, so it is not the
+submitted system; it is listed only because merging is mathematically equivalent yet not
+bit-identical in bf16 (`W + BA` rounded once, versus `Wx + B(Ax)`), which resolves about 1% of
+near-tied questions differently.
+
+## Files
+- `inference.py` – the exact script that produced our final Codabench submission
+- `predictions.jsonl` – the exact file we submitted to the final leaderboard (macro 57.81),
+  byte-for-byte reproducible by running `inference.py`
+- `predictions_ablation_base_v3b.jsonl` – an earlier submission from the ablation model
+  (`team-qbits-base-v3b`, macro 57.15), included for transparency. To regenerate it, set
+  `MODEL_ID` to `SinhalaMMLU-SharedTask/team-qbits-base-v3b` and change nothing else.
+- `requirements.txt` – package versions we used
+- `training/` – training, fine-tuning and data-preparation code
+  - `sft_gemma4_mcq.py` – LoRA SFT with answer-only loss and candidate-scoring evaluation
+  - `mcq_dataset.py` – extraction cleanup, validation, decontamination, splitting, prompt building
+  - `add_new_data.py`, `normalize_subjects.py` – dataset assembly and subject-name alignment
+  - `merge_adapter.py`, `push_merged_to_hf.py` – adapter merging and publishing
+  - `eval_sinhalammlu.py` – read-only analysis; also builds the decontamination blocklist
+  - `check_env.py` – dependency verification
+  - `run_data_pipeline.sh`, `run_training.sh`, `run_inference.sh` – end-to-end drivers
+  - `REPRODUCE.md` – exact commands, expected counts and the training curve
+
+## Data and decontamination
+Source: publicly available Sri Lankan examination papers (`pastpapers.wiki`,
+`e-thaksalawa.moe.gov.lk`), O/L and Grades 6–11. We do not redistribute the source PDFs; the derived
+question set is published at `team-qbits-mcq-data`.
+
+The released `TEST/easy` split was used **only** as a decontamination blocklist. It matched and
+removed **26** questions from our training data. It was never used for training, checkpoint
+selection, prompt design or hyperparameter choice. Model selection used a held-out split of our own
+data, grouped so that no source paper appears in both training and validation.
+
+## Citation
+This work evaluates on the SinhalaMMLU benchmark:
+
+> Ashmari Pramodya, Nirasha Nelki, Heshan Shalinda, Chamila Liyanage, Yusuke Sakai, Randil
+> Pushpananda, Ruvan Weerasinghe, Hidetaka Kamigaito, and Taro Watanabe. 2025. **SinhalaMMLU: A
+> Comprehensive Benchmark for Evaluating Multitask Language Understanding in Sinhala.** In
+> *Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing*, pages
+> 32943–32961. https://aclanthology.org/2025.emnlp-main.1673/

inference.py

@@ -0,0 +1,277 @@
+"""
+Sinhala MMLU Shared Task - Final inference script
+Team: QBitsLK
+
+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
+
+System: one merged 16-bit model (7,996,157,418 parameters, <= 8B). Gemma 4 E4B with Sinhala
+continued pretraining, LoRA-fine-tuned on past-paper MCQs and merged back into the base weights.
+
+Answers are chosen by per-option log-likelihood under a 5-prompt ensemble. No text is generated,
+nothing is sampled, and the seed is fixed, so a re-run gives identical answers.
+"""
+
+# ======================= 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-qbits-cpt-v3"   # merged model, self-contained
+SEED = 42
+BATCH = 16                 # candidate sequences scored per forward pass
+MAX_SEQ_LEN = 1024
+
+# Prompt variants the ensemble averages over. Each is a configuration the organizers' own
+# reference harness supports (sinhala/english x subject-intro on/off), plus two rotations of the
+# option order to cancel answer-position bias. CHANGING THIS LIST CHANGES THE PREDICTIONS.
+VARIANTS = [
+    {"lang": "si", "intro": True,  "rot": 0},
+    {"lang": "si", "intro": False, "rot": 0},
+    {"lang": "en", "intro": True,  "rot": 0},
+    {"lang": "si", "intro": True,  "rot": 1},
+    {"lang": "si", "intro": True,  "rot": 2},
+]
+
+import json
+import math
+import os
+import random
+import re
+
+import numpy as np
+import torch
+from huggingface_hub import HfFileSystem, list_repo_files
+
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+
+
+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. OUR SYSTEM ===============
+# The two function names are kept. predict_all() returns one answer ("1".."4") per question,
+# in the same order as `rows`.
+
+# Instruction text taken verbatim from the organizers' util_prompt.py.
+SINHALA_INSTRUCTIONS = {
+    "easy": "පහත ප්‍රශ්නයට 1, 2, 3, 4 යන පිළිතුරුවලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරන්න.\n\n",
+    "medium": "පහත ප්‍රශ්නයට 1, 2, 3, 4  යන පිළිතුරුවලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරාගන්න.\n\n",
+    "hard": "පහත ප්‍රශ්නයට 1, 2, 3, 4, 5 යන පිළිතුරුවලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරාගන්න.\n\n",
+}
+ENGLISH_INSTRUCTIONS = {
+    "easy": "Choose the correct or most appropriate answer from 1, 2, 3, 4 for the question below.\n\n",
+    "medium": "Choose the correct or most appropriate answer from 1, 2, 3, 4 for the question below.\n\n",
+    "hard": "Select the correct answer from 1, 2, 3, 4, 5 for the given question.\n\n",
+}
+
+
+def fit_numbers(instruction, n):
+    """The reference prompt hard-codes "1, 2, 3, 4, 5" for the HARD level because the original
+    SinhalaMMLU hard split had five options. SinhalaMMLU 2.0 drops one incorrect option, so hard
+    questions have four. Telling the model to choose from five when four are shown is strictly
+    worse, so the number run is rewritten to match the actual option count."""
+    return re.sub(r"1,\s*2,\s*3,\s*4(?:,\s*5)?",
+                  ", ".join(str(i + 1) for i in range(n)), instruction, count=1)
+
+
+def build_prompt(level, subject, subject_original, question, choices, spec):
+    """Returns (prompt, rot). Displayed position j holds original option (j + rot) % n."""
+    n = len(choices)
+    rot = spec["rot"] % n if n else 0
+    shown = choices[rot:] + choices[:rot]
+    body = "\n".join(f"{i + 1}. {str(c).strip()}" for i, c in enumerate(shown))
+    if spec["lang"] == "en":
+        intro = (f"This is a multiple choice question related to the subject {subject}. "
+                 "Given in Sinhala Language.\n") if spec["intro"] else ""
+        instr = fit_numbers(ENGLISH_INSTRUCTIONS[level], n)
+        return intro + instr + f"Question: {question}\n{body}\nAnswer:", rot
+    intro = f"මෙය {subject_original} විෂයයට අදාළ බහුවරණ ප්‍රශ්නයකි.\n" if spec["intro"] else ""
+    instr = fit_numbers(SINHALA_INSTRUCTIONS[level], n)
+    return intro + instr + f"ප්‍රශ්නය: {question}\n{body}\nපිළිතුර:", rot
+
+
+def load_system():
+    """Loads the merged model. Gemma 4 is a multimodal architecture, so the right AutoClass
+    depends on the transformers version; we try the plausible ones in order."""
+    try:
+        from unsloth import FastModel
+        model, proc = FastModel.from_pretrained(model_name=MODEL_ID, dtype=None,
+                                                max_seq_length=MAX_SEQ_LEN,
+                                                load_in_4bit=False, full_finetuning=False)
+        tok = getattr(proc, "tokenizer", proc)
+        if getattr(tok, "pad_token", None) is None:
+            tok.pad_token = tok.eos_token
+        model.eval()
+        print("[model] loaded with unsloth.FastModel (the loader used for our submission)")
+        return tok, model
+    except Exception as e:
+        print(f"[model] unsloth unavailable or failed ({type(e).__name__}: {e});"
+              " falling back to plain transformers - expect ~1% of answers to differ")
+
+    from transformers import AutoConfig, AutoProcessor, AutoTokenizer
+
+    try:
+        proc = AutoProcessor.from_pretrained(MODEL_ID)
+        tok = getattr(proc, "tokenizer", proc)
+    except Exception:
+        tok = AutoTokenizer.from_pretrained(MODEL_ID)
+    if getattr(tok, "pad_token", None) is None:
+        tok.pad_token = tok.eos_token
+
+    cfg = AutoConfig.from_pretrained(MODEL_ID)
+    print("[model] architectures:", getattr(cfg, "architectures", None))
+
+    candidates = []
+    import transformers
+    for name in ("AutoModelForCausalLM", "AutoModelForImageTextToText",
+                 "AutoModelForVision2Seq", "AutoModel"):
+        if hasattr(transformers, name):
+            candidates.append(getattr(transformers, name))
+    last = None
+    for cls in candidates:
+        try:
+            model = cls.from_pretrained(MODEL_ID, dtype=torch.bfloat16, device_map="auto")
+            print(f"[model] loaded with {cls.__name__}")
+            break
+        except Exception as e:                                           # noqa: BLE001
+            last = e
+            print(f"[model] {cls.__name__} failed: {type(e).__name__}")
+    else:
+        raise RuntimeError(f"could not load {MODEL_ID}: {last}")
+    model.eval()
+    return tok, model
+
+
+def _common_prefix_len(a, b):
+    n = 0
+    for x, y in zip(a, b):
+        if x != y:
+            break
+        n += 1
+    return n
+
+
+def predict_all(system, rows):
+    """For every (question, variant, option) the answer token is scored; within each variant the
+    option scores become a log-softmax over displayed positions (putting variants of different
+    prompt lengths on one scale), are mapped back to the original option order, and summed."""
+    tok, model = system
+    dev = next(model.parameters()).device
+    pad = tok.pad_token_id if tok.pad_token_id is not None else 0
+
+    items, meta = [], {}
+    for ri, r in enumerate(rows):
+        choices = [c for c in (r.get("choices") or []) if str(c).strip()]
+        if len(choices) < 2:
+            continue
+        level = get_level(r)
+        if level not in SINHALA_INSTRUCTIONS:
+            level = "easy"
+        md = r.get("metadata") or {}
+        for vi, spec in enumerate(VARIANTS):
+            prompt, rot = build_prompt(level, r.get("subject"),
+                                       md.get("subject_original") or r.get("subject"),
+                                       str(r.get("question") or "").strip(), choices, spec)
+            meta[(ri, vi)] = (rot, len(choices))
+            p_ids = tok(prompt)["input_ids"]
+            for j in range(len(choices)):
+                f = tok(prompt + f" {j + 1}")["input_ids"]
+                if len(f) > MAX_SEQ_LEN:
+                    f = f[-MAX_SEQ_LEN:]
+                items.append((ri, vi, j, f, max(1, _common_prefix_len(p_ids, f))))
+
+    print(f"{len(rows)} questions x {len(VARIANTS)} variants = {len(items)} candidates")
+    scores = {}
+    with torch.no_grad():
+        for i in range(0, len(items), BATCH):
+            chunk = items[i:i + BATCH]
+            L = max(len(x[3]) for x in chunk)
+            ids = torch.full((len(chunk), L), pad, dtype=torch.long)
+            mask = torch.zeros((len(chunk), L), dtype=torch.long)
+            for b, (_, _, _, sq, _) in enumerate(chunk):
+                ids[b, :len(sq)] = torch.tensor(sq)
+                mask[b, :len(sq)] = 1
+            out = model(input_ids=ids.to(dev), attention_mask=mask.to(dev))
+            logits = out.logits if hasattr(out, "logits") else out[0]
+            for b, (ri, vi, j, sq, k) in enumerate(chunk):
+                lp = logits[b, k - 1:len(sq) - 1].float().log_softmax(-1)
+                tgt = ids[b, k:len(sq)].to(dev)
+                scores[(ri, vi, j)] = lp.gather(-1, tgt[:, None]).sum().item()
+            del logits, out
+            print(f"{min(i + BATCH, len(items))}/{len(items)}", end="\r")
+
+    answers = []
+    for ri, r in enumerate(rows):
+        n = len([c for c in (r.get("choices") or []) if str(c).strip()])
+        if n < 2:
+            answers.append("1")                 # a prediction is still required
+            continue
+        total, used = [0.0] * n, 0
+        for vi in range(len(VARIANTS)):
+            if (ri, vi) not in meta:
+                continue
+            rot, nv = meta[(ri, vi)]
+            raw = [scores.get((ri, vi, j)) for j in range(nv)]
+            if any(x is None for x in raw):
+                continue
+            m = max(raw)
+            z = m + math.log(sum(math.exp(x - m) for x in raw))
+            for j in range(nv):
+                total[(j + rot) % n] += raw[j] - z
+            used += 1
+        answers.append("1" if used == 0 else str(1 + max(range(n), key=lambda i: total[i])))
+    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)

requirements.txt

@@ -0,0 +1,15 @@
+unsloth
+# Sinhala MMLU Shared Task - Team QBitsLK
+# Verified install: pip install -r requirements.txt && python inference.py
+#
+# Inference needs ONE GPU with >= 20 GB of memory (the model is 16 GB in bf16).
+# Reference machine: 1x NVIDIA A40 48 GB, driver 595.91, CUDA 13.2, Python 3.12.
+
+torch>=2.6
+transformers==5.5.0
+tokenizers>=0.22.0,<=0.23.0
+huggingface_hub>=1.5.0,<2.0
+accelerate>=1.0
+timm            # Gemma 4 ships a multimodal processor; loading it requires timm
+numpy
+hf_transfer     # faster weight download; optional

training/REPRODUCE.md

@@ -0,0 +1,172 @@
+# Reproducing our SinhalaMMLU Shared Task submission
+
+Team: QBitsLK · ICTer 2026 SinhalaMMLU Shared Task
+
+Everything below runs with no human in the loop. Inference is deterministic: no sampling, no
+generation, fixed seeds, and the prediction is an argmax over summed log-softmax scores.
+
+## 1. Quick start — reproduce the submitted predictions
+
+```bash
+pip install unsloth
+pip install --no-deps transformers==5.5.0 "tokenizers>=0.22.0,<=0.23.0"
+pip install "huggingface_hub>=1.5.0,<2.0"
+pip install --no-deps --upgrade timm
+pip install hf_transfer
+
+python check_env.py                 # verifies the whole stack; exits non-zero on a problem
+
+bash run_inference.sh /path/to/test_root ./out
+```
+
+`test_root` is the directory holding `easy/`, `medium/` and `hard/`. The script writes
+`out/predictions.jsonl` and `out/predictions.zip` (with `predictions.jsonl` at the archive root).
+
+Hardware used for every reported run: one **NVIDIA A40, 48 GB**, driver 595.91, CUDA 13.2.
+Inference over the 2,238-question test set takes about **18 minutes** and peaks around 27 GB of VRAM.
+
+## 2. The submitted system in one line
+
+| | |
+|---|---|
+| model | `QBitsLK/gemma4-e4b-sinhala-mcq-cpt-v3` (merged 16-bit) |
+| parameters | **7,996,156,448** — the entire inference pipeline, one model, ≤ 8B |
+| base | `QBitsLK/gemma4-e4b-sinhala-cpt` (Gemma 4 E4B + Sinhala continued pretraining) |
+| adapter | LoRA r=32, α=32, merged into the base; merging adds no parameters |
+| scoring | per-option log-likelihood, 5-prompt ensemble, argmax |
+| prompt variants | `si_intro si_nointro en_intro si_rot1 si_rot2` |
+| seed | 3407 |
+
+The parameter count can be checked without downloading the weights:
+
+```bash
+python -c "
+from huggingface_hub import get_safetensors_metadata
+t = sum(get_safetensors_metadata('QBitsLK/gemma4-e4b-sinhala-mcq-cpt-v3').parameter_count.values())
+print(f'{t:,}', 'OK, <= 8B' if t <= 8_000_000_000 else 'OVER 8B')"
+```
+
+**The adapter is merged on purpose.** Base + adapter loaded separately is 8,069,556,768 parameters,
+which exceeds the cap; merging folds the LoRA into the existing weight matrices and returns the model
+to 7,996,156,448.
+
+## 3. Results reproduced by this package
+
+Codabench, hidden test set (2,238 questions: easy 642, medium 1,111, hard 485):
+
+| run | easy | medium | hard | **macro** |
+|---|---|---|---|---|
+| **submitted** — CPT + v3 data, 5-prompt ensemble | 58.26 | 62.02 | **52.58** | **57.62** |
+| ablation — base Gemma 4 + v3 data, 5-prompt ensemble | 57.01 | 63.10 | 51.34 | 57.15 |
+| ablation — base Gemma 4 + v3 data, single prompt | — | — | — | 55.66 |
+
+To reproduce the first ablation, change `MODEL` in `run_inference.sh` to
+`QBitsLK/gemma4-e4b-sinhala-mcq-base-v3b`. For the second, set `VARIANTS="si_intro"`.
+
+## 4. Full reproduction from scratch
+
+### 4.1 Rebuild the training data (CPU, a few minutes)
+
+```bash
+bash run_data_pipeline.sh
+```
+
+Requires the raw extraction files in `./data` (see section 5). Expected output: **7,872** questions
+after deduplication, split into **7,151 train → 21,235 examples** and **721 validation**.
+
+The public `TEST/easy` split is downloaded by this script and used **only** as a decontamination
+blocklist. It removed **26 questions** that our past-paper extractions shared with the released test
+data. It is never used for training, checkpoint selection, prompt design or hyperparameter choice.
+
+### 4.2 Re-run the fine-tune (1 GPU, 88 minutes)
+
+```bash
+bash run_training.sh /workspace/outputs/mcq_sft_cpt_v3
+```
+
+This downloads the published SFT files from
+`https://huggingface.co/datasets/QBitsLK/sinhala-mcq-sft-v3`, so section 4.1 is optional if you only
+want to reproduce training.
+
+Expected validation curve (721 held-out questions, accuracy):
+
+| step | 0 | 150 | 300 | 450 | 600 | 664 | 750 | 900 | 1050 | 1200 | 1328 |
+|---|---|---|---|---|---|---|---|---|---|---|---|
+| acc | .2316 | .5520 | .5395 | .5728 | .5506 | .5562 | .5867 | .5645 | .5742 | .5825 | **.5811** |
+
+**The submitted checkpoint is `final_adapter` (step 1328), not `best_adapter`.** Rationale in
+`SYSTEM_DESCRIPTION.md` §5.2: on a parallel run the validation-best checkpoint scored 0.7 points
+*worse* on held-out test data than the final one, because our validation split shares the training
+distribution and does not detect this failure mode.
+
+### 4.3 Publish the merged model
+
+```bash
+export HF_TOKEN=...
+python push_merged_to_hf.py \
+  --adapter /workspace/outputs/mcq_sft_cpt_v3/final_adapter \
+  --repo <your-namespace>/gemma4-e4b-sinhala-mcq-cpt-v3
+```
+
+Streams the merge directly to the Hub, so no 16 GB local copy is needed, and reads the parameter
+count back for the ≤8B check.
+
+## 5. Files
+
+| file | purpose |
+|---|---|
+| `run_inference.sh` | reproduces the submitted predictions (the critical path) |
+| `run_training.sh` | reproduces the fine-tune |
+| `run_data_pipeline.sh` | rebuilds the SFT dataset from raw extractions |
+| `make_predictions.py` | prompt builder, ensemble scorer, submission writer, format checks |
+| `sft_gemma4_mcq.py` | LoRA SFT with answer-only loss; per-step candidate-scoring evaluation |
+| `mcq_dataset.py` | extraction cleanup, validation, decontamination, splitting, prompt building |
+| `add_new_data.py` | merges per-subject extractions into the dataset |
+| `normalize_subjects.py` | aligns subject names with the official `metadata.subject_original` |
+| `eval_sinhalammlu.py` | read-only analysis; also builds the decontamination blocklist |
+| `push_merged_to_hf.py` | merges the adapter and publishes it |
+| `check_env.py` | verifies the dependency stack before a run |
+| `requirements.txt` | exact pins and install order |
+
+**Not included:** the source examination PDFs. They are third-party past papers obtained from
+publicly available sources (`pastpapers.wiki`, `e-thaksalawa.moe.gov.lk`), and we do not hold
+redistribution rights. The derived question set — questions, options and answer keys — is published
+at `QBitsLK/sinhala-mcq-sft-v3`, which is sufficient to reproduce training and inference. The raw
+per-subject extraction files are included so `run_data_pipeline.sh` runs end-to-end.
+
+## 6. Determinism
+
+- No sampling anywhere; `model.generate` is never called on the submission path.
+- `random.seed(3407)` and `torch.manual_seed(3407)`.
+- The prediction is `argmax` over log-softmax sums; ties resolve to the lowest option index.
+- Prompt variants are explicit in `run_inference.sh`; changing them changes the predictions.
+- `make_predictions.py` refuses to write the zip unless the line count equals the question count,
+  every line carries exactly `{q_no, subject, category, level, answer}`, `answer` is a `"1"`–`"5"`
+  string, and every `level` is `easy`/`medium`/`hard`.
+
+One deliberate deviation from the organisers' reference code: `evaluate.py` hard-codes the *hard*
+instruction as "1, 2, 3, 4, 5" because the original SinhalaMMLU hard split had five options.
+SinhalaMMLU 2.0 drops one incorrect option, so hard questions now have four. `fit_numbers()` in
+`make_predictions.py` rewrites the number list to match the actual option count (verified: 0 of 2,238
+prompts disagree). Without this, all 485 hard questions would be asked to choose from five options
+when only four are shown.
+
+## 7. Citation
+
+This work evaluates on the SinhalaMMLU benchmark:
+
+```bibtex
+@inproceedings{pramodya-etal-2025-sinhalammlu,
+    title = "SinhalaMMLU: A Comprehensive Benchmark for Evaluating Multitask Language Understanding in Sinhala",
+    author = "Pramodya, Ashmari and Nelki, Nirasha and Shalinda, Heshan and Liyanage, Chamila and
+              Sakai, Yusuke and Pushpananda, Randil and Weerasinghe, Ruvan and
+              Kamigaito, Hidetaka and Watanabe, Taro",
+    booktitle = "Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing",
+    year = "2025",
+    pages = "32943--32961",
+    url = "https://aclanthology.org/2025.emnlp-main.1673/"
+}
+```
+
+Verify this against the BibTeX on <https://sinhalammlu.ucsc.cmb.ac.lk/task.html> before submitting;
+it is reproduced here from that page and the arXiv preprint (arXiv:2509.03162).

training/add_new_data.py

@@ -0,0 +1,174 @@
+#!/usr/bin/env python
+"""
+Add newly extracted subject files to the MCQ dataset, safely.
+
+Why a script instead of `cat a b > c`: on Windows PowerShell the `>` redirect re-encodes the text and
+corrupts Sinhala. This reads and writes UTF-8 explicitly and validates every line.
+
+It figures out each input file's format on its own:
+  * raw extraction format  (original_question / answer_1..5 / correct_answer)  -> runs mcq_dataset
+    convert on it (same cleaning, junk-stripping and filters as the original pipeline)
+  * MCQ schema             (question / choices / answer)                        -> passed through
+It then merges everything into one converted file, ready for clean -> dedupe -> split -> build-sft.
+
+    python add_new_data.py --base data/converted.jsonl \
+        --new data/art.jsonl data/buddhim.jsonl data/dancing.jsonl data/sinhala.jsonl \
+        --out data/all_converted.jsonl
+
+Prints a per-file report (rows kept, subjects, exam types/grades) so you can see what you actually
+added before spending GPU time. Put mcq_dataset.py in the same folder.
+"""
+import argparse
+import collections
+import importlib.util
+import json
+import os
+import sys
+from argparse import Namespace
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+
+
+def load_mcq_module():
+    path = os.path.join(HERE, "mcq_dataset.py")
+    if not os.path.exists(path):
+        sys.exit(f"mcq_dataset.py not found next to this script ({HERE})")
+    spec = importlib.util.spec_from_file_location("mcq_dataset", path)
+    m = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(m)
+    return m
+
+
+def read_jsonl(path):
+    rows, bad = [], 0
+    with open(path, encoding="utf-8-sig") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                rows.append(json.loads(line))
+            except json.JSONDecodeError:
+                bad += 1
+    if bad:
+        print(f"  [warn] {bad} unparseable line(s) skipped in {os.path.basename(path)}")
+    return rows
+
+
+def detect_format(rows):
+    if not rows:
+        return "empty"
+    r = rows[0]
+    if "choices" in r and "answer" in r and "question" in r:
+        return "mcq"
+    if "original_question" in r or "answer_1" in r or "correct_answer" in r:
+        return "raw"
+    return "unknown"
+
+
+def describe(rows, label):
+    subj = collections.Counter(r.get("subject") or "?" for r in rows)
+    lvl = collections.Counter()
+    for r in rows:
+        md = r.get("metadata") or {}
+        lvl[str(md.get("exam_type") or r.get("level") or "?") + "/" +
+            str(md.get("grade") or r.get("grade") or "?")] += 1
+    nch = collections.Counter(len(r.get("choices") or []) for r in rows)
+    print(f"  {label}: {len(rows)} rows")
+    print(f"    subjects : {dict(subj.most_common(8))}")
+    print(f"    level/grade: {dict(lvl.most_common(8))}")
+    print(f"    n_choices: {dict(sorted(nch.items()))}")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--base", default="data/converted.jsonl",
+                    help="existing converted file to append to ('' = start fresh)")
+    ap.add_argument("--new", nargs="+", required=True, help="newly extracted subject files")
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--tmp_dir", default="", help="where to write per-file converted output "
+                                                  "(default: alongside each input)")
+    ap.add_argument("--keep_latin", action="store_true")
+    ap.add_argument("--latin_allow", nargs="*", default=[])
+    ap.add_argument("--min_confidence", choices=["low", "medium", "high"], default="low")
+    args = ap.parse_args()
+    for stream in (sys.stdout, sys.stderr):
+        try:
+            stream.reconfigure(encoding="utf-8")
+        except Exception:                                                # noqa: BLE001
+            pass
+
+    mcq = load_mcq_module()
+    merged, seen_ids = [], set()
+
+    if args.base:
+        if not os.path.exists(args.base):
+            sys.exit(f"--base {args.base} not found (pass --base '' to start fresh)")
+        base_rows = read_jsonl(args.base)
+        if detect_format(base_rows) != "mcq":
+            sys.exit(f"--base {args.base} is not in MCQ schema; pass the CONVERTED file")
+        print(f"[base] {args.base}")
+        describe(base_rows, "existing")
+        merged.extend(base_rows)
+
+    for path in args.new:
+        if not os.path.exists(path):
+            print(f"[skip] {path} not found")
+            continue
+        rows = read_jsonl(path)
+        fmt = detect_format(rows)
+        print(f"\n[new] {os.path.basename(path)}  format={fmt}")
+        if fmt == "raw":
+            out_dir = args.tmp_dir or os.path.dirname(path) or "."
+            os.makedirs(out_dir, exist_ok=True)
+            stem = os.path.splitext(os.path.basename(path))[0]
+            conv = os.path.join(out_dir, f"{stem}_conv.jsonl")
+            rej = os.path.join(out_dir, f"{stem}_latin_rejects.jsonl")
+            mcq.cmd_convert(Namespace(inp=path, out=conv, min_confidence=args.min_confidence,
+                                      skip_needs_review=False, keep_shared_stem=False,
+                                      keep_latin=args.keep_latin, latin_allow=args.latin_allow,
+                                      rejects=rej))
+            rows = read_jsonl(conv)
+        elif fmt == "mcq":
+            print("  already in MCQ schema: passed through without conversion")
+        else:
+            print(f"  [skip] unrecognised format; first keys: {sorted(rows[0])[:8] if rows else '-'}")
+            continue
+        describe(rows, "added")
+        merged.extend(rows)
+
+    # ids must be unique, or later stages (and any per-question analysis) get confused
+    out_rows, dup = [], 0
+    for i, r in enumerate(merged):
+        rid = r.get("id")
+        if rid is None or rid in seen_ids:
+            dup += 1
+            rid = f"{rid or 'row'}__{i}"
+        seen_ids.add(rid)
+        r["id"] = rid
+        out_rows.append(r)
+    if dup:
+        print(f"\n[ids] {dup} missing/duplicate id(s) were made unique")
+
+    os.makedirs(os.path.dirname(os.path.abspath(args.out)) or ".", exist_ok=True)
+    with open(args.out, "w", encoding="utf-8") as f:
+        for r in out_rows:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+
+    print(f"\n[out] {len(out_rows)} rows -> {args.out}")
+    describe(out_rows, "combined")
+    print("\nNext (run from the same folder):")
+    print(f"  python mcq_dataset.py validate {args.out}")
+    print(f"  python mcq_dataset.py clean {args.out} data/all_cleaned.jsonl --rejects data/all_clean_rejects.jsonl")
+    print("  python mcq_dataset.py dedupe data/all_cleaned.jsonl data/all_deduped.jsonl "
+          "--blocklist test_easy_block.jsonl --conflicts data/all_conflicts.jsonl")
+    print("  python mcq_dataset.py stats data/all_deduped.jsonl")
+    print("  python mcq_dataset.py split data/all_deduped.jsonl --train data/train2.jsonl "
+          "--val data/val2.jsonl --val_frac 0.1")
+    print("  python mcq_dataset.py build-sft data/train2.jsonl data/sft_train2.jsonl --shuffle_copies 2")
+    print("  python mcq_dataset.py build-sft data/val2.jsonl data/sft_val2.jsonl "
+          "--shuffle_copies 0 --strip_zwj_prob 0 --templates si_instruction")
+
+
+if __name__ == "__main__":
+    main()

training/check_env.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python
+"""
+Environment check for cpt_gemma4_sinhala.py (Unsloth + Gemma 4 E4B on an A40).
+
+    python check_env.py            # quick checks (no model download)
+    python check_env.py --model    # also loads Gemma 4 E4B, adds LoRA, runs 1 train step (~16 GB download)
+
+Exit code 0 = everything passed.
+"""
+import argparse
+import importlib
+import importlib.metadata as md
+import subprocess
+import sys
+
+OK, BAD, WARN = "[ OK ]", "[FAIL]", "[WARN]"
+failures = []
+
+
+def report(tag, msg):
+    print(f"{tag} {msg}")
+    if tag == BAD:
+        failures.append(msg)
+
+
+def ver(pkg):
+    try:
+        return md.version(pkg)
+    except md.PackageNotFoundError:
+        return None
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", action="store_true", help="load model + 1 training step")
+    ap.add_argument("--model_name", default="unsloth/gemma-4-E4B")
+    args = ap.parse_args()
+
+    print(f"Python {sys.version.split()[0]}")
+
+    # 1. GPU / torch  (import unsloth FIRST, as in the training script)
+    try:
+        import unsloth  # noqa: F401
+        report(OK, f"unsloth imported (v{ver('unsloth')})")
+    except Exception as e:  # noqa: BLE001
+        report(BAD, f"unsloth import failed: {type(e).__name__}: {e}")
+
+    try:
+        import torch
+        report(OK, f"torch {torch.__version__}, CUDA build {torch.version.cuda}")
+        if torch.cuda.is_available():
+            p = torch.cuda.get_device_properties(0)
+            report(OK, f"GPU: {p.name}, {p.total_memory / 1024**3:.1f} GB, bf16={torch.cuda.is_bf16_supported()}")
+        else:
+            report(BAD, "torch.cuda.is_available() is False (Unsloth needs a GPU)")
+    except Exception as e:  # noqa: BLE001
+        report(BAD, f"torch import failed: {e}")
+
+    # 2. Pinned versions
+    transformers_v = ver("transformers")
+    if transformers_v == "5.5.0":
+        report(OK, "transformers 5.5.0")
+    else:
+        report(BAD, f"transformers is {transformers_v}, expected 5.5.0 "
+                    "(re-run: pip install --no-deps transformers==5.5.0)")
+
+    tok_v = ver("tokenizers")
+    try:
+        from packaging.version import Version
+        good = tok_v and Version("0.22.0") <= Version(tok_v) <= Version("0.23.0")
+    except Exception:  # noqa: BLE001
+        good = False
+    report(OK if good else BAD, f"tokenizers {tok_v} (need >=0.22.0,<=0.23.0)")
+
+    hub_v = ver("huggingface_hub")
+    try:
+        from packaging.version import Version
+        good = hub_v and Version("1.5.0") <= Version(hub_v) < Version("2.0")
+    except Exception:  # noqa: BLE001
+        good = False
+    report(OK if good else BAD, f"huggingface_hub {hub_v} (need >=1.5.0,<2.0)")
+
+    # 3. Other packages the script imports
+    for pkg in ["trl", "peft", "datasets", "pandas", "accelerate", "bitsandbytes", "timm", "hf_transfer"]:
+        v = ver(pkg)
+        report(OK if v else BAD, f"{pkg} {v or 'NOT INSTALLED'}")
+
+    # 4. Imports the training script actually does
+    for stmt in ["from unsloth import FastModel",
+                 "from trl import SFTConfig, SFTTrainer",
+                 "from datasets import Dataset",
+                 "from huggingface_hub import hf_hub_download"]:
+        try:
+            exec(stmt, {})  # noqa: S102
+            report(OK, stmt)
+        except Exception as e:  # noqa: BLE001
+            report(BAD, f"{stmt} -> {type(e).__name__}: {e}")
+
+    # 5. Do the SFTConfig args used by the script exist in this TRL version?
+    try:
+        from trl import SFTConfig
+        import inspect
+        fields = set(inspect.signature(SFTConfig.__init__).parameters)
+        for a in ["max_length", "packing", "dataset_text_field", "bf16", "optim",
+                  "eval_strategy", "warmup_ratio"]:
+            report(OK if a in fields else BAD, f"SFTConfig accepts '{a}'")
+    except Exception as e:  # noqa: BLE001
+        report(WARN, f"could not inspect SFTConfig: {e}")
+
+    # 6. pip's own dependency check (informational: --no-deps installs can show conflicts)
+    out = subprocess.run([sys.executable, "-m", "pip", "check"], capture_output=True, text=True).stdout.strip()
+    if out and "No broken requirements" not in out:
+        print(f"{WARN} pip check reported:\n{out}\n       (conflicts about transformers/tokenizers are often "
+              "harmless with --no-deps; only worry if the model test below fails)")
+    else:
+        report(OK, "pip check: no broken requirements")
+
+    # 7. Optional end-to-end test
+    if args.model:
+        try:
+            import torch
+            from unsloth import FastModel
+            from datasets import Dataset
+            from trl import SFTConfig, SFTTrainer
+
+            print(f"\nLoading {args.model_name} ...")
+            model, tok = FastModel.from_pretrained(
+                model_name=args.model_name, dtype=None, max_seq_length=512,
+                load_in_4bit=False, full_finetuning=False)
+            report(OK, "model + tokenizer loaded")
+
+            model = FastModel.get_peft_model(
+                model, finetune_vision_layers=False, finetune_language_layers=True,
+                finetune_attention_modules=True, finetune_mlp_modules=True,
+                r=8, lora_alpha=8, lora_dropout=0, bias="none", use_rslora=True, random_state=3407)
+            report(OK, "LoRA attached (use_rslora accepted)")
+
+            sin = "ශ්‍රී ලංකාව ඉන්දියන් සාගරයේ පිහිටි ද්වීපයකි."
+            ids = tok(sin, add_special_tokens=False)["input_ids"]
+            print(f"       Sinhala sample -> {len(ids)} tokens for {len(sin)} chars")
+
+            ds = Dataset.from_dict({"text": [(sin + "\n") * 20] * 16})
+            cfg = SFTConfig(output_dir="/tmp/envcheck", dataset_text_field="text", max_length=512,
+                            packing=True, per_device_train_batch_size=2, gradient_accumulation_steps=1,
+                            max_steps=2, learning_rate=5e-5, bf16=True, optim="adamw_8bit",
+                            logging_steps=1, save_strategy="no", report_to="none")
+            tr = SFTTrainer(model=model, tokenizer=tok, train_dataset=ds, args=cfg)
+            tr.train()
+            report(OK, f"2 training steps ran; peak VRAM {torch.cuda.max_memory_reserved() / 1024**3:.1f} GB")
+        except Exception as e:  # noqa: BLE001
+            import traceback
+            traceback.print_exc()
+            report(BAD, f"model test failed: {type(e).__name__}: {e}")
+
+    print("\n" + ("ALL CHECKS PASSED" if not failures else f"{len(failures)} PROBLEM(S):"))
+    for f in failures:
+        print("  -", f)
+    sys.exit(1 if failures else 0)
+
+
+if __name__ == "__main__":
+    main()

training/eval_sinhalammlu.py

@@ -0,0 +1,396 @@
+#!/usr/bin/env python
+"""
+SinhalaMMLU evaluation with the ORGANISERS' own prompt (copied from their util_prompt.py).
+
+Two subcommands:
+
+  blocklist  - convert the official TEST json files into the mcq_dataset.py schema, so they can be
+               used for DECONTAMINATION:
+                   python eval_sinhalammlu.py blocklist --data /workspace/data/TEST/easy \
+                          --out test_easy_block.jsonl
+                   python mcq_dataset.py dedupe data/cleaned.jsonl data/deduped.jsonl \
+                          --blocklist test_easy_block.jsonl --conflicts data/conflicts.jsonl
+               Any non-zero "overlap with blocklist" means your training data contained test
+               questions and they have now been removed. Re-run split/build-sft after this.
+
+  eval       - score a model on those files. Reports BOTH:
+                 loglik   = pick the option whose answer-token has the highest log-probability
+                            (this is how the model was trained and how we measured validation)
+                 generate = greedy-decode after "පිළිතුර:" and read the first digit
+                            (this is what the organisers' evaluate.py does)
+               If generate is much worse than loglik, the model knows the answer but is not
+               emitting a clean digit, which would cost real marks.
+
+                   python eval_sinhalammlu.py eval --data /workspace/data/TEST/easy \
+                          --model /workspace/outputs/mcq_sft_base/best_adapter \
+                          --out_csv results_easy.csv
+
+NOTE ON THE RULES: the shared task forbids "training, tuning or prompt design on test questions in
+any form". Use this as a one-off read-only sanity check. Do NOT use it to choose a checkpoint, a
+prompt or hyperparameters: use your own held-out val.jsonl for that.
+"""
+import argparse
+import csv
+import glob
+import json
+import os
+import re
+import sys
+import time
+from collections import defaultdict
+
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+
+# --- the organisers' prompt text, verbatim from util_prompt.py -------------------------------
+SINHALA_INSTRUCTIONS = {
+    "easy": "පහත ප්‍රශ්නයට 1, 2, 3, 4 යන පිළිතුරුවලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරන්න.\n\n",
+    "medium": "පහත ප්‍රශ්නයට 1, 2, 3, 4  යන පිළිතුරුවලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරාගන්න.\n\n",
+    "hard": "පහත ප්‍රශ්නයට 1, 2, 3, 4, 5 යන පිළිතුරුවලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරාගන්න.\n\n",
+}
+
+
+def fit_numbers(instruction, n):
+    """The organisers hard-code "1, 2, 3, 4, 5" in the HARD instruction (the original SinhalaMMLU
+    hard split had five options). SinhalaMMLU 2.0 dropped one, so hard questions now have four;
+    the number run is rewritten to match the actual option count."""
+    return re.sub(r"1,\s*2,\s*3,\s*4(?:,\s*5)?",
+                  ", ".join(str(i + 1) for i in range(n)), instruction, count=1)
+
+
+def official_prompt(level, subject_original, question, choices, intro_type=True):
+    """Zero-shot Sinhala prompt, identical to generate_instruction_prompt(..., language='sinhala')."""
+    intro = f"මෙය {subject_original} විෂයයට අදාළ බහුවරණ ප්‍රශ්නයකි.\n" if intro_type else ""
+    body = "\n".join(f"{i + 1}. {str(c).strip()}" for i, c in enumerate(choices))
+    instr = fit_numbers(SINHALA_INSTRUCTIONS[level.lower()], len(choices))
+    return intro + instr + f"ප්‍රශ්නය: {question}\n{body}\nපිළිතුර:"
+
+
+def load_items(data_dir, level_override=""):
+    """Read every *.json in data_dir (each a list of question dicts)."""
+    files = sorted(glob.glob(os.path.join(data_dir, "**", "*.json"), recursive=True))
+    if not files:
+        sys.exit(f"no .json files found under {data_dir}")
+    rows = []
+    for path in files:
+        try:
+            data = json.load(open(path, encoding="utf-8"))
+        except Exception as e:                                           # noqa: BLE001
+            print(f"[warn] skipping {os.path.basename(path)}: {type(e).__name__}: {e}")
+            continue
+        if not isinstance(data, list):
+            print(f"[warn] skipping {os.path.basename(path)}: not a JSON list")
+            continue
+        for it in data:
+            md = it.get("metadata") or {}
+            level = level_override or md.get("difficulty") or "easy"
+            if level.lower() not in SINHALA_INSTRUCTIONS:
+                level = "easy"
+            ch = [c for c in (it.get("choices") or []) if str(c).strip()]
+            ans = it.get("answer")
+            if len(ch) < 2 or not isinstance(ans, int) or not (1 <= ans <= len(ch)):
+                continue
+            rows.append({"file": os.path.basename(path), "q_no": it.get("q_no"),
+                         "subject": it.get("subject"), "category": it.get("category"),
+                         "question": str(it.get("question") or "").strip(), "choices": ch,
+                         "answer": ans, "level": level.lower(),
+                         "subject_original": md.get("subject_original") or it.get("subject"),
+                         "source": md.get("source")})
+    print(f"[data] {len(rows)} questions from {len(files)} file(s) in {data_dir}")
+    return rows
+
+
+# ------------------------------------------------------------------ blocklist (no GPU needed)
+def cmd_blocklist(a):
+    rows = load_items(a.data, a.level)
+    with open(a.out, "w", encoding="utf-8") as f:
+        for i, r in enumerate(rows):
+            f.write(json.dumps({"id": f"mmlu_{r['level']}_{r['file']}_{r['q_no']}",
+                                "subject": r["subject"], "question": r["question"],
+                                "choices": r["choices"], "answer": r["answer"],
+                                "metadata": {"source": r["source"], "split": "official_test"}},
+                               ensure_ascii=False) + "\n")
+    print(f"[out] {len(rows)} questions -> {a.out}")
+    print("Now decontaminate your TRAINING data (this is the compliant use of these files):")
+    print(f"  python mcq_dataset.py dedupe data/cleaned.jsonl data/deduped.jsonl --blocklist {a.out}")
+    print("Any non-zero 'overlap with blocklist' = test questions were in your training set.")
+
+
+# ------------------------------------------------------------------ eval (needs the GPU)
+def common_prefix_len(x, y):
+    n = 0
+    for p, q in zip(x, y):
+        if p != q:
+            break
+        n += 1
+    return n
+
+
+def score_loglik(model, tok, rows, prompts, max_len, batch_seqs):
+    """Total log-prob of each candidate answer token; returns a predicted 1-based option per row."""
+    import torch
+    dev = next(model.parameters()).device
+    pad = tok.pad_token_id if tok.pad_token_id is not None else 0
+    items, skipped = [], set()
+    for ri, (r, p) in enumerate(zip(rows, prompts)):
+        p_ids = tok(p)["input_ids"]
+        seqs = []
+        for ci in range(len(r["choices"])):
+            f = tok(p + f" {ci + 1}")["input_ids"]
+            seqs.append((f, max(1, common_prefix_len(p_ids, f))))
+        if max(len(s) for s, _ in seqs) > max_len:
+            skipped.add(ri)
+            continue
+        for ci, (s, k) in enumerate(seqs):
+            items.append((ri, ci, s, k))
+
+    scores = {}
+    with torch.no_grad():
+        for i in range(0, len(items), batch_seqs):
+            chunk = items[i:i + batch_seqs]
+            L = max(len(x[2]) for x in chunk)
+            ids = torch.full((len(chunk), L), pad, dtype=torch.long)
+            mask = torch.zeros((len(chunk), L), dtype=torch.long)
+            for b, (_, _, s, _) in enumerate(chunk):
+                ids[b, :len(s)] = torch.tensor(s)
+                mask[b, :len(s)] = 1
+            logits = model(input_ids=ids.to(dev), attention_mask=mask.to(dev)).logits
+            for b, (ri, ci, s, k) in enumerate(chunk):
+                lp = logits[b, k - 1:len(s) - 1].float().log_softmax(-1)
+                tgt = ids[b, k:len(s)].to(dev)
+                scores[(ri, ci)] = lp.gather(-1, tgt[:, None]).sum().item()
+            del logits
+
+    preds = []
+    for ri, r in enumerate(rows):
+        if ri in skipped:
+            preds.append(None)
+            continue
+        cs = [scores.get((ri, ci)) for ci in range(len(r["choices"]))]
+        preds.append(None if any(c is None for c in cs)
+                     else 1 + max(range(len(cs)), key=lambda i: cs[i]))
+    return preds, len(skipped)
+
+
+def score_generate(model, tok, prompts, max_new_tokens, batch_rows):
+    """Greedy-decode and read the first digit, the way the organisers' evaluate.py does."""
+    import torch
+    dev = next(model.parameters()).device
+    pad = tok.pad_token_id if tok.pad_token_id is not None else 0
+    preds, raw = [], []
+    with torch.no_grad():
+        for i in range(0, len(prompts), batch_rows):
+            batch = prompts[i:i + batch_rows]
+            enc = [tok(p)["input_ids"] for p in batch]
+            L = max(len(e) for e in enc)
+            ids = torch.full((len(enc), L), pad, dtype=torch.long)
+            mask = torch.zeros((len(enc), L), dtype=torch.long)
+            for b, e in enumerate(enc):                      # LEFT pad so generation continues correctly
+                ids[b, L - len(e):] = torch.tensor(e)
+                mask[b, L - len(e):] = 1
+            out = model.generate(input_ids=ids.to(dev), attention_mask=mask.to(dev),
+                                 max_new_tokens=max_new_tokens, do_sample=False,
+                                 pad_token_id=pad)
+            for b in range(len(enc)):
+                text = tok.decode(out[b, L:], skip_special_tokens=True)
+                raw.append(text)
+                m = re.search(r"[1-5]", text)
+                preds.append(int(m.group()) if m else None)
+    return preds, raw
+
+
+def score_official(model, tok, prompts, n_choices, system="You must output the Answer",
+                   batch_rows=8, no_chat=False):
+    """Exact replica of the organisers' predict_classification_causal_by_number (utils.py):
+    wrap in the chat template with a system message, ONE forward pass, take the logits at the final
+    position, and argmax over the token ids of the BARE digits '1'..'N' (no leading space).
+    This is what actually decides their score - there is no generation involved."""
+    import torch
+    dev = next(model.parameters()).device
+    pad = tok.pad_token_id if tok.pad_token_id is not None else 0
+
+    mode = {"how": "raw prompt (no chat template requested)"} if no_chat else {}
+
+    def wrap(instruction):
+        if no_chat:
+            return instruction
+        msgs = [{"role": "system", "content": system}, {"role": "user", "content": instruction}]
+        try:
+            out = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+            mode.setdefault("how", "chat template, system + user roles")
+            return out
+        except Exception:
+            pass
+        try:
+            # many Gemma templates reject a system role: fold it into the first user turn
+            msgs = [{"role": "user", "content": system + "\n\n" + instruction}]
+            out = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+            mode.setdefault("how", "chat template, system folded into the user turn")
+            return out
+        except Exception:
+            # NO chat template on this tokenizer at all (base Gemma 4 has none).
+            # The organisers' evaluate.py calls apply_chat_template unconditionally, so it would
+            # CRASH here too -> scoring the raw prompt is the closest thing we can measure.
+            mode.setdefault("how", "NO CHAT TEMPLATE on this tokenizer: scored the raw prompt. "
+                                   "The organisers' evaluate.py would fail on this model as-is.")
+            return instruction
+
+    texts = [wrap(p) for p in prompts]
+    # digit token ids, exactly as they compute them
+    dig = {k: tok.encode(str(k), add_special_tokens=False)[-1] for k in range(1, 6)}
+    preds = []
+    with torch.no_grad():
+        for i in range(0, len(texts), batch_rows):
+            batch = texts[i:i + batch_rows]
+            enc = [tok(t)["input_ids"] for t in batch]          # add_special_tokens default, as they do
+            L = max(len(e) for e in enc)
+            ids = torch.full((len(enc), L), pad, dtype=torch.long)
+            mask = torch.zeros((len(enc), L), dtype=torch.long)
+            for b, e in enumerate(enc):                          # LEFT pad: last column is the last real token
+                ids[b, L - len(e):] = torch.tensor(e)
+                mask[b, L - len(e):] = 1
+            logits = model(input_ids=ids.to(dev), attention_mask=mask.to(dev)).logits[:, -1, :].float()
+            for b in range(len(enc)):
+                nc = n_choices[i + b]
+                cand = [dig[k] for k in range(1, nc + 1)]
+                row = logits[b, torch.tensor(cand, device=logits.device)]
+                preds.append(1 + int(row.argmax().item()))
+            del logits
+    print(f"[official] prompt construction: {mode.get('how', 'unknown')}")
+    return preds, texts[0]
+
+
+def report(rows, preds, label):
+    n = correct = unparsed = 0
+    by_cat, by_subj = defaultdict(lambda: [0, 0]), defaultdict(lambda: [0, 0])
+    for r, p in zip(rows, preds):
+        if p is None:
+            unparsed += 1
+            continue
+        n += 1
+        ok = int(p == r["answer"])
+        correct += ok
+        by_cat[r["category"] or "?"][0] += ok
+        by_cat[r["category"] or "?"][1] += 1
+        by_subj[r["subject"] or "?"][0] += ok
+        by_subj[r["subject"] or "?"][1] += 1
+    print(f"\n=== {label} ===")
+    print(f"accuracy {correct / max(1, n):.4f} ({correct}/{n}); no prediction: {unparsed}")
+    print("  by category:")
+    for k, (c, t) in sorted(by_cat.items()):
+        print(f"    {k:<18} {c / t:.4f}  ({c}/{t})")
+    print("  by subject:")
+    for k, (c, t) in sorted(by_subj.items()):
+        if t >= 20:
+            print(f"    {k:<18} {c / t:.4f}  ({c}/{t})")
+    return {"accuracy": correct / max(1, n), "n": n, "no_prediction": unparsed,
+            "by_category": {k: v[0] / v[1] for k, v in by_cat.items()},
+            "by_subject": {k: v[0] / v[1] for k, v in by_subj.items()}}
+
+
+def cmd_eval(a):
+    rows = load_items(a.data, a.level)
+    if a.limit:
+        rows = rows[:a.limit]
+    prompts = [official_prompt(r["level"], r["subject_original"], r["question"], r["choices"],
+                               intro_type=not a.no_intro) for r in rows]
+    print("\n--- prompt actually sent to the model (first question) ---")
+    print(prompts[0])
+    print("--- end of prompt ---\n")
+    if a.dry_run:
+        return
+
+    from unsloth import FastModel
+    print(f"[model] loading {a.model}")
+    model, processor = FastModel.from_pretrained(model_name=a.model, dtype=None,
+                                                 max_seq_length=a.max_seq_length,
+                                                 load_in_4bit=False, full_finetuning=False)
+    tok = getattr(processor, "tokenizer", processor)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    model.eval()
+
+    out = {"model": a.model, "data": a.data, "n_questions": len(rows)}
+    t0 = time.time()
+    ll_preds, skipped = score_loglik(model, tok, rows, prompts, a.max_seq_length, a.eval_batch)
+    print(f"[loglik] {time.time() - t0:.0f}s; skipped (too long) {skipped}")
+    out["loglik"] = report(rows, ll_preds, "loglik (how the model was trained/validated)")
+
+    if not a.skip_official:
+        t0 = time.time()
+        off_preds, shown = score_official(model, tok, prompts, [len(r["choices"]) for r in rows],
+                                          system=a.system, batch_rows=a.gen_batch, no_chat=a.no_chat)
+        print(f"[official] {time.time() - t0:.0f}s")
+        print("--- prompt as the organisers' harness builds it (chat template applied) ---")
+        print(shown[:900] + ("..." if len(shown) > 900 else ""))
+        print("--- end ---")
+        out["official"] = report(rows, off_preds, "OFFICIAL (bare-digit logits at last position)")
+
+    gen_preds, gen_raw = [None] * len(rows), [""] * len(rows)
+    if not a.skip_generate:
+        t0 = time.time()
+        gen_preds, gen_raw = score_generate(model, tok, prompts, a.max_new_tokens, a.gen_batch)
+        print(f"[generate] {time.time() - t0:.0f}s")
+        out["generate"] = report(rows, gen_preds, "generate (how the organisers score)")
+        bad = [t for t, p in zip(gen_raw, gen_preds) if p is None][:5]
+        if bad:
+            print(f"  examples with no digit in the output: {bad}")
+
+    if a.out_csv:
+        # columns match the organisers' get_accuracy_all.py so their aggregator can read this
+        with open(a.out_csv, "w", encoding="utf-8", newline="") as f:
+            w = csv.writer(f)
+            w.writerow(["q_no", "subject", "category", "question", "choices", "answer",
+                        "generated_text", "model_predicted_answer", "loglik_pred", "raw_output"])
+            for r, lp, gp, g in zip(rows, ll_preds, gen_preds, gen_raw):
+                w.writerow([r["q_no"], r["subject"], r["category"], r["question"],
+                            json.dumps(r["choices"], ensure_ascii=False), r["answer"],
+                            gp if gp is not None else "", gp if gp is not None else "", lp, g])
+        print(f"[out] per-question results -> {a.out_csv}")
+    if a.out_json:
+        json.dump(out, open(a.out_json, "w", encoding="utf-8"), ensure_ascii=False, indent=1)
+        print(f"[out] summary -> {a.out_json}")
+    print("\nREMINDER: these are official test questions. Do not use this score to choose a "
+          "checkpoint, prompt or hyperparameter; the task rules forbid tuning on test questions.")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    sub = ap.add_subparsers(dest="cmd", required=True)
+
+    p = sub.add_parser("blocklist", help="convert official TEST json -> jsonl for decontamination")
+    p.add_argument("--data", required=True)
+    p.add_argument("--out", required=True)
+    p.add_argument("--level", default="")
+    p.set_defaults(f=cmd_blocklist)
+
+    p = sub.add_parser("eval", help="score a model with the organisers' prompt")
+    p.add_argument("--data", required=True)
+    p.add_argument("--model", required=True, help="adapter dir, merged dir or HF repo")
+    p.add_argument("--level", default="", help="force easy/medium/hard (default: read from metadata)")
+    p.add_argument("--no_intro", action="store_true", help="drop the subject intro line (intro_type=False)")
+    p.add_argument("--max_seq_length", type=int, default=1024)
+    p.add_argument("--eval_batch", type=int, default=16)
+    p.add_argument("--gen_batch", type=int, default=8)
+    p.add_argument("--max_new_tokens", type=int, default=5)
+    p.add_argument("--skip_generate", action="store_true")
+    p.add_argument("--skip_official", action="store_true")
+    p.add_argument("--no_chat", action="store_true", help="official scoring WITHOUT the chat template")
+    p.add_argument("--system", default="You must output the Answer",
+                   help="system message the organisers' evaluate.py uses")
+    p.add_argument("--limit", type=int, default=0)
+    p.add_argument("--dry_run", action="store_true", help="print the prompt and exit (no model load)")
+    p.add_argument("--out_csv", default="")
+    p.add_argument("--out_json", default="")
+    p.set_defaults(f=cmd_eval)
+
+    a = ap.parse_args()
+    for stream in (sys.stdout, sys.stderr):
+        try:
+            stream.reconfigure(encoding="utf-8")
+        except Exception:                                                # noqa: BLE001
+            pass
+    a.f(a)
+
+
+if __name__ == "__main__":
+    main()

training/make_predictions.py

@@ -0,0 +1,487 @@
+#!/usr/bin/env python
+"""
+Produce the shared-task submission: predictions.zip containing predictions.jsonl.
+
+Required line format (one JSON object per question):
+    {"q_no":1,"subject":"Arts","category":"humanities","level":"easy","answer":"2"}
+
+`answer` is a STRING holding the 1-based option number.
+
+    python make_predictions.py \
+        --data /workspace/data/TEST \
+        --model /workspace/outputs/mcq_sft_base_v2/best_adapter \
+        --out predictions.jsonl --zip predictions.zip
+
+--data may be either a single folder of *.json files, or a root containing easy/ medium/ hard/
+subfolders; the level comes from metadata.difficulty, falling back to the folder name.
+
+Scoring: for every option, score the joint tokenisation of prompt+" N" and take the highest total
+log-probability ("loglik"). This is deterministic, always yields a valid 1..N answer, and measured
+highest on TEST/easy. Because the task is submitted as PREDICTIONS, we are free to choose the
+decoding method - we are not restricted to the organisers' bare-digit logit comparison.
+
+Reproducibility: the organisers re-run submitted code on their own hardware and only verified runs
+enter the final ranking. Nothing here samples; seeds are fixed and generation is not used. Keep this
+file, the adapter/model reference and the prompt builder together in the submitted code.
+"""
+import argparse
+import glob
+import json
+import os
+import re
+import random
+import sys
+import time
+import zipfile
+
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+
+# --- the organisers' prompt, verbatim from their util_prompt.py -------------------------------
+SINHALA_INSTRUCTIONS = {
+    "easy": "පහත ප්‍රශ්නයට 1, 2, 3, 4 යන පිළිතුරුවලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරන්න.\n\n",
+    "medium": "පහත ප්‍රශ්නයට 1, 2, 3, 4  යන පිළිතුරුවලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරාගන්න.\n\n",
+    "hard": "පහත ප්‍රශ්නයට 1, 2, 3, 4, 5 යන පිළිතුරුවලින් නිවැරදි හෝ ඉතාමත් ගැළපෙන පිළිතුර තෝරාගන්න.\n\n",
+}
+ENGLISH_INSTRUCTIONS = {
+    "easy": "Choose the correct or most appropriate answer from 1, 2, 3, 4 for the question below.\n\n",
+    "medium": "Choose the correct or most appropriate answer from 1, 2, 3, 4 for the question below.\n\n",
+    "hard": "Select the correct answer from 1, 2, 3, 4, 5 for the given question.\n\n",
+}
+LEVELS = ("easy", "medium", "hard")
+
+# Prompt variants to ensemble over. Each is a real configuration the organisers' harness supports
+# (sinhala/english x intro true/false), plus option rotations to cancel answer-position bias.
+# "rot" rotates the DISPLAYED option order; predictions are mapped back to the original indices.
+VARIANTS = {
+    "si_intro":    {"lang": "si", "intro": True,  "rot": 0},
+    "si_nointro":  {"lang": "si", "intro": False, "rot": 0},
+    "en_intro":    {"lang": "en", "intro": True,  "rot": 0},
+    "si_rot1":     {"lang": "si", "intro": True,  "rot": 1},
+    "si_rot2":     {"lang": "si", "intro": True,  "rot": 2},
+    "si_rot3":     {"lang": "si", "intro": True,  "rot": 3},
+    "en_nointro":  {"lang": "en", "intro": False, "rot": 0},
+    "en_rot1":     {"lang": "en", "intro": True,  "rot": 1},
+    "si_nointro_r2": {"lang": "si", "intro": False, "rot": 2},
+}
+DEFAULT_ENSEMBLE = ["si_intro", "si_nointro", "en_intro", "si_rot1", "si_rot2"]
+
+
+def fit_numbers(instruction, n):
+    """The organisers hard-code "1, 2, 3, 4, 5" in the HARD instruction because the original
+    SinhalaMMLU hard split had five options. SinhalaMMLU 2.0 dropped one incorrect option, so hard
+    questions now have four. Telling the model to choose from five options when four are shown is
+    strictly worse, so the number run is rewritten to match the actual count."""
+    return re.sub(r"1,\s*2,\s*3,\s*4(?:,\s*5)?",
+                  ", ".join(str(i + 1) for i in range(n)), instruction, count=1)
+
+
+def official_prompt(level, subject_original, question, choices, intro_type=True):
+    intro = f"මෙය {subject_original} විෂයයට අදාළ බහුවරණ ප්‍රශ්නයකි.\n" if intro_type else ""
+    body = "\n".join(f"{i + 1}. {str(c).strip()}" for i, c in enumerate(choices))
+    instr = fit_numbers(SINHALA_INSTRUCTIONS[level], len(choices))
+    return intro + instr + f"ප්‍රශ්නය: {question}\n{body}\nපිළිතුර:"
+
+
+def english_prompt(level, subject, question, choices, intro_type=True):
+    """The organisers' prompt with language="english" (their english_intro uses the ENGLISH subject)."""
+    intro = (f"This is a multiple choice question related to the subject {subject}. "
+             "Given in Sinhala Language.\n") if intro_type else ""
+    body = "\n".join(f"{i + 1}. {str(c).strip()}" for i, c in enumerate(choices))
+    instr = fit_numbers(ENGLISH_INSTRUCTIONS[level], len(choices))
+    return intro + instr + f"Question: {question}\n{body}\nAnswer:"
+
+
+def rotate(choices, rot):
+    """Displayed order; displayed position j holds original option (j + rot) % n."""
+    n = len(choices)
+    r = rot % n if n else 0
+    return choices[r:] + choices[:r], r
+
+
+def variant_prompt(row, spec):
+    shown, r = rotate(row["choices"], spec["rot"])
+    if spec["lang"] == "en":
+        p = english_prompt(row["level"], row["subject"], row["question"], shown, spec["intro"])
+    else:
+        p = official_prompt(row["level"], row["subject_original"], row["question"], shown, spec["intro"])
+    return p, r, len(shown)
+
+
+def content_free_prompt(row, spec, filler="N/A"):
+    """Same template and option count, but the question and every option replaced by a neutral
+    filler. Scoring the digits here measures the model's a-priori preference among option POSITIONS,
+    independent of content (contextual calibration / PMI-DC). Depends only on
+    (variant, level, subject, n_choices), so it is cached across questions."""
+    stub = dict(row)
+    stub["question"] = filler
+    stub["choices"] = [filler] * len(row["choices"])
+    return variant_prompt(stub, spec)
+
+
+def level_of(item, path, override=""):
+    """metadata.difficulty, else a top-level 'level'/'difficulty' field, else the file name or any
+    path component containing easy/medium/hard, else 'easy'."""
+    if override:
+        return override
+    md = item.get("metadata") or {}
+    for cand in (md.get("difficulty"), md.get("level"), item.get("level"), item.get("difficulty")):
+        d = str(cand or "").strip().lower()
+        if d in LEVELS:
+            return d
+    low = path.replace("\\", "/").lower()
+    # substring match so "hard.json", "test_hard.jsonl" and ".../hard/..." all work
+    hits = [lv for lv in LEVELS if lv in low]
+    if len(hits) == 1:
+        return hits[0]
+    if hits:                                  # e.g. ".../easy/medium_extra.json": trust the last one
+        return max(hits, key=low.rfind)
+    return "easy"
+
+
+def read_any(path):
+    """A JSON array, a JSONL file, or a dict wrapping the list under a common key."""
+    raw = open(path, encoding="utf-8-sig").read().strip()
+    if not raw:
+        return []
+    if raw[0] == "[":
+        return json.loads(raw)
+    if raw[0] == "{":
+        try:
+            obj = json.loads(raw)
+        except json.JSONDecodeError:
+            obj = None
+        if isinstance(obj, dict):
+            for k in ("data", "questions", "items", "rows", "test"):
+                if isinstance(obj.get(k), list):
+                    return obj[k]
+            return [obj]
+        # not one object -> treat as JSONL
+    out = []
+    for line in raw.splitlines():
+        line = line.strip()
+        if line:
+            out.append(json.loads(line))
+    return out
+
+
+def load_items(data_dir, level_override=""):
+    if os.path.isfile(data_dir):
+        files = [data_dir]
+    else:
+        files = sorted(set(glob.glob(os.path.join(data_dir, "**", "*.json"), recursive=True))
+                       | set(glob.glob(os.path.join(data_dir, "**", "*.jsonl"), recursive=True)))
+    if not files:
+        sys.exit(f"no .json/.jsonl files found under {data_dir}")
+    rows, skipped, per_file = [], 0, []
+    for path in files:
+        try:
+            data = read_any(path)
+        except Exception as e:                                           # noqa: BLE001
+            print(f"[warn] {os.path.basename(path)}: {type(e).__name__}: {e}")
+            continue
+        if not isinstance(data, list) or not data:
+            print(f"[warn] {os.path.basename(path)}: no question list found")
+            continue
+        before = len(rows)
+        for it in data:
+            md = it.get("metadata") or {}
+            ch = [c for c in (it.get("choices") or []) if str(c).strip()]
+            if len(ch) < 2:
+                skipped += 1          # never drop silently: these still need a prediction line
+            ans = it.get("answer")
+            rows.append({
+                "file": os.path.basename(path),
+                "gold": ans if isinstance(ans, int) and 1 <= ans <= max(2, len(ch)) else None,
+                "q_no": it.get("q_no"),
+                "subject": it.get("subject"),
+                "category": it.get("category"),
+                "level": level_of(it, path, level_override),
+                "question": str(it.get("question") or "").strip(),
+                "choices": ch,
+                "subject_original": md.get("subject_original") or it.get("subject"),
+            })
+        lv = {r["level"] for r in rows[before:]}
+        per_file.append((os.path.basename(path), len(rows) - before, sorted(lv)))
+    by_level = {}
+    for r in rows:
+        by_level[r["level"]] = by_level.get(r["level"], 0) + 1
+    print(f"[data] {len(rows)} questions from {len(files)} file(s); by level: {by_level}")
+    for name, n, lv in per_file:
+        print(f"    {name:<44} {n:>6} questions   level -> {','.join(lv)}")
+    missing = [lv for lv in LEVELS if lv not in by_level]
+    if missing:
+        print(f"[warn] no questions detected for: {missing}. If that is wrong, the level could not be "
+              "read from metadata.difficulty or the file name - check the table above.")
+    if skipped:
+        print(f"[warn] {skipped} question(s) have fewer than 2 usable options; they still get a line")
+    return rows
+
+
+def common_prefix_len(a, b):
+    n = 0
+    for x, y in zip(a, b):
+        if x != y:
+            break
+        n += 1
+    return n
+
+
+def predict(model, tok, rows, max_len, batch_seqs, variants, calibrate=False):
+    """Deterministic prompt ensemble. For every (question, variant) the options are scored, converted
+    to a log-softmax over the DISPLAYED positions (so variants with different prompt lengths are on a
+    common scale), mapped back to the original option order, and summed across variants.
+    No sampling, no generation: the same inputs always give the same predictions."""
+    import math
+    import torch
+    model.eval()
+    dev = next(model.parameters()).device
+    pad = tok.pad_token_id if tok.pad_token_id is not None else 0
+
+    items, meta = [], {}                 # (ri, vi, shown_pos, ids, answer_start) ; (ri,vi)->(rot,n)
+    for ri, r in enumerate(rows):
+        if len(r["choices"]) < 2:
+            continue
+        for vi, vname in enumerate(variants):
+            prompt, rot, n = variant_prompt(r, VARIANTS[vname])
+            meta[(ri, vi)] = (rot, n)
+            p_ids = tok(prompt)["input_ids"]
+            for j in range(n):
+                f = tok(prompt + f" {j + 1}")["input_ids"]
+                if len(f) > max_len:
+                    f = f[-max_len:]
+                items.append((ri, vi, j, f, max(1, common_prefix_len(p_ids, f))))
+
+    # ---- calibration priors: one short prompt per (variant, level, subject, n_choices)
+    prior_items, prior_key = [], {}
+    if calibrate:
+        for ri, r in enumerate(rows):
+            if len(r["choices"]) < 2:
+                continue
+            for vi, vname in enumerate(variants):
+                key = (vi, r["level"], r["subject_original"], len(r["choices"]))
+                prior_key[(ri, vi)] = key
+                if key in [k for _, k, _, _, _ in prior_items]:
+                    continue
+                cf, rot, n = content_free_prompt(r, VARIANTS[vname])
+                p_ids = tok(cf)["input_ids"]
+                for j in range(n):
+                    f = tok(cf + f" {j + 1}")["input_ids"]
+                    prior_items.append((len(prior_items), key, j, f, max(1, common_prefix_len(p_ids, f))))
+        # de-duplicate: keep the first occurrence of each (key, j)
+        seen, uniq = set(), []
+        for _, key, j, f, k in prior_items:
+            if (key, j) in seen:
+                continue
+            seen.add((key, j))
+            uniq.append((key, j, f, k))
+        prior_items = uniq
+        print(f"[calibrate] {len(prior_items)} content-free candidates "
+              f"({len({k for k, _, _, _ in prior_items})} distinct prompts)")
+
+    print(f"[predict] {len(rows)} questions x {len(variants)} variant(s) = {len(items)} candidates")
+    scores, t0 = {}, time.time()
+    with torch.no_grad():
+        for i in range(0, len(items), batch_seqs):
+            chunk = items[i:i + batch_seqs]
+            L = max(len(x[3]) for x in chunk)
+            ids = torch.full((len(chunk), L), pad, dtype=torch.long)
+            mask = torch.zeros((len(chunk), L), dtype=torch.long)
+            for b, (_, _, _, sq, _) in enumerate(chunk):
+                ids[b, :len(sq)] = torch.tensor(sq)
+                mask[b, :len(sq)] = 1
+            logits = model(input_ids=ids.to(dev), attention_mask=mask.to(dev)).logits
+            for b, (ri, vi, j, sq, k) in enumerate(chunk):
+                lp = logits[b, k - 1:len(sq) - 1].float().log_softmax(-1)
+                tgt = ids[b, k:len(sq)].to(dev)
+                scores[(ri, vi, j)] = lp.gather(-1, tgt[:, None]).sum().item()
+            del logits
+            if (i // max(1, batch_seqs)) % 100 == 0:
+                done = min(i + batch_seqs, len(items))
+                print(f"  scored {done}/{len(items)} ({time.time() - t0:.0f}s)", flush=True)
+
+    priors = {}
+    if calibrate and prior_items:
+        with torch.no_grad():
+            for i in range(0, len(prior_items), batch_seqs):
+                chunk = prior_items[i:i + batch_seqs]
+                L = max(len(x[2]) for x in chunk)
+                ids = torch.full((len(chunk), L), pad, dtype=torch.long)
+                mask = torch.zeros((len(chunk), L), dtype=torch.long)
+                for b, (_, _, sq, _) in enumerate(chunk):
+                    ids[b, :len(sq)] = torch.tensor(sq)
+                    mask[b, :len(sq)] = 1
+                logits = model(input_ids=ids.to(dev), attention_mask=mask.to(dev)).logits
+                for b, (key, j, sq, k) in enumerate(chunk):
+                    lp = logits[b, k - 1:len(sq) - 1].float().log_softmax(-1)
+                    tgt = ids[b, k:len(sq)].to(dev)
+                    priors[(key, j)] = lp.gather(-1, tgt[:, None]).sum().item()
+                del logits
+        print(f"[calibrate] priors computed in {time.time() - t0:.0f}s")
+
+    preds, fallback, agree = [], 0, 0
+    for ri, r in enumerate(rows):
+        n = len(r["choices"])
+        if n < 2:
+            preds.append(1)
+            fallback += 1
+            continue
+        total, used, per_variant = [0.0] * n, 0, []
+        for vi in range(len(variants)):
+            if (ri, vi) not in meta:
+                continue
+            rot, nv = meta[(ri, vi)]
+            raw = [scores.get((ri, vi, j)) for j in range(nv)]
+            if any(x is None for x in raw):
+                continue
+            m = max(raw)                                     # log-softmax over displayed positions
+            z = m + math.log(sum(math.exp(x - m) for x in raw))
+            adj = [raw[j] - z for j in range(nv)]
+            if calibrate:
+                key = prior_key.get((ri, vi))
+                pr = [priors.get((key, j)) for j in range(nv)]
+                if all(x is not None for x in pr):
+                    pm = max(pr)
+                    pz = pm + math.log(sum(math.exp(x - pm) for x in pr))
+                    adj = [adj[j] - (pr[j] - pz) for j in range(nv)]   # subtract the position prior
+            best_j = max(range(nv), key=lambda j: adj[j])
+            per_variant.append(1 + (best_j + rot) % n)
+            for j in range(nv):
+                total[(j + rot) % n] += adj[j]                 # displayed j == original (j+rot)%n
+            used += 1
+        if used == 0:
+            preds.append(1)
+            fallback += 1
+            continue
+        preds.append(1 + max(range(n), key=lambda i: total[i]))
+        if len(set(per_variant)) == 1:
+            agree += 1
+
+    if fallback:
+        print(f"[warn] {fallback} question(s) fell back to answer 1")
+    if len(variants) > 1:
+        print(f"[predict] all variants agreed on {agree}/{len(rows)} ({agree / max(1, len(rows)):.1%}) "
+              "- the rest are where the ensemble changes the answer")
+    print(f"[predict] done in {time.time() - t0:.0f}s")
+    return preds, None
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", required=True, help="folder of test *.json files, or a root with easy/medium/hard")
+    ap.add_argument("--model", required=True, help="adapter dir, merged dir, or HF repo")
+    ap.add_argument("--out", default="predictions.jsonl")
+    ap.add_argument("--zip", default="predictions.zip", help="'' to skip zipping")
+    ap.add_argument("--level", default="", choices=["", *LEVELS], help="force the level for every item")
+    ap.add_argument("--max_seq_length", type=int, default=1024)
+    ap.add_argument("--eval_batch", type=int, default=16)
+    ap.add_argument("--calibrate", action="store_true",
+                    help="contextual calibration: measure the model's a-priori preference among "
+                         "option POSITIONS on a content-free prompt and subtract it. ~50 extra "
+                         "forward passes in total; standard fix for log-likelihood MCQ position bias.")
+    ap.add_argument("--score", action="store_true",
+                    help="if the input carries gold 'answer' values (e.g. your val3.jsonl), report "
+                         "accuracy as well. Use this to compare variant sets on HELD-OUT data "
+                         "instead of spending Codabench submissions.")
+    ap.add_argument("--ensemble", nargs="*", default=None,
+                    help=f"prompt variants to average over (default: {' '.join(DEFAULT_ENSEMBLE)}). "
+                         f"Choices: {' '.join(VARIANTS)}. Pass 'si_intro' alone for single-prompt scoring.")
+    ap.add_argument("--seed", type=int, default=3407)
+    ap.add_argument("--dry_run", action="store_true", help="show the prompt and exit, no model load")
+    args = ap.parse_args()
+    for stream in (sys.stdout, sys.stderr):
+        try:
+            stream.reconfigure(encoding="utf-8")
+        except Exception:                                                # noqa: BLE001
+            pass
+
+    random.seed(args.seed)
+    variants = args.ensemble if args.ensemble else list(DEFAULT_ENSEMBLE)
+    unknown = [v for v in variants if v not in VARIANTS]
+    if unknown:
+        sys.exit(f"unknown variant(s) {unknown}; choose from {list(VARIANTS)}")
+    print(f"[ensemble] {len(variants)} variant(s): {variants}")
+    rows = load_items(args.data, args.level)
+
+    if args.dry_run:
+        r = next((x for x in rows if len(x["choices"]) >= 2), None)
+        if r:
+            for v in variants:
+                prompt, rot, _ = variant_prompt(r, VARIANTS[v])
+                print(f"\n--- variant {v} (rot={rot}) ---")
+                print(prompt)
+            print("--- end ---")
+        return
+
+    import torch
+    torch.manual_seed(args.seed)
+    from unsloth import FastModel
+    print(f"[model] loading {args.model}")
+    model, processor = FastModel.from_pretrained(model_name=args.model, dtype=None,
+                                                 max_seq_length=args.max_seq_length,
+                                                 load_in_4bit=False, full_finetuning=False)
+    tok = getattr(processor, "tokenizer", processor)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+
+    preds, _ = predict(model, tok, rows, args.max_seq_length, args.eval_batch, variants,
+                       calibrate=args.calibrate)
+
+    with open(args.out, "w", encoding="utf-8") as f:
+        for r, p in zip(rows, preds):
+            f.write(json.dumps({"q_no": r["q_no"], "subject": r["subject"],
+                                "category": r["category"], "level": r["level"],
+                                "answer": str(p)}, ensure_ascii=False) + "\n")
+    print(f"[out] {len(rows)} lines -> {args.out}")
+
+    # self-checks: the submission is rejected or mis-scored if any of these are wrong
+    lines = [json.loads(l) for l in open(args.out, encoding="utf-8") if l.strip()]
+    assert len(lines) == len(rows), "line count does not match the question count"
+    bad_keys = [i for i, d in enumerate(lines)
+                if set(d) != {"q_no", "subject", "category", "level", "answer"}]
+    bad_ans = [i for i, d in enumerate(lines)
+               if not (isinstance(d["answer"], str) and d["answer"] in {"1", "2", "3", "4", "5"})]
+    bad_lvl = [i for i, d in enumerate(lines) if d["level"] not in LEVELS]
+    for name, bad in [("wrong key set", bad_keys), ("answer not a '1'..'5' string", bad_ans),
+                      ("level not easy/medium/hard", bad_lvl)]:
+        print(f"[check] {name}: {len(bad)}" + (f"  e.g. line {bad[0] + 1}" if bad else "  OK"))
+    dist = {}
+    for d in lines:
+        dist[d["answer"]] = dist.get(d["answer"], 0) + 1
+    print(f"[check] answer distribution: {dict(sorted(dist.items()))}")
+    dup = len(lines) - len({(d["level"], d["subject"], d["q_no"]) for d in lines})
+    print(f"[check] duplicate (level, subject, q_no) keys: {dup}")
+    if bad_keys or bad_ans or bad_lvl:
+        sys.exit("fix the problems above before submitting")
+
+    if args.score:
+        import collections
+        have = [(r, p) for r, p in zip(rows, preds) if r["gold"] is not None]
+        if not have:
+            print("[score] the input has no gold 'answer' field, nothing to score")
+        else:
+            n = sum(1 for _ in have)
+            c = sum(int(p == r["gold"]) for r, p in have)
+            print(f"\n[score] accuracy {c / n:.4f}  ({c}/{n})   variants: {variants}")
+            for key, label in (("level", "level"), ("subject", "subject")):
+                agg = collections.defaultdict(lambda: [0, 0])
+                for r, p in have:
+                    k = r.get(key) or "?"
+                    agg[k][0] += int(p == r["gold"])
+                    agg[k][1] += 1
+                print(f"  by {label}:")
+                for k, (cc, tt) in sorted(agg.items(), key=lambda x: -x[1][1]):
+                    if tt >= 20 or label == "level":
+                        print(f"    {str(k):<34} {cc / tt:.4f}  ({cc}/{tt})")
+
+    if args.zip:
+        with zipfile.ZipFile(args.zip, "w", zipfile.ZIP_DEFLATED) as z:
+            z.write(args.out, arcname="predictions.jsonl")   # must sit at the archive root
+        with zipfile.ZipFile(args.zip) as z:
+            print(f"[zip] {args.zip} contains: {z.namelist()}")
+    print("\nFirst two lines:")
+    for l in open(args.out, encoding="utf-8").readlines()[:2]:
+        print("  " + l.rstrip())
+
+
+if __name__ == "__main__":
+    main()

training/mcq_dataset.py

@@ -0,0 +1,639 @@
+#!/usr/bin/env python
+"""
+Toolkit for building a Sinhala multiple-choice fine-tuning dataset (SinhalaMMLU-style).
+
+Record schema (one JSON object per line, UTF-8). Mirrors the test file; `answer` is 1-based:
+  {"id": "ol2019_sci_p1_q07", "subject": "science", "category": "stem",
+   "question": "...", "choices": ["...", "...", "...", "..."], "answer": 2,
+   "metadata": {"source": "<paper id or URL>", "exam_type": "OL", "year": 2019, "grade": 11,
+                "difficulty": "medium", "extraction": "text_pdf|ocr|synthetic|translated"}}
+
+Commands (run in this order):
+  python mcq_dataset.py example  example.jsonl            # writes 2 dummy records to see the schema
+  python mcq_dataset.py validate raw.jsonl                # report problems, changes nothing
+  python mcq_dataset.py clean    raw.jsonl clean.jsonl --rejects rejects.jsonl
+  python mcq_dataset.py dedupe   clean.jsonl dedup.jsonl --blocklist public_test.json   # also decontaminates
+  python mcq_dataset.py split    dedup.jsonl --train train.jsonl --val val.jsonl --val_frac 0.1
+  python mcq_dataset.py stats    train.jsonl
+  python mcq_dataset.py build-sft train.jsonl sft_train.jsonl --shuffle_copies 1 --strip_zwj_prob 0.5
+  python mcq_dataset.py build-sft val.jsonl   sft_val.jsonl   --shuffle_copies 0 --strip_zwj_prob 0 --templates si_instruction
+
+The --blocklist file is ONLY used to find and drop overlapping questions. It is never copied into the output.
+"""
+import argparse
+import collections
+import hashlib
+import json
+import random
+import re
+import sys
+import unicodedata
+
+ZWJ = "‍"
+
+# options that refer to other options: never reorder these
+POSITIONAL = [r"ඉහත", r"සියල්ල", r"කිසිවක්", r"[1-5]\s*(සහ|හා|and|&)\s*[1-5]", r"\ball of the above\b", r"\bnone of\b"]
+FIGURE = [r"රූපය", r"රූප සටහන", r"රූපයේ", r"රූපවල", r"ප්රස්තාර", r"සිතියම", r"\bfigure\b", r"\bdiagram\b", r"\bgraph\b"]
+
+SI_INSTR = "පහත බහුවරණ ප්‍රශ්නයට නිවැරදි පිළිතුරේ අංකය ලියන්න."
+EN_INSTR = "Answer the following multiple-choice question. Reply with only the number of the correct option."
+
+
+# ------------------------------------------------------------------ io / helpers
+def load(path):
+    with open(path, encoding="utf-8") as f:
+        txt = f.read().strip()
+    if not txt:
+        return []
+    if txt[0] == "[":
+        return json.loads(txt)
+    return [json.loads(l) for l in txt.splitlines() if l.strip()]
+
+
+def dump(path, rows):
+    with open(path, "w", encoding="utf-8") as f:
+        for r in rows:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+
+
+def getpath(d, path, default=None):
+    for k in path.split("."):
+        if not isinstance(d, dict) or k not in d:
+            return default
+        d = d[k]
+    return d
+
+
+def strip_zwj(s):
+    return s.replace(ZWJ, "")
+
+
+def norm_key(s):
+    """Comparison key: letters/marks/digits only, no ZWJ, no spaces/punctuation, lower-case."""
+    s = unicodedata.normalize("NFC", s).lower()
+    return "".join(ch for ch in s if unicodedata.category(ch)[0] in "LMN")
+
+
+def clean_text(s):
+    s = unicodedata.normalize("NFC", str(s))
+    s = s.replace("\\r\\n", " ").replace("\\n", " ").replace("\\r", " ").replace("\\t", " ")  # literal escapes
+    # U+200C (ZWNJ) is a legacy-font-converter artifact after word-final virama; real test text has none
+    s = s.replace("​", "").replace("‌", "").replace("", "").replace("�", "")
+    s = "".join(ch if (ch in "\n\t" or unicodedata.category(ch) != "Cc") else " " for ch in s)
+    return re.sub(r"\s+", " ", s).strip()
+
+
+def strip_numbering(choices):
+    """Remove '(1) ', '2. ', '3) ' prefixes, but only if ALL choices carry the matching sequential prefix."""
+    pats = [re.compile(rf"^\s*[\(\[]?{i + 1}[\)\]\.]\s*") for i in range(len(choices))]
+    if all(p.match(c) for p, c in zip(pats, choices)):
+        return [p.sub("", c, count=1) for p, c in zip(pats, choices)]
+    return choices
+
+
+def garbage_ratio(s):
+    if not s:
+        return 0.0
+    bad = sum(1 for ch in s if "" <= ch <= "" or ch == "�" or "À" <= ch <= "ÿ")
+    return bad / len(s)
+
+
+def sinhala_ratio(s):
+    letters = [ch for ch in s if unicodedata.category(ch).startswith("L")]
+    if not letters:
+        return 1.0
+    return sum(1 for ch in letters if "඀" <= ch <= "෿") / len(letters)
+
+
+def any_match(patterns, text):
+    t = strip_zwj(text)
+    return any(re.search(p, t, flags=re.I) for p in patterns)
+
+
+def problems_of(r):
+    """Hard errors (record unusable) and soft warnings."""
+    errs, warns = [], []
+    q, ch, a = r.get("question"), r.get("choices"), r.get("answer")
+    if not isinstance(q, str) or len(q.strip()) < 5:
+        errs.append("question missing/too short")
+    if not isinstance(ch, list) or not (2 <= len(ch) <= 6) or not all(isinstance(c, str) and c.strip() for c in ch):
+        errs.append("choices must be 2-6 non-empty strings")
+    elif len({norm_key(c) for c in ch}) < len(ch):
+        errs.append("duplicate choices")
+    if not isinstance(a, int) or isinstance(a, bool):
+        errs.append("answer must be an integer (1-based)")
+    elif isinstance(ch, list) and not (1 <= a <= len(ch)):
+        errs.append("answer out of range (must be 1-based)")
+    if not r.get("subject"):
+        warns.append("no subject")
+    if not getpath(r, "metadata.source"):
+        warns.append("no metadata.source (needed for splitting and for removing overlaps later)")
+    blob = (q or "") + " " + " ".join(ch if isinstance(ch, list) else [])
+    if "\\n" in blob:
+        warns.append("literal backslash-n in text")
+    if re.search(r"<[a-zA-Z/][^>]*>", blob):
+        warns.append("HTML tags in text")
+    if garbage_ratio(blob) > 0.05:
+        errs.append("garbled text (legacy font / OCR garbage)")
+    if sinhala_ratio(blob) < 0.25:
+        warns.append("little Sinhala text (English medium or garbled?)")
+    if isinstance(q, str) and any_match(FIGURE, q):
+        warns.append("refers to a figure/diagram")
+    return errs, warns
+
+
+# ------------------------------------------------------------------ commands
+def cmd_example(a):
+    rows = [
+        {"id": "demo_1", "subject": "science", "category": "stem", "question": "උදාහරණ ප්‍රශ්නය එක කුමක්ද?",
+         "choices": ["පිළිතුර එක", "පිළිතුර දෙක", "පිළිතුර තුන", "පිළිතුර හතර"], "answer": 2,
+         "metadata": {"source": "demo_paper_A", "exam_type": "OL", "year": 2019, "grade": 11,
+                      "difficulty": "medium", "extraction": "text_pdf"}},
+        {"id": "demo_2", "subject": "science", "category": "stem", "question": "උදාහරණ ප්‍රශ්නය දෙක කුමක්ද?",
+         "choices": ["එක", "දෙක", "තුන", "ඉහත සියල්ලම"], "answer": 4,
+         "metadata": {"source": "demo_paper_B", "exam_type": "Scholarship", "year": 2018, "grade": 5,
+                      "difficulty": "easy", "extraction": "text_pdf"}},
+    ]
+    dump(a.out, rows)
+    print(f"wrote {a.out}")
+
+
+JUNK_CHARS = set("-—–~|।*!'\"“”‘’.•_:;,")
+LATIN_OK = {"UN", "UNO", "UNESCO", "NATO", "SAARC", "ASEAN", "GDP", "IT", "ICT", "TV", "DNA", "HIV", "AIDS", "WHO",
+            "CD", "DVD", "BC", "AD", "pH", "CO2", "H2O", "SMS", "PDF", "USA", "UK", "EU", "AI"}
+
+
+LATIN_OK |= {"kg", "cm", "mm", "km", "ml", "ms", "cm3", "m2", "m3", "mg", "Hz", "kHz", "MHz", "GHz", "KB", "MB", "GB", "TB",
+             "bps", "Kbps", "Mbps", "Ctrl", "Alt", "Shift", "Enter", "Esc", "Tab", "Del", "Delete", "Home", "End", "Page",
+             "key", "Key", "ASCII", "USB", "RAM", "ROM", "CPU", "ALU", "BIOS", "OS", "LAN", "WAN", "IP", "MSB", "LSB", "BCD",
+             "www", "http", "https", "HTML", "URL", "Data", "Duplex", "control", "AND", "OR", "NOT", "XOR", "NAND", "NOR",
+             "II", "III", "IV", "VI", "VII", "VIII", "IX", "Excel", "Word", "Windows", "Access", "PowerPoint", "Internet",
+             "Email", "Wi", "Fi", "pH", "DC", "AC", "LED", "LCD", "PIN", "ATM", "VAT", "PAYE", "CEO", "NGO", "SMS", "GPS"}
+
+
+def chem_like(t):
+    """Chemical formula such as HCl, NaOH, CaCO, CO, Mg: element symbols with <=2-digit counts, >=2 distinct chars
+    (so 'OO' and 'S888' are NOT accepted)."""
+    return bool(re.fullmatch(r"(?:[A-Z][a-z]?\d{0,2}){1,6}", t)) and len(set(t)) > 1
+
+
+LATIN_OK |= {"of", "to", "in", "and", "or", "the", "gov", "lk", "com", "moe", "edu", "org", "net", "www", "mol", "aq", "ppm", "kgm",
+             "Web", "Hub", "Mem", "Sri", "TT", "tt", "CC", "Cc", "Tt", "ii", "iii", "iv", "vi", "vii", "viii", "ix"}
+
+
+def looks_like_ocr_garbage(t):
+    """Real English words / acronyms / formulas are fine inside Sinhala exam text (ICT, science, URLs...).
+    Garbage looks different: tokens containing digits ('S888'), very short stray tokens ('wo', 'Qed'),
+    or letter soup with no vowel / a long consonant run."""
+    if t in LATIN_OK or chem_like(t):
+        return False
+    if re.search(r"\d", t):
+        return True
+    if len(t) <= 3:
+        return True
+    if not re.search(r"[AEIOUaeiouYy]", t):
+        return True
+    if re.search(r"[^AEIOUaeiouYy]{5,}", t):
+        return True
+    return False
+
+
+def latin_garbage_tokens(s):
+    """Latin tokens (2+ chars) that look like OCR garbage (e.g. 'S888', 'Qed', 'OO', 'wo')."""
+    return [t for t in re.findall(r"[A-Za-z][A-Za-z0-9]*", s) if len(t) >= 2 and looks_like_ocr_garbage(t)]
+
+
+def trim_tail(s):
+    """Remove OCR debris from the END of a question/option: lone punctuation, '|', '।', short quote-led tokens,
+    and Latin garbage tokens such as 'Qed,'."""
+    toks = s.split(" ")
+    while toks:
+        t = toks[-1]
+        junk_only = t == "" or all(ch in JUNK_CHARS for ch in t)
+        quote_led = t[:1] in "“”‘’'\"|।*~" and len(t) <= 3
+        latin_junk = bool(latin_garbage_tokens(t)) and len(re.sub(r"[A-Za-z0-9]", "", t)) <= 1
+        if junk_only or quote_led or latin_junk:
+            toks.pop()
+        else:
+            break
+    return " ".join(toks)
+
+
+SHARED_STEM = [r"ප්‍රශ්නවලට", r"ප්රශ්නවලට", r"ඇසුරෙන්", r"ඉහත\s*(ප්‍රකාශ|ප්රකාශ|වගුව|සටහන|කොටස|ඡේදය)", r"පහත\s*(ඡේදය|කොටස|ප්‍රකාශ|ප්රකාශ)"]
+CONF = {"low": 0, "medium": 1, "high": 2}
+
+
+def cmd_convert(a):
+    """Convert the PDF-extraction format (answer_1..answer_5 + correct_answer letter) to the MCQ schema."""
+    rows, out, rej = load(a.inp), [], collections.Counter()
+    seen = collections.Counter()
+    latin_tok, latin_rows = collections.Counter(), []
+    for r in rows:
+        conf = str(r.get("extraction_confidence") or "unknown").lower()
+        seen[("confidence", conf)] += 1
+        seen[("needs_review", bool(r.get("needs_review")))] += 1
+        qt = r.get("question_type")
+        if qt not in (None, "", "mcq"):
+            rej[f"question_type={qt}"] += 1
+            continue
+        if CONF.get(conf, 0) < CONF[a.min_confidence]:
+            rej["below min_confidence"] += 1
+            continue
+        if a.skip_needs_review and r.get("needs_review"):
+            rej["needs_review"] += 1
+            continue
+        # the training label must exist: rows with a null/empty correct_answer are unusable
+        if not str(r.get("correct_answer") or "").strip():
+            rej["no correct_answer (null)"] += 1
+            continue
+        q = clean_text(r.get("original_question") or "")
+        # OCR junk after the final question mark, e.g.  "...ද? ී i “"
+        if "?" in q:
+            head, tail = q.rsplit("?", 1)
+            if len(tail.strip()) <= 8:
+                q = head + "?"
+        q = trim_tail(q)
+        raw_choices, spill = [], False
+        filled = [i for i in range(1, 6) if r.get(f"answer_{i}") is not None and str(r.get(f"answer_{i}")).strip()]
+        last_i = filled[-1] if filled else 0
+        for i in range(1, 6):
+            v = r.get(f"answer_{i}")
+            if v is None or not str(v).strip():
+                raw_choices.append(None)
+                continue
+            lines = [ln.strip() for ln in str(v).strip().split("\n") if ln.strip()]
+            keep = [lines[0]]
+            for ln in lines[1:]:
+                if i < last_i:
+                    # another option follows, so nothing can have leaked in: a newline is a wrapped line
+                    keep.append(ln)
+                else:
+                    # last option: continuation lines are short; next-question text is long / a question / numbered
+                    looks_like_leak = (len(ln) > 80 or "?" in ln or re.match(r"^\(?\d+[\.\)]", ln)
+                                       or any_match(SHARED_STEM, ln))
+                    if looks_like_leak:
+                        spill = True
+                        break
+                    keep.append(ln)
+            first = clean_text(" ".join(keep))
+            first = trim_tail(first).rstrip(" ,;:")
+            raw_choices.append(first or None)
+        # options must be contiguous from 1
+        n = 0
+        while n < 5 and raw_choices[n]:
+            n += 1
+        if any(raw_choices[n:]):
+            rej["non-contiguous options"] += 1
+            continue
+        choices = raw_choices[:n]
+        if n < 3:
+            rej["fewer than 3 options"] += 1
+            continue
+        ca = str(r.get("correct_answer") or "").strip().upper()
+        if ca in "ABCDE" and ca:
+            ans = "ABCDE".index(ca) + 1
+        elif ca.isdigit():
+            ans = int(ca)
+        else:
+            rej["missing/invalid correct_answer"] += 1
+            continue
+        if not (1 <= ans <= n):
+            rej["correct_answer outside the options"] += 1
+            continue
+        if len(q) < 8:
+            rej["question too short"] += 1
+            continue
+        if any_match(SHARED_STEM, q) and not a.keep_shared_stem:
+            rej["depends on a shared stem/passage"] += 1
+            continue
+        if max(len(c) for c in choices) > 250:
+            rej["option too long (extraction spill)"] += 1
+            continue
+        if not a.keep_latin:
+            allow = {t for t in (a.latin_allow or [])}
+            bad = [t for t in latin_garbage_tokens(q + " " + " ".join(choices)) if t not in allow]
+            if bad:
+                rej["OCR garbage (Latin letters/digits inside Sinhala text)"] += 1
+                latin_tok.update(bad)
+                latin_rows.append({"id": r.get("id"), "subject": r.get("subject"), "tokens": bad,
+                                   "question": q, "choices": choices})
+                continue
+        out.append({
+            "id": r.get("id"), "subject": r.get("subject"), "question": q, "choices": choices, "answer": ans,
+            "metadata": {"source": r.get("source_pdf"), "page": r.get("source_page"), "year": r.get("year"),
+                         "exam_type": r.get("level"), "medium": r.get("medium"),
+                         "extraction_confidence": conf, "needs_review": bool(r.get("needs_review")),
+                         "option_spill_fixed": spill}})
+    dump(a.out, out)
+    print(f"read {len(rows)} -> kept {len(out)} ({len(out) / max(1, len(rows)):.0%}) -> {a.out}")
+    for (k, v), c in sorted(seen.items(), key=lambda kv: str(kv[0])):
+        print(f"  input {k}={v}: {c}")
+    for k, v in rej.most_common():
+        print(f"  dropped: {k}: {v}")
+    if latin_rows:
+        print("  most common Latin tokens in the rows dropped as 'OCR garbage' (legit terms like CPU/Excel/NaCl "
+              "should be allowed with --latin_allow):")
+        print("   ", ", ".join(f"{t}:{c}" for t, c in latin_tok.most_common(40)))
+        if a.rejects:
+            dump(a.rejects, latin_rows)
+            print(f"  dropped Latin rows saved to {a.rejects}")
+    if out and all(x["metadata"]["needs_review"] for x in out):
+        print("WARNING: every kept record is flagged needs_review. Label noise from extraction is the main risk: "
+              "hand-check ~100 random records (python mcq_dataset.py sample ...) before training.")
+
+
+def cmd_sample(a):
+    rows = load(a.inp)
+    random.Random(a.seed).shuffle(rows)
+    lines = []
+    for r in rows[:a.n]:
+        md = r.get("metadata") or {}
+        lines.append("-" * 60 + f" {r.get('id')}  [{md.get('extraction_confidence', '?')}]")
+        lines.append("Q: " + r["question"])
+        for i, c in enumerate(r["choices"], 1):
+            lines.append(f"  {'*' if i == r['answer'] else ' '} {i}. {c}")
+    text = "\n".join(lines) + "\n"
+    if a.out:
+        # write the file directly as UTF-8 (a PowerShell '>' redirect re-encodes and garbles Sinhala)
+        with open(a.out, "w", encoding="utf-8-sig") as f:
+            f.write(text)
+        print(f"wrote {min(a.n, len(rows))} records -> {a.out}")
+    else:
+        print(text, end="")
+
+
+def cmd_validate(a):
+    rows = load(a.inp)
+    ec, wc = collections.Counter(), collections.Counter()
+    bad = 0
+    for r in rows:
+        errs, warns = problems_of(r)
+        bad += bool(errs)
+        ec.update(errs)
+        wc.update(warns)
+    print(f"{len(rows)} records, {bad} with hard errors ({bad / max(1, len(rows)):.1%})")
+    for k, v in ec.most_common():
+        print(f"  ERROR  {v:6d}  {k}")
+    for k, v in wc.most_common():
+        print(f"  warn   {v:6d}  {k}")
+    zw = sum(ZWJ in (r.get("question", "") + "".join(r.get("choices", []) or [])) for r in rows)
+    print(f"records containing ZWJ (U+200D): {zw} ({zw / max(1, len(rows)):.1%})")
+
+
+def cmd_clean(a):
+    out, rej = [], []
+    for r in load(a.inp):
+        r = dict(r)
+        if isinstance(r.get("answer"), str) and r["answer"].strip().isdigit():
+            r["answer"] = int(r["answer"].strip())
+        if isinstance(r.get("question"), str):
+            r["question"] = clean_text(r["question"])
+        if isinstance(r.get("choices"), list):
+            ch = [clean_text(c) for c in r["choices"] if isinstance(c, str)]
+            r["choices"] = strip_numbering(ch)
+        errs, warns = problems_of(r)
+        if "refers to a figure/diagram" in warns and not a.keep_figures:
+            errs.append("needs a figure/diagram")
+        if errs:
+            rej.append({**r, "_reject": errs})
+        else:
+            out.append(r)
+    dump(a.out, out)
+    if a.rejects:
+        dump(a.rejects, rej)
+    print(f"kept {len(out)}, rejected {len(rej)}  -> {a.out}")
+    for k, v in collections.Counter(e for x in rej for e in x["_reject"]).most_common():
+        print(f"  rejected for: {k}: {v}")
+
+
+def shingles(key, n=8):
+    return {hash(key[i:i + n]) for i in range(max(1, len(key) - n + 1))}
+
+
+def rec_key(r):
+    return norm_key(r.get("question", "") + " " + " ".join(r.get("choices", [])))
+
+
+class Index:
+    def __init__(self):
+        self.post = collections.defaultdict(list)
+        self.n = 0
+
+    def add(self, sh):
+        for h in sh:
+            self.post[h].append(self.n)
+        self.n += 1
+
+    def best_containment(self, sh):
+        if not sh:
+            return 0.0
+        cnt = collections.Counter()
+        for h in sh:
+            for d in self.post.get(h, ()):
+                cnt[d] += 1
+        return max(cnt.values()) / len(sh) if cnt else 0.0
+
+
+def cmd_dedupe(a):
+    rows = load(a.inp)
+    # 1) same question text: identical answer -> keep first, conflicting answers -> drop all (likely label noise)
+    groups = collections.defaultdict(list)
+    for r in rows:
+        groups[norm_key(r["question"])].append(r)
+    kept, conflicts = [], []
+    for k, g in groups.items():
+        ans = {norm_key(x["choices"][x["answer"] - 1]) for x in g}
+        if len(ans) > 1:
+            conflicts += g
+        else:
+            kept.append(g[0])
+    dups = len(rows) - len(kept) - len(conflicts)
+    # 2) near-duplicates inside the set
+    idx, uniq, near = Index(), [], 0
+    for r in kept:
+        sh = shingles(rec_key(r))
+        if idx.best_containment(sh) >= a.near_threshold:
+            near += 1
+            continue
+        idx.add(sh)
+        uniq.append(r)
+    # 3) overlap with blocklist (e.g. public test questions): drop, never copied
+    overlap = 0
+    if a.blocklist:
+        bfull, bq = Index(), Index()
+        for p in a.blocklist:
+            for b in load(p):
+                if isinstance(b, dict) and b.get("question"):
+                    bfull.add(shingles(norm_key(b["question"] + " " + " ".join(b.get("choices", []) or []))))
+                    bq.add(shingles(norm_key(b["question"])))
+        final = []
+        for r in uniq:
+            qk = norm_key(r["question"])
+            # drop if the whole item matches, OR the question text alone matches (robust to differently formatted options)
+            hit = bfull.best_containment(shingles(rec_key(r))) >= a.block_threshold or \
+                (len(qk) >= 25 and bq.best_containment(shingles(qk)) >= 0.8)
+            if hit:
+                overlap += 1
+            else:
+                final.append(r)
+        uniq = final
+    dump(a.out, uniq)
+    if conflicts and a.conflicts:
+        dump(a.conflicts, conflicts)
+    print(f"in {len(rows)} -> out {len(uniq)} | exact dup {dups} | answer conflicts {len(conflicts)} | "
+          f"near-dup {near} | overlap with blocklist {overlap}")
+
+
+def cmd_split(a):
+    rows = load(a.inp)
+    tr, va = [], []
+    groups = set()
+    for r in rows:
+        g = str(getpath(r, a.group_key) or r.get("id", "")).strip() or "none"
+        groups.add(g)
+        h = int(hashlib.sha1(f"{a.seed}:{g}".encode()).hexdigest()[:8], 16) / 2 ** 32
+        (va if h < a.val_frac else tr).append(r)
+    dump(a.train, tr)
+    dump(a.val, va)
+    print(f"{len(groups)} groups ({a.group_key}) -> train {len(tr)}, val {len(va)}")
+    if len(groups) < 10:
+        print("WARNING: few groups; the split may be lopsided. Add metadata.source to every record.")
+
+
+def cmd_stats(a):
+    rows = load(a.inp)
+    print(f"{len(rows)} records")
+    for title, f in [("subject", lambda r: r.get("subject")), ("category", lambda r: r.get("category")),
+                     ("difficulty", lambda r: getpath(r, "metadata.difficulty")),
+                     ("exam_type", lambda r: getpath(r, "metadata.exam_type")),
+                     ("year", lambda r: getpath(r, "metadata.year")),
+                     ("n_choices", lambda r: len(r.get("choices", []))),
+                     ("answer position", lambda r: r.get("answer")),
+                     ("extraction", lambda r: getpath(r, "metadata.extraction_confidence") or getpath(r, "metadata.extraction"))]:
+        c = collections.Counter(f(r) for r in rows)
+        print(f"\n{title}:")
+        for k, v in sorted(c.items(), key=lambda kv: -kv[1])[:20]:
+            print(f"  {str(k):28s} {v:6d}  {v / len(rows):6.1%}")
+
+
+def render(r, tpl, rng):
+    ch = r["choices"]
+    letters = "ABCDEF"
+    subj = r.get("subject") or ""
+    if tpl == "bare":
+        body = r["question"] + "\n" + "\n".join(f"{i + 1}. {c}" for i, c in enumerate(ch)) + "\nපිළිතුර:"
+        return body, f" {r['answer']}"
+    if tpl == "si_instruction":
+        head = SI_INSTR + "\n\n" + (f"විෂය: {subj}\n" if subj else "")
+        body = head + f"ප්‍රශ්නය: {r['question']}\n" + "\n".join(f"{i + 1}. {c}" for i, c in enumerate(ch)) + "\nපිළිතුර:"
+        return body, f" {r['answer']}"
+    if tpl == "en_instruction":
+        head = EN_INSTR + "\n\n" + (f"Subject: {subj}\n" if subj else "")
+        body = head + f"Question: {r['question']}\n" + "\n".join(f"{i + 1}. {c}" for i, c in enumerate(ch)) + "\nAnswer:"
+        return body, f" {r['answer']}"
+    if tpl == "letters":
+        body = f"Question: {r['question']}\n" + "\n".join(f"{letters[i]}. {c}" for i, c in enumerate(ch)) + "\nAnswer:"
+        return body, f" {letters[r['answer'] - 1]}"
+    raise ValueError(tpl)
+
+
+TEMPLATE_WEIGHTS = {"si_instruction": 0.35, "bare": 0.25, "en_instruction": 0.25, "letters": 0.15}
+
+
+def shuffle_safe(r):
+    return not any(any_match(POSITIONAL, c) for c in r["choices"])
+
+
+def cmd_build_sft(a):
+    rng = random.Random(a.seed)
+    names = a.templates or list(TEMPLATE_WEIGHTS)
+    weights = [TEMPLATE_WEIGHTS[n] for n in names]
+    rows, out, n_shuf_skip = load(a.inp), [], 0
+    for r in rows:
+        variants = [r]
+        if a.shuffle_copies > 0:
+            if shuffle_safe(r):
+                for _ in range(a.shuffle_copies):
+                    perm = list(range(len(r["choices"])))
+                    rng.shuffle(perm)
+                    new_choices = [r["choices"][i] for i in perm]
+                    new_ans = perm.index(r["answer"] - 1) + 1
+                    variants.append({**r, "choices": new_choices, "answer": new_ans})
+            else:
+                n_shuf_skip += 1
+        for v in variants:
+            v = dict(v)
+            if rng.random() < a.strip_zwj_prob:
+                v["question"] = strip_zwj(v["question"])
+                v["choices"] = [strip_zwj(c) for c in v["choices"]]
+                v["subject"] = strip_zwj(v.get("subject") or "")
+            tpl = rng.choices(names, weights)[0]
+            prompt, completion = render(v, tpl, rng)
+            if a.format == "chat":
+                rec = {"prompt": [{"role": "user", "content": prompt}],
+                       "completion": [{"role": "assistant", "content": completion.strip()}]}
+            else:
+                lab = "ABCDEF" if tpl == "letters" else "123456"
+                rec = {"prompt": prompt, "completion": completion,
+                       "candidates": [f" {lab[i]}" for i in range(len(v["choices"]))],
+                       "subject": r.get("subject")}
+            out.append(rec)
+    rng.shuffle(out)
+    dump(a.out, out)
+    pos = collections.Counter(
+        (x["completion"] if isinstance(x["completion"], str) else x["completion"][0]["content"]).strip()
+        for x in out)
+    print(f"{len(rows)} questions -> {len(out)} training examples ({a.format}); "
+          f"not shuffled (positional options): {n_shuf_skip}")
+    print("answer label distribution:", dict(sorted(pos.items())))
+
+
+def main():
+    # Windows consoles/redirects default to cp1252, which cannot encode Sinhala: force UTF-8
+    for stream in (sys.stdout, sys.stderr):
+        try:
+            stream.reconfigure(encoding="utf-8")
+        except Exception:
+            pass
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = ap.add_subparsers(dest="cmd", required=True)
+    p = sub.add_parser("example"); p.add_argument("out"); p.set_defaults(f=cmd_example)
+    p = sub.add_parser("convert", help="PDF-extraction format -> MCQ schema")
+    p.add_argument("inp"); p.add_argument("out")
+    p.add_argument("--min_confidence", choices=["low", "medium", "high"], default="low")
+    p.add_argument("--skip_needs_review", action="store_true")
+    p.add_argument("--keep_shared_stem", action="store_true")
+    p.add_argument("--rejects", default="", help="save the rows dropped for Latin garbage to this file")
+    p.add_argument("--latin_allow", nargs="*", default=[], help="Latin tokens that are legitimate (e.g. CPU Excel RAM)")
+    p.add_argument("--keep_latin", action="store_true",
+                   help="do not drop records with Latin tokens (use for science/maths/English papers)")
+    p.set_defaults(f=cmd_convert)
+    p = sub.add_parser("sample", help="print random records for hand-checking"); p.add_argument("inp")
+    p.add_argument("-n", type=int, default=20); p.add_argument("--seed", type=int, default=1)
+    p.add_argument("--out", default="", help="write to this file as UTF-8 (use this instead of '>' on Windows)")
+    p.set_defaults(f=cmd_sample)
+    p = sub.add_parser("validate"); p.add_argument("inp"); p.set_defaults(f=cmd_validate)
+    p = sub.add_parser("clean"); p.add_argument("inp"); p.add_argument("out")
+    p.add_argument("--rejects", default=""); p.add_argument("--keep_figures", action="store_true"); p.set_defaults(f=cmd_clean)
+    p = sub.add_parser("dedupe"); p.add_argument("inp"); p.add_argument("out")
+    p.add_argument("--blocklist", nargs="*", default=[], help="files of questions to remove overlaps with")
+    p.add_argument("--near_threshold", type=float, default=0.9)
+    p.add_argument("--block_threshold", type=float, default=0.5)
+    p.add_argument("--conflicts", default="conflicts.jsonl"); p.set_defaults(f=cmd_dedupe)
+    p = sub.add_parser("split"); p.add_argument("inp"); p.add_argument("--train", required=True)
+    p.add_argument("--val", required=True); p.add_argument("--val_frac", type=float, default=0.1)
+    p.add_argument("--group_key", default="metadata.source"); p.add_argument("--seed", type=int, default=3407)
+    p.set_defaults(f=cmd_split)
+    p = sub.add_parser("stats"); p.add_argument("inp"); p.set_defaults(f=cmd_stats)
+    p = sub.add_parser("build-sft"); p.add_argument("inp"); p.add_argument("out")
+    p.add_argument("--format", choices=["plain", "chat"], default="plain")
+    p.add_argument("--templates", nargs="*", choices=list(TEMPLATE_WEIGHTS), default=None)
+    p.add_argument("--shuffle_copies", type=int, default=1)
+    p.add_argument("--strip_zwj_prob", type=float, default=0.5)
+    p.add_argument("--seed", type=int, default=3407); p.set_defaults(f=cmd_build_sft)
+    a = ap.parse_args()
+    a.f(a)
+
+
+if __name__ == "__main__":
+    main()

training/merge_adapter.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python
+"""
+Merge a CPT LoRA adapter into its base model and save a plain 16-bit model folder.
+
+Needed when the CPT run was uploaded as an ADAPTER (adapter_config.json + adapter_model.safetensors)
+rather than as merged weights. After merging, the folder loads like any ordinary model, so
+sft_gemma4_mcq.py can train on it with no changes:
+
+    python merge_adapter.py --adapter ovinduG/gemma4-e4b-sinhala-cpt-d3-c3ucsc \
+                            --out /workspace/models/cpt_d3c3_merged
+    python sft_gemma4_mcq.py --train sft_train.jsonl --val sft_val.jsonl \
+        --model /workspace/models/cpt_d3c3_merged --out /workspace/outputs/mcq_sft_cpt_d3c3
+
+Why merge instead of training the adapter directly: loading the adapter gives a PeftModel, and the
+SFT script then adds a SECOND LoRA on top. Saving/merging that stack is error-prone, and the
+parameter count for the 8B rule stops being obvious. Merging folds the CPT LoRA into the existing
+weight matrices, so the result is still ~7.996B parameters.
+
+Needs the GPU free (stop any training job first) and ~16 GB of free disk for the output.
+"""
+import argparse
+import json
+import os
+import struct
+import sys
+from pathlib import Path
+
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+
+LIMIT = 8_000_000_000          # shared-task parameter cap
+
+
+def adapter_info(src):
+    """Print what we are about to merge: base model and LoRA rank, read from adapter_config.json."""
+    cfg = None
+    p = Path(src) / "adapter_config.json"
+    if p.exists():
+        cfg = json.loads(p.read_text(encoding="utf-8"))
+    else:
+        try:
+            from huggingface_hub import hf_hub_download
+            cfg = json.loads(Path(hf_hub_download(src, "adapter_config.json")).read_text(encoding="utf-8"))
+        except Exception as e:                                        # noqa: BLE001
+            print(f"[warn] could not read adapter_config.json ({type(e).__name__}: {e}); "
+                  "if this is already a merged model you do not need this script")
+            return None
+    print(f"[adapter] base model : {cfg.get('base_model_name_or_path')}")
+    print(f"[adapter] LoRA r={cfg.get('r')}  alpha={cfg.get('lora_alpha')}  "
+          f"rslora={cfg.get('use_rslora')}  targets={cfg.get('target_modules')}")
+    return cfg
+
+
+def count_params(folder):
+    """Sum parameter counts from safetensors headers (no need to load the weights)."""
+    total = 0
+    for f in sorted(Path(folder).glob("*.safetensors")):
+        with open(f, "rb") as fh:
+            n = struct.unpack("<Q", fh.read(8))[0]
+            head = json.loads(fh.read(n))
+        for name, meta in head.items():
+            if name == "__metadata__":
+                continue
+            shape = meta.get("shape") or []
+            c = 1
+            for d in shape:
+                c *= d
+            total += c
+    return total
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--adapter", required=True, help="HF repo id or local dir of the CPT LoRA adapter")
+    ap.add_argument("--out", required=True, help="local folder to write the merged 16-bit model to")
+    ap.add_argument("--max_seq_length", type=int, default=2048)
+    args = ap.parse_args()
+
+    if Path(args.out).exists() and list(Path(args.out).glob("*.safetensors")):
+        print(f"[skip] {args.out} already holds safetensors; delete it to rebuild")
+        n = count_params(args.out)
+        print(f"[check] parameters: {n:,} -> {'OK, under 8B' if n <= LIMIT else 'OVER THE 8B LIMIT'}")
+        return
+
+    adapter_info(args.adapter)
+
+    from unsloth import FastModel          # imported late, as in the other scripts
+    print(f"[load] {args.adapter} (base weights download on first use, ~16 GB)")
+    model, tok = FastModel.from_pretrained(model_name=args.adapter, max_seq_length=args.max_seq_length,
+                                           dtype=None, load_in_4bit=False)
+    print(f"[merge] writing merged 16-bit model to {args.out}")
+    model.save_pretrained_merged(args.out, tok, save_method="merged_16bit")
+
+    files = sorted(p.name for p in Path(args.out).iterdir())
+    print(f"[done] {len(files)} files: {', '.join(files)}")
+    n = count_params(args.out)
+    print(f"[check] parameters: {n:,} -> {'OK, under 8B' if n <= LIMIT else 'OVER THE 8B LIMIT'}")
+    for need in ["config.json", "tokenizer.json", "tokenizer_config.json"]:
+        if need not in files:
+            print(f"[warn] {need} is missing from {args.out}; loading it may fail")
+    if not any(f.endswith(".safetensors") for f in files):
+        sys.exit("no safetensors were written: the merge failed")
+
+
+if __name__ == "__main__":
+    main()

training/normalize_subjects.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python
+"""
+Make the `subject` field consistent across the dataset.
+
+Why it matters: build-sft puts the subject INTO the prompt ("විෂය: {subject}" / "Subject: {subject}"),
+and the organisers' prompt always supplies the SINHALA subject name (metadata.subject_original, e.g.
+"චිත්‍ර" for Arts). If some training rows carry an English or short name ("Dancing", "සිංහල") the model
+sees a label at train time that never appears at test time, for exactly those subjects.
+
+Three steps:
+
+  1. see what the official test files actually use (the authoritative names):
+         python normalize_subjects.py --test_dir TEST/easy
+  2. see what your dataset currently uses:
+         python normalize_subjects.py --inp data/all_converted.jsonl --list
+  3. rewrite with the mapping below (edit MAPPING first if step 1 disagrees):
+         python normalize_subjects.py --inp data/all_converted.jsonl --out data/all_converted_norm.jsonl
+
+Nothing else in the record is touched. Use --map_file to pass your own {"from": "to"} JSON instead.
+"""
+import argparse
+import collections
+import glob
+import json
+import os
+import sys
+
+# English / short / variant name  ->  the SINHALA name the official test files use as
+# metadata.subject_original. These were read off TEST/easy, not guessed.
+MAPPING = {
+    # --- corrections to names already in the dataset -------------------------------------
+    "සිංහල භාෂාව හා සාහිත්‍යය": "සිංහල භාෂාව සහ සාහිත්‍යය",   # හා -> සහ
+    "සෞඛ්‍ය හා ශාරීරික අධ්‍යාපනය": "සෞඛ්‍ය හා ශාරීරික විද්‍යාව",  # අධ්‍යාපනය -> විද්‍යාව
+    "නැටුම්": "නර්තනය",
+    "සිංහල": "සිංහල භාෂාව සහ සාහිත්‍යය",
+    # --- English / variant labels from newer extractions ---------------------------------
+    "Arts": "චිත්‍ර",
+    "Art": "චිත්‍ර",
+    "art": "චිත්‍ර",
+    "Dancing": "නර්තනය",
+    "dancing": "නර්තනය",
+    "Buddhism": "බුද්ධ ධර්මය",
+    "Sinhala": "සිංහල භාෂාව සහ සාහිත්‍යය",
+    "Sinhala language and literature": "සිංහල භාෂාව සහ සාහිත්‍යය",
+    "Sinhala_language_and_literature": "සිංහල භාෂාව සහ සාහිත්‍යය",
+    "Eastern music": "පෙරදිග සංගීතය",
+    "Eastern_music": "පෙරදිග සංගීතය",
+    "drama and Theatre": "නාට්‍ය හා රංග කලාව",
+    "drama_and_theatre": "නාට්‍ය හා රංග කලාව",
+    "drama and theatre": "නාට්‍ය හා රංග කලාව",
+    "drama": "නාට්‍ය හා රංග කලාව",
+    "drama and theater": "නාට්‍ය හා රංග කලාව",
+    "theatre and drama": "නාට්‍ය හා රංග කලාව",
+    "eastern music": "පෙරදිග සංගීතය",
+    "Health and Physical Science": "සෞඛ්‍ය හා ශාරීරික විද්‍යාව",
+    "History": "ඉතිහාසය",
+    "Civics": "පුරවැසි අධ්‍යාපනය",
+    "Geography": "භූගෝල විද්‍යාව",
+    "Catholicism": "කතෝලික ධර්මය",
+    "Christianity": "ක්‍රිස්තියානි ධර්මය",
+    "Islam": "ඉස්ලාම් ධර්මය",
+    "science": "විද්‍යාව",
+    "Science": "විද්‍යාව",
+}
+
+
+def load_map(path):
+    """Read a {"from": "to"} JSON file, tolerating a UTF-8 BOM or UTF-16 (what Windows editors write)."""
+    raw = open(path, "rb").read()
+    if not raw.strip():
+        sys.exit(f"{path} is empty")
+    for enc in ("utf-8-sig", "utf-16", "utf-8"):
+        try:
+            txt = raw.decode(enc)
+        except (UnicodeDecodeError, UnicodeError):
+            continue
+        try:
+            m = json.loads(txt.lstrip("\ufeff"))
+        except json.JSONDecodeError:
+            continue
+        if not isinstance(m, dict):
+            sys.exit(f"{path} must contain a JSON object of {{\"from\": \"to\"}} pairs")
+        print(f"[map] {len(m)} override(s) read from {path} ({enc})")
+        return {str(k): str(v) for k, v in m.items()}
+    sys.exit(f"could not parse {path} as JSON in utf-8/utf-8-sig/utf-16. First bytes: {raw[:40]!r}")
+
+
+def read_jsonl(path):
+    rows = []
+    with open(path, encoding="utf-8-sig") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                rows.append(json.loads(line))
+    return rows
+
+
+def show_test_names(test_dir):
+    """Print the subject / subject_original pairs the official test files use."""
+    files = sorted(glob.glob(os.path.join(test_dir, "**", "*.json"), recursive=True))
+    if not files:
+        sys.exit(f"no .json files under {test_dir}")
+    c = collections.Counter()
+    for p in files:
+        try:
+            data = json.load(open(p, encoding="utf-8"))
+        except Exception as e:                                           # noqa: BLE001
+            print(f"[warn] {os.path.basename(p)}: {type(e).__name__}")
+            continue
+        if isinstance(data, list):
+            for it in data:
+                md = it.get("metadata") or {}
+                c[(it.get("subject"), md.get("subject_original"))] += 1
+    print(f"official test subjects in {test_dir} "
+          f"(subject_original is what goes into THEIR prompt):")
+    print(f"{'n':>6}  {'subject':<34}{'subject_original'}")
+    for (s, so), n in sorted(c.items(), key=lambda x: -x[1]):
+        print(f"{n:>6}  {str(s):<34}{so}")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--inp", default="")
+    ap.add_argument("--out", default="")
+    ap.add_argument("--list", action="store_true", help="just print the current subject counts")
+    ap.add_argument("--test_dir", default="", help="print the official test files' subject names")
+    ap.add_argument("--map_file", default="", help='JSON file of {"from": "to"} overrides')
+    args = ap.parse_args()
+    for stream in (sys.stdout, sys.stderr):
+        try:
+            stream.reconfigure(encoding="utf-8")
+        except Exception:                                                # noqa: BLE001
+            pass
+
+    if args.test_dir:
+        show_test_names(args.test_dir)
+        if not args.inp:
+            return
+        print()
+
+    if not args.inp:
+        sys.exit("give --inp (and --out, unless --list)")
+    rows = read_jsonl(args.inp)
+    before = collections.Counter(r.get("subject") or "?" for r in rows)
+
+    if args.list or not args.out:
+        print(f"{len(rows)} rows in {args.inp}; subjects:")
+        for s, n in before.most_common():
+            flag = "  <-- would be renamed" if s in MAPPING else ""
+            print(f"{n:>6}  {s}{flag}")
+        if not args.out:
+            return
+
+    mapping = dict(MAPPING)
+    if args.map_file:
+        mapping.update(load_map(args.map_file))
+
+    ci = {k.lower(): v for k, v in mapping.items()}
+    changed = collections.Counter()
+    for r in rows:
+        s = r.get("subject")
+        # normalise stray whitespace too: the official files contain "drama and Theatre " with a
+        # trailing space, which would otherwise be a separate bucket
+        if isinstance(s, str):
+            s2 = " ".join(s.split())
+            if s2 != s:
+                r["subject"] = s = s2
+        tgt = mapping.get(s)
+        if tgt is None and isinstance(s, str):
+            tgt = ci.get(s.lower())          # Sinhala has no case, so lowering is safe
+        if tgt and tgt != s:
+            changed[(s, tgt)] += 1
+            r["subject"] = tgt
+
+    with open(args.out, "w", encoding="utf-8") as f:
+        for r in rows:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+
+    if changed:
+        print("renamed:")
+        for (a, b), n in changed.most_common():
+            print(f"{n:>6}  {a!r} -> {b!r}")
+    else:
+        print("nothing to rename (subjects already consistent)")
+    after = collections.Counter(r.get("subject") or "?" for r in rows)
+    print(f"\n[out] {len(rows)} rows -> {args.out}")
+    print(f"distinct subjects: {len(before)} -> {len(after)}")
+
+
+if __name__ == "__main__":
+    main()

training/push_merged_to_hf.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python
+"""
+Merge a LoRA adapter into its base model and push the merged 16-bit weights STRAIGHT TO THE HUB,
+without ever writing a full ~16 GB copy to local disk.
+
+Use this when `upload_to_hf.py --make_merged` fails with:
+    RuntimeError: Unsloth: Failed saving locally - no disk space left. Uploading can work luckily!
+
+    export HF_TOKEN=hf_...
+    python push_merged_to_hf.py \
+        --adapter /workspace/outputs/mcq_sft_cpt_d3c3/best_adapter \
+        --repo    ovinduG/gemma4-e4b-sinhala-mcq
+
+Add --public to make the repo public (default: private).
+Afterwards it reads the parameter count back FROM THE HUB (no download) to check the 8B rule.
+
+The adapter's base model must still be resolvable: either the local folder recorded in
+adapter_config.json ("base_model_name_or_path") or a Hub repo id.
+"""
+import argparse
+import inspect
+import json
+import os
+import sys
+from pathlib import Path
+
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+
+LIMIT = 8_000_000_000
+
+
+def show_base(adapter):
+    """`adapter` may be a local folder OR a Hub repo id."""
+    p = Path(adapter) / "adapter_config.json"
+    if p.exists():
+        cfg = json.loads(p.read_text(encoding="utf-8"))
+    else:
+        try:
+            from huggingface_hub import hf_hub_download
+            cfg = json.loads(Path(hf_hub_download(adapter, "adapter_config.json"))
+                             .read_text(encoding="utf-8"))
+            print(f"[adapter] read adapter_config.json from the Hub repo {adapter}")
+        except Exception as e:
+            sys.exit(f"could not read adapter_config.json from {adapter} "
+                     f"({type(e).__name__}: {e}). Give a local adapter folder or a Hub repo id "
+                     "that contains adapter_config.json + adapter_model.safetensors.")
+    base = cfg.get("base_model_name_or_path")
+    print(f"[adapter] LoRA r={cfg.get('r')} alpha={cfg.get('lora_alpha')}")
+    print(f"[adapter] base model: {base}")
+    if base and ("/" in str(base)) and Path(base).is_absolute():
+        if not Path(base).exists():
+            sys.exit(f"base model folder {base} is missing: re-create it (merge_adapter.py) or edit "
+                     "base_model_name_or_path in adapter_config.json to a Hub repo id")
+        print("[adapter] base folder found locally")
+    return base
+
+
+def verify_on_hub(repo, token):
+    """Read parameter counts from the uploaded safetensors headers without downloading the weights."""
+    try:
+        from huggingface_hub import get_safetensors_metadata
+        md = get_safetensors_metadata(repo, token=token)
+        counts = getattr(md, "parameter_count", None) or {}
+        total = sum(counts.values())
+        print(f"[verify] parameter count on the Hub: {total:,} ({dict(counts)})")
+        print(f"[verify] 8B rule: {'OK, under 8,000,000,000' if total <= LIMIT else 'OVER THE LIMIT'}")
+        return total
+    except Exception as e:                                               # noqa: BLE001
+        print(f"[verify] could not read parameter count ({type(e).__name__}: {e}); "
+              "check the repo's Files tab by hand")
+        return None
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--adapter", required=True, help="LoRA adapter folder (e.g. .../best_adapter)")
+    ap.add_argument("--repo", required=True, help="target repo, e.g. ovinduG/gemma4-e4b-sinhala-mcq")
+    ap.add_argument("--max_seq_length", type=int, default=2048)
+    ap.add_argument("--public", action="store_true", help="make the repo public (default: private)")
+    ap.add_argument("--token", default="")
+    args = ap.parse_args()
+
+    token = args.token or os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+    if not token:
+        sys.exit("No token. Run:  export HF_TOKEN=hf_...   (or pass --token)")
+
+    show_base(args.adapter)
+
+    from unsloth import FastModel
+    print(f"[load] {args.adapter}")
+    model, tok = FastModel.from_pretrained(model_name=args.adapter, max_seq_length=args.max_seq_length,
+                                           dtype=None, load_in_4bit=False)
+
+    fn = getattr(model, "push_to_hub_merged", None)
+    if fn is None:
+        sys.exit("this Unsloth build has no push_to_hub_merged; upgrade unsloth, or free ~16 GB and "
+                 "use upload_to_hf.py --make_merged")
+
+    # build kwargs from whatever this Unsloth version actually accepts
+    params = set(inspect.signature(fn).parameters)
+    kw = {}
+    for name, val in [("save_method", "merged_16bit"), ("token", token),
+                      ("private", not args.public), ("tokenizer", tok)]:
+        if name in params:
+            kw[name] = val
+    print(f"[push] -> {args.repo} ({'public' if args.public else 'private'}); "
+          f"streaming shards, nothing large written locally")
+    print(f"[push] extra args accepted by this build: {sorted(kw)}")
+    try:
+        if "tokenizer" in kw:
+            fn(args.repo, **kw)
+        else:                                   # older signature takes the tokenizer positionally
+            kw.pop("tokenizer", None)
+            fn(args.repo, tok, **kw)
+    except TypeError as e:
+        print(f"[push] signature mismatch ({e}); retrying with the minimal call")
+        fn(args.repo, tok, save_method="merged_16bit", token=token)
+
+    print("[push] done")
+    verify_on_hub(args.repo, token)
+    print("\nSubmit the MERGED repo (not the adapter): merging adds no parameters, so it stays "
+          "under 8B, whereas base + a 73.4M adapter does not.")
+
+
+if __name__ == "__main__":
+    main()

training/run_data_pipeline.sh

@@ -0,0 +1,48 @@
+#!/usr/bin/env bash
+# Rebuilds the SFT dataset from the raw PDF extractions. CPU only, a few minutes.
+#
+#   bash run_data_pipeline.sh
+#
+# Inputs expected in ./data :
+#   candidates.jsonl                     bulk extraction (9,409 rows, O/L + Grade 10-11)
+#   art.jsonl buddhim.jsonl dancing.jsonl sinhala.jsonl Eastern_music.jsonl drama.jsonl
+#                                        later per-subject extractions
+# Plus the public SinhalaMMLU TEST/easy split, used ONLY as a decontamination blocklist.
+set -euo pipefail
+
+# 0. public test split -> decontamination blocklist (never used for training or tuning)
+hf download naist-nlp/SinhalaMMLU --repo-type dataset --include "TEST/easy/*" --local-dir .
+python eval_sinhalammlu.py blocklist --data TEST/easy --out test_easy_block.jsonl
+
+# 1. bulk extraction -> MCQ schema
+python mcq_dataset.py convert data/candidates.jsonl data/converted.jsonl \
+  --rejects data/latin_rejects.jsonl
+
+# 2. merge the per-subject extractions
+python add_new_data.py --base data/converted.jsonl \
+  --new data/art.jsonl data/buddhim.jsonl data/dancing.jsonl data/sinhala.jsonl \
+        data/Eastern_music.jsonl data/drama.jsonl \
+  --out data/all_converted_v3.jsonl
+
+# 3. align subject names with the official metadata.subject_original values
+python normalize_subjects.py --inp data/all_converted_v3.jsonl --out data/all_converted_norm_v3.jsonl
+
+# 4. clean, decontaminate, split
+python mcq_dataset.py validate data/all_converted_norm_v3.jsonl
+python mcq_dataset.py clean    data/all_converted_norm_v3.jsonl data/all_cleaned_v3.jsonl \
+  --rejects data/all_clean_rejects_v3.jsonl
+python mcq_dataset.py dedupe   data/all_cleaned_v3.jsonl data/all_deduped_v3.jsonl \
+  --blocklist test_easy_block.jsonl --conflicts data/all_conflicts_v3.jsonl
+python mcq_dataset.py stats    data/all_deduped_v3.jsonl
+python mcq_dataset.py split    data/all_deduped_v3.jsonl \
+  --train data/train3.jsonl --val data/val3.jsonl --val_frac 0.1 --seed 3407
+
+# 5. build the prompt/completion files
+python mcq_dataset.py build-sft data/train3.jsonl data/sft_train3.jsonl \
+  --shuffle_copies 2 --templates si_instruction bare en_instruction --seed 3407
+python mcq_dataset.py build-sft data/val3.jsonl data/sft_val3.jsonl \
+  --shuffle_copies 0 --strip_zwj_prob 0 --templates si_instruction --seed 3407
+
+echo
+echo "== expected: 7,872 questions after dedupe (26 removed as test overlaps),"
+echo "==           7,151 train -> 21,235 examples, 721 validation."

training/run_inference.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+# Reproduces the submitted predictions end-to-end. No human in the loop.
+#
+#   bash run_inference.sh /path/to/test_root [output_dir]
+#
+# test_root is the directory holding easy/ medium/ hard/ (or any folder of the test .json files).
+# Writes predictions.jsonl and predictions.zip into output_dir (default: ./out).
+#
+# Deterministic: no sampling, no generation, fixed seed, argmax over summed log-softmax scores.
+# Total inference pipeline = one 7,996,156,448-parameter model (<= 8B).
+set -euo pipefail
+
+TEST_ROOT="${1:?usage: bash run_inference.sh <test_root> [out_dir]}"
+OUT_DIR="${2:-./out}"
+MODEL="QBitsLK/gemma4-e4b-sinhala-mcq-cpt-v3"
+VARIANTS="si_intro si_nointro en_intro si_rot1 si_rot2"
+
+mkdir -p "$OUT_DIR"
+echo "== model    : $MODEL"
+echo "== variants : $VARIANTS"
+echo "== test root: $TEST_ROOT"
+
+python make_predictions.py \
+  --data "$TEST_ROOT" \
+  --model "$MODEL" \
+  --ensemble $VARIANTS \
+  --seed 3407 \
+  --out "$OUT_DIR/predictions.jsonl" \
+  --zip "$OUT_DIR/predictions.zip"
+
+echo
+echo "== wrote $OUT_DIR/predictions.zip (contains predictions.jsonl at the archive root)"

training/run_training.sh

@@ -0,0 +1,34 @@
+#!/usr/bin/env bash
+# Reproduces the fine-tuning run behind the submitted model.
+#
+#   bash run_training.sh [out_dir]
+#
+# Needs: one 48 GB GPU (reference: NVIDIA A40), ~17 GB disk for the base weights.
+# Runtime on the reference machine: 88 minutes, peak VRAM 28.0 GB.
+#
+# The SFT files are published at:
+#   https://huggingface.co/datasets/QBitsLK/sinhala-mcq-sft-v3
+# and are downloaded here. run_data_pipeline.sh rebuilds them from the raw extractions.
+set -euo pipefail
+
+OUT_DIR="${1:-/workspace/outputs/mcq_sft_cpt_v3}"
+BASE_MODEL="QBitsLK/gemma4-e4b-sinhala-cpt"
+DATA_REPO="QBitsLK/sinhala-mcq-sft-v3"
+
+hf download "$DATA_REPO" --repo-type dataset --local-dir ./sft_data
+TRAIN=./sft_data/sft_train3.jsonl
+VAL=./sft_data/sft_val3.jsonl
+wc -l "$TRAIN" "$VAL"      # expect 21235 and 721
+
+python sft_gemma4_mcq.py \
+  --train "$TRAIN" --val "$VAL" \
+  --model "$BASE_MODEL" \
+  --out "$OUT_DIR" \
+  --lora_r 32 --lora_alpha 32 --lr 1e-4 \
+  --epochs 2 --batch_size 8 --grad_accum 4 \
+  --max_seq_length 1024 --warmup_steps 10 \
+  --eval_steps 150 --seed 3407
+
+echo
+echo "== the submitted checkpoint is $OUT_DIR/final_adapter (step 1328), NOT best_adapter."
+echo "== see SYSTEM_DESCRIPTION.md section 5.2 for why."

training/sft_gemma4_mcq.py

@@ -0,0 +1,282 @@
+#!/usr/bin/env python
+"""
+Supervised fine-tuning of Gemma 4 E4B for Sinhala multiple-choice QA (SinhalaMMLU-style), Unsloth + LoRA.
+
+Input: the prompt/completion files made by `mcq_dataset.py build-sft` (plain format, with a `candidates` field).
+The loss is computed ONLY on the answer tokens. Accuracy is measured separately by scoring every answer
+option on the held-out validation file (before training = zero-shot baseline, then after every epoch).
+
+  # base model
+  python sft_gemma4_mcq.py --train sft_train.jsonl --val sft_val.jsonl \
+      --model unsloth/gemma-4-E4B --out /workspace/outputs/mcq_sft_base
+  # CPT model (local merged folder from the earlier run)
+  python sft_gemma4_mcq.py --train sft_train.jsonl --val sft_val.jsonl \
+      --model /workspace/outputs/gemma4_e4b_sinhala_cpt/merged_16bit --out /workspace/outputs/mcq_sft_cpt
+  # zero-shot baseline only (no training)
+  python sft_gemma4_mcq.py --val sft_val.jsonl --model unsloth/gemma-4-E4B --eval_only
+
+Outputs: <out>/best_adapter (best validation accuracy), <out>/final_adapter, <out>/history.json
+To merge + upload the best adapter:  python upload_to_hf.py --adapter_dir <out>/best_adapter --make_merged ...
+"""
+import argparse
+import inspect
+import json
+import os
+import random
+import sys
+import time
+
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+
+
+# ------------------------------------------------------------------ helpers (no GPU needed)
+def load_rows(path):
+    rows = [json.loads(l) for l in open(path, encoding="utf-8") if l.strip()]
+    for r in rows[:3]:
+        if not isinstance(r.get("prompt"), str) or not isinstance(r.get("candidates"), list):
+            sys.exit(f"{path}: need plain prompt/completion rows with a 'candidates' list "
+                     "(re-run mcq_dataset.py build-sft without --format chat)")
+    return rows
+
+
+def common_prefix_len(a, b):
+    n = 0
+    for x, y in zip(a, b):
+        if x != y:
+            break
+        n += 1
+    return n
+
+
+def encode_example(tok, prompt, completion, add_eos=True):
+    """Tokenise prompt+completion jointly (as a test harness would); mask everything before the first
+    token that differs from the prompt-only tokenisation. Returns input_ids, labels."""
+    p_ids = tok(prompt)["input_ids"]
+    f_ids = tok(prompt + completion)["input_ids"]
+    k = max(1, common_prefix_len(p_ids, f_ids))
+    ids = list(f_ids) + ([tok.eos_token_id] if add_eos and tok.eos_token_id is not None else [])
+    labels = [-100] * k + ids[k:]
+    return ids, labels
+
+
+def evaluate(model, tok, rows, max_len=1024, limit=0, batch_seqs=16, verbose=True):
+    """Accuracy of picking the option whose answer tokens get the highest total log-probability.
+    Every (question, candidate) pair is scored; pairs are batched together for speed."""
+    import torch
+    was_training = model.training
+    model.eval()
+    dev = next(model.parameters()).device
+    pad = tok.pad_token_id if tok.pad_token_id is not None else 0
+    use = rows[:limit] if limit else rows
+    t0 = time.time()
+
+    items, skipped = [], 0            # (row_idx, cand_idx, token_ids, answer_start)
+    for ri, r in enumerate(use):
+        p_ids = tok(r["prompt"])["input_ids"]
+        seqs, starts = [], []
+        for c in r["candidates"]:
+            f = tok(r["prompt"] + c)["input_ids"]
+            seqs.append(f)
+            starts.append(max(1, common_prefix_len(p_ids, f)))
+        if max(len(x) for x in seqs) > max_len:
+            skipped += 1
+            continue
+        for ci, (sq, k) in enumerate(zip(seqs, starts)):
+            items.append((ri, ci, sq, k))
+
+    scores = {}
+    with torch.no_grad():
+        for i in range(0, len(items), batch_seqs):
+            chunk = items[i:i + batch_seqs]
+            L = max(len(x[2]) for x in chunk)
+            ids = torch.full((len(chunk), L), pad, dtype=torch.long)
+            mask = torch.zeros((len(chunk), L), dtype=torch.long)
+            for b, (_, _, sq, _) in enumerate(chunk):
+                ids[b, :len(sq)] = torch.tensor(sq)
+                mask[b, :len(sq)] = 1
+            logits = model(input_ids=ids.to(dev), attention_mask=mask.to(dev)).logits
+            for b, (ri, ci, sq, k) in enumerate(chunk):
+                lp = logits[b, k - 1:len(sq) - 1].float().log_softmax(-1)
+                tgt = ids[b, k:len(sq)].to(dev)
+                scores[(ri, ci)] = lp.gather(-1, tgt[:, None]).sum().item()
+            del logits
+
+    n = correct = 0
+    by_subj, gold_pos = {}, {}
+    for ri, r in enumerate(use):
+        cs = [scores.get((ri, ci)) for ci in range(len(r["candidates"]))]
+        if any(c is None for c in cs):
+            continue
+        pred = max(range(len(cs)), key=lambda i: cs[i])
+        gold = r["candidates"].index(r["completion"])
+        n += 1
+        correct += int(pred == gold)
+        gold_pos[gold] = gold_pos.get(gold, 0) + 1
+        sj = r.get("subject") or "?"
+        c_, t_ = by_subj.get(sj, (0, 0))
+        by_subj[sj] = (c_ + int(pred == gold), t_ + 1)
+
+    if was_training:
+        model.train()
+    acc = correct / max(1, n)
+    majority = max(gold_pos.values()) / max(1, n) if gold_pos else 0.0
+    if verbose:
+        print(f"[eval] accuracy {acc:.4f} ({correct}/{n}); always-guess-most-common-position {majority:.4f}; "
+              f"skipped(too long) {skipped}; {time.time() - t0:.0f}s", flush=True)
+    return {"acc": acc, "n": n, "majority_baseline": majority,
+            "by_subject": {k: round(v[0] / v[1], 4) for k, v in sorted(by_subj.items()) if v[1] >= 20}}
+
+
+# ------------------------------------------------------------------ main
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--train", default="")
+    ap.add_argument("--val", required=True)
+    ap.add_argument("--model", default="unsloth/gemma-4-E4B")
+    ap.add_argument("--out", default="/workspace/outputs/mcq_sft")
+    ap.add_argument("--max_seq_length", type=int, default=1024)
+    ap.add_argument("--lora_r", type=int, default=32)
+    ap.add_argument("--lora_alpha", type=int, default=32)
+    ap.add_argument("--lr", type=float, default=1e-4)
+    ap.add_argument("--epochs", type=float, default=3.0)
+    ap.add_argument("--max_steps", type=int, default=-1, help=">0 overrides epochs (smoke test)")
+    ap.add_argument("--batch_size", type=int, default=8)
+    ap.add_argument("--grad_accum", type=int, default=4)
+    ap.add_argument("--warmup_steps", type=int, default=10)
+    ap.add_argument("--eval_limit", type=int, default=0, help="evaluate on the first N val rows only (0 = all)")
+    ap.add_argument("--eval_batch", type=int, default=16, help="candidate sequences scored per forward pass")
+    ap.add_argument("--eval_steps", type=int, default=0,
+                    help="also evaluate every N optimizer steps (0 = only at epoch ends). "
+                         "Use it when the model converges early: the best checkpoint may be mid-epoch.")
+    ap.add_argument("--skip_baseline", action="store_true")
+    ap.add_argument("--eval_only", action="store_true", help="just score the model on --val (zero-shot baseline)")
+    ap.add_argument("--seed", type=int, default=3407)
+    args = ap.parse_args()
+    if not args.eval_only and not args.train:
+        sys.exit("--train is required unless --eval_only")
+
+    from unsloth import FastModel  # must come before transformers / trl / peft
+    import torch
+    from datasets import Dataset
+    from transformers import DataCollatorForSeq2Seq, TrainerCallback
+    from trl import SFTConfig, SFTTrainer
+
+    random.seed(args.seed)
+    os.makedirs(args.out, exist_ok=True)
+    json.dump(vars(args), open(os.path.join(args.out, "run_config.json"), "w"), indent=1)
+    print(f"[model] loading {args.model}")
+    model, processor = FastModel.from_pretrained(model_name=args.model, dtype=None,
+                                                 max_seq_length=args.max_seq_length,
+                                                 load_in_4bit=False, full_finetuning=False)
+    tok = getattr(processor, "tokenizer", processor)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+
+    val_rows = load_rows(args.val)
+    print(f"[data] val rows: {len(val_rows)}")
+
+    if args.eval_only:
+        res = evaluate(model, tok, val_rows, args.max_seq_length, args.eval_limit, args.eval_batch)
+        json.dump(res, open(os.path.join(args.out, "eval_only.json"), "w"), ensure_ascii=False, indent=1)
+        return
+
+    model = FastModel.get_peft_model(
+        model, finetune_vision_layers=False, finetune_language_layers=True,
+        finetune_attention_modules=True, finetune_mlp_modules=True,
+        r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=0, bias="none", random_state=args.seed,
+        use_gradient_checkpointing="unsloth")
+    try:
+        model.print_trainable_parameters()
+    except Exception:
+        pass
+
+    train_rows = load_rows(args.train)
+    enc, dropped = [], 0
+    for r in train_rows:
+        ids, labels = encode_example(tok, r["prompt"], r["completion"])
+        if len(ids) > args.max_seq_length:
+            dropped += 1
+            continue
+        enc.append({"input_ids": ids, "attention_mask": [1] * len(ids), "labels": labels})
+    print(f"[data] train examples: {len(enc)} (dropped {dropped} longer than {args.max_seq_length} tokens)")
+    train_ds = Dataset.from_list(enc)
+
+    coll_kw = dict(padding=True, label_pad_token_id=-100, pad_to_multiple_of=8)
+    coll_params = inspect.signature(DataCollatorForSeq2Seq.__init__).parameters
+    if "tokenizer" in coll_params:
+        collator = DataCollatorForSeq2Seq(tokenizer=tok, **coll_kw)
+    else:                                   # transformers >=5 renamed it
+        collator = DataCollatorForSeq2Seq(processing_class=tok, **coll_kw)
+    # --- masking sanity check: only the answer tokens may carry a label
+    b = collator([enc[0], enc[min(1, len(enc) - 1)]])
+    frac = float((b["labels"] != -100).float().mean())
+    shown = tok.decode([t for t in b["labels"][0].tolist() if t != -100])
+    print(f"[mask] labelled-token fraction {frac:.3f}; example 0 trains on: {shown!r}")
+    if frac > 0.3:
+        sys.exit("Masking looks wrong (too many labelled tokens). Stop and check the data format.")
+
+    history = []
+    best = {"acc": -1.0}
+
+    def record(tag, res, step):
+        history.append({"tag": tag, "step": step, **res})
+        json.dump(history, open(os.path.join(args.out, "history.json"), "w"), ensure_ascii=False, indent=1)
+
+    if not args.skip_baseline:
+        print("[baseline] zero-shot accuracy before training")
+        record("baseline", evaluate(model, tok, val_rows, args.max_seq_length, args.eval_limit, args.eval_batch), 0)
+
+    class EpochEval(TrainerCallback):
+        last = -1
+
+        def _run(self, tag, state, model):
+            if state.global_step == self.last:
+                return
+            self.last = state.global_step
+            res = evaluate(model, tok, val_rows, args.max_seq_length, args.eval_limit, args.eval_batch)
+            record(tag, res, state.global_step)
+            if res["acc"] > best["acc"]:
+                best.update(res, step=state.global_step)
+                d = os.path.join(args.out, "best_adapter")
+                model.save_pretrained(d)
+                processor.save_pretrained(d)
+                print(f"[save] new best accuracy {res['acc']:.4f} -> {d}")
+
+        def on_step_end(self, a, state, control, model=None, **kw):
+            if args.eval_steps > 0 and state.global_step > 0 and state.global_step % args.eval_steps == 0:
+                self._run(f"step {state.global_step}", state, model)
+
+        def on_epoch_end(self, a, state, control, model=None, **kw):
+            self._run(f"epoch {state.epoch:.2f}", state, model)
+
+        def on_train_end(self, a, state, control, model=None, **kw):
+            self._run("end", state, model)
+
+    cfg_kw = dict(
+        output_dir=os.path.join(args.out, "trainer"), per_device_train_batch_size=args.batch_size,
+        gradient_accumulation_steps=args.grad_accum, learning_rate=args.lr, lr_scheduler_type="cosine",
+        warmup_steps=args.warmup_steps, weight_decay=0.01, max_grad_norm=1.0,
+        num_train_epochs=args.epochs, max_steps=args.max_steps, bf16=True, optim="adamw_8bit",
+        logging_steps=10, save_strategy="no", seed=args.seed, report_to="none",
+        remove_unused_columns=False, dataset_text_field="", max_length=args.max_seq_length)
+    if "dataset_kwargs" in inspect.signature(SFTConfig.__init__).parameters:
+        cfg_kw["dataset_kwargs"] = {"skip_prepare_dataset": True}
+    trainer = SFTTrainer(model=model, processing_class=tok, train_dataset=train_ds,
+                         data_collator=collator, args=SFTConfig(**cfg_kw), callbacks=[EpochEval()])
+    stats = trainer.train()
+    print(f"[done] {stats.metrics.get('train_runtime', 0) / 60:.1f} min; "
+          f"peak VRAM {torch.cuda.max_memory_reserved() / 1024 ** 3:.1f} GB")
+
+    final = os.path.join(args.out, "final_adapter")
+    model.save_pretrained(final)
+    processor.save_pretrained(final)
+    print(f"[save] final adapter -> {final}")
+    print("\nValidation accuracy history:")
+    for h in history:
+        print(f"  {h['tag']:>12s}  step {h['step']:5d}  acc {h['acc']:.4f}  (n={h['n']}, "
+              f"majority-position baseline {h['majority_baseline']:.3f})")
+    print(f"best: {best['acc']:.4f} at step {best.get('step')}  -> {os.path.join(args.out, 'best_adapter')}")
+
+
+if __name__ == "__main__":
+    main()