Blame for inference.py

"""
Reconstructed runner for the stages contributing to Codabench submission #138.
See README.md and REPRODUCIBILITY.md before using this runner for final ranking.
The saved-pass merge is byte-exact; fresh model reproduction has NOT been verified.
Original Kaggle source snapshots are under training/kaggle_cells/.

Default: regenerate passes with Ollama, without reading archived test predictions.
Pass 3 retains historical temperature 0.1 and no explicit seed. This conflicts with
no-sampling guidance in the supplied template; it is disclosed, not silently changed.
The historical Mac run stopped after 1160 items of a full priority queue; a fresh
run completes that queue, since a hard-coded test-count cutoff is not generalizable.

Run from a separate output directory to preserve the submitted predictions.jsonl.
--merge-only is an archival audit, not fresh model inference.
"""

# ======================= ORGANIZERS CHANGE ONLY THIS LINE =======================
DATA_REPO = "SinhalaMMLU-SharedTask/Sinhala-MMLU-SharedTask-Test-Dataset"
# ================================================================================

OUTPUT_FILE = "predictions.jsonl"

# ----------------------- TEAM SETTINGS (exactly as used for the final submission)
MODEL_TAG = "gemma4:e4b"
SEED = None  # Original archived requests did not specify an Ollama seed.
STAGES_DEFAULT = "baselines,pass1,pass2,pass3,pass45,mac"

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import time
import urllib.request
from collections import Counter
import hashlib
import platform
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

HERE = Path(__file__).resolve().parent
LEVEL_RANK = {"easy": 0, "medium": 1, "hard": 2}


# =============== 1. LOAD THE DATASET (template loader; `--local-data` is for offline checks only) ===============
def load_questions(data_repo, local_data=None):
    rows = []
    if local_data:
        for lvl in ("easy", "medium", "hard"):
            for f in sorted((Path(local_data) / lvl).rglob("*.json")):
                raw = json.load(open(f, encoding="utf-8"))
                if isinstance(raw, dict):
                    raw = next(v for v in raw.values() if isinstance(v, list))
                for r in raw:
                    r["_file"] = f"{lvl}/{f.name}"
                    rows.append(r)
        return rows
    from huggingface_hub import HfFileSystem, list_repo_files
    fs = HfFileSystem()
    files = sorted(f for f in list_repo_files(data_repo, repo_type="dataset") if f.endswith(".json"))
    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 str((r.get("metadata") or {}).get("difficulty") or r["_file"].split("/")[0]).lower().strip()


def prepare(rows):
    """Normalise rows (list 'choices', lower-case '_level') and order them easy -> medium -> hard."""
    for r in rows:
        ch = r.get("choices")
        if ch is None:
            ch = r.get("options", [])
        if isinstance(ch, dict):
            ch = [ch[k] for k in sorted(ch.keys(), key=lambda x: str(x))]
        r["choices"] = list(ch)
        r["_level"] = get_level(r)
    # Preserve the organiser loader's row order. Scheduling is handled separately.
    seen = set()
    for r in rows:
        key = kof(r)
        if key in seen:
            raise ValueError(f"Duplicate question key: {key}")
        seen.add(key)
        if r["_level"] not in LEVEL_RANK or len(r["choices"]) != 4:
            raise ValueError(f"Expected known difficulty and four choices: {key}")
    if not rows:
        raise ValueError("No questions were loaded")
    return rows


def kof(q):
    return (q["q_no"], q["subject"], q["category"], q["_level"])


# =============== 2. YOUR SYSTEM ===============
SYSTEM_THINK = (
    "You are an expert teacher of the Sri Lankan national school curriculum (grades 6-13), "
    "fluent in Sinhala. You will receive a multiple-choice question written in Sinhala with four "
    "numbered options. Exactly one option is correct. Read the Sinhala carefully, including "
    "negations such as 'නොවන' / 'නොවේ' (NOT). You may reason in English. "
    "Finish with the final answer in the form 'Answer: <number>'."
)
SYSTEM_THINK_HINT = (  # pass 3 only
    "You are an expert teacher of the Sri Lankan national school curriculum (grades 6-13), "
    "fluent in Sinhala. You will receive a multiple-choice question written in Sinhala with four "
    "numbered options. Exactly one option is correct. Read the Sinhala carefully, including "
    "negations such as 'නොවන' / 'නොවේ' (NOT). If the question or options contain Sanskrit, "
    "Hindustani, Pali, Tamil, or technical terms, carefully map them to their standard curriculum "
    "meaning. You may reason in English. Finish with the final answer in the form 'Answer: <number>'."
)
ANS_RE = re.compile(r"(?:Answer|පිළිතුර)\s*[:：]?\s*\**\s*\(?([1-4])\)?", re.I)
POS_RE = re.compile(r"ඉහත|above|\b[1-4]\s*(?:හා|සහ|හෝ|,|and|or)\s*[1-4]\b", re.I)
STOP_NONTHINK = ["\nප්‍රශ්නය:", "\nQuestion:", "<|im_end|>", "<|endoftext|>", "<|im_start|>"]
_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({}))


def think_user_prompt(q, shift=0):
    lvl = q["_level"]
    grade = "6-9" if lvl == "easy" else ("10-11 (O/L)" if lvl == "medium" else "12-13 (A/L)")
    ch = list(q["choices"])
    ch = ch[shift:] + ch[:shift]
    opts = "\n".join(f"{i}. {c}" for i, c in enumerate(ch, 1))
    return (
        f"විෂය (Subject): {q.get('subject', 'General')} | ශ්‍රේණිය (Grade): {grade}\n\n"
        f"ප්‍රශ්නය (Question):\n{str(q['question']).strip()}\n\n"
        f"විකල්ප (Options):\n{opts}\n\n"
        "Which option (1, 2, 3 or 4) is correct? End your reply with 'Answer: <number>'."
    )


# ---------------------------------------------------------------- Ollama plumbing
def _alive(host):
    try:
        urllib.request.urlopen(f"{host}/api/tags", timeout=3)
        return True
    except Exception:
        return False


def ensure_ollama(hosts):
    """Check/start existing Ollama; model provisioning is an explicit setup step."""
    if shutil.which("ollama") is None:
        raise RuntimeError("Install Ollama first (see README.md). No automatic shell installer is executed.")
    for idx, host in enumerate(hosts):
        if not _alive(host):
            if not host.startswith(("http://127.0.0.1:", "http://localhost:")):
                raise RuntimeError(f"Only local Ollama hosts are supported: {host}")
            env = {**os.environ, "OLLAMA_HOST": host.replace("http://", ""), "OLLAMA_NUM_PARALLEL": "2",
                   "OLLAMA_FLASH_ATTENTION": "1", "OLLAMA_KEEP_ALIVE": "24h"}
            if len(hosts) > 1:
                env["CUDA_VISIBLE_DEVICES"] = str(idx)
            subprocess.Popen(["ollama", "serve"], env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            for _ in range(60):
                if _alive(host):
                    break
                time.sleep(1)
        with _OPENER.open(f"{host}/api/tags", timeout=10) as response:
            tags = json.load(response).get("models", [])
        model = next((m for m in tags if MODEL_TAG in (m.get("name"), m.get("model"))), None)
        if model is None:
            raise RuntimeError(f"{MODEL_TAG} is not installed on {host}. Provision the exact original model first.")
        print("Ollama model:", model.get("name"), "digest:", model.get("digest"), flush=True)



def default_hosts():
    env = os.environ.get("OLLAMA_HOSTS")
    if env:
        return [h.strip() if h.strip().startswith("http") else f"http://{h.strip()}" for h in env.split(",")]
    try:
        n = len(subprocess.run(["nvidia-smi", "-L"], capture_output=True, text=True).stdout.strip().splitlines())
    except Exception:
        n = 0
    return ["http://127.0.0.1:11434", "http://127.0.0.1:11435"] if n >= 2 else ["http://127.0.0.1:11434"]


def chat(host, system, user, think, options, timeout=420):
    payload = json.dumps({
        "model": MODEL_TAG,
        "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}],
        "stream": False, "think": think, "options": options,
    }).encode("utf-8")
    req = urllib.request.Request(f"{host}/api/chat", data=payload, headers={"Content-Type": "application/json"})
    with _OPENER.open(req, timeout=timeout) as resp:
        msg = json.loads(resp.read().decode("utf-8")).get("message", {})
    return msg.get("content", "") or "", msg.get("thinking", "") or ""


def extract_think(content, thinking, digit_fallback=True):
    """Answer digit from a thinking-mode reply, or None."""
    m = ANS_RE.findall(content) or ANS_RE.findall(thinking)
    if m:
        return m[-1]
    if digit_fallback:
        d = re.findall(r"\b([1-4])\b", content if content.strip() else thinking)
        if d:
            return d[-1]
    return None


def extract_answer_digit(text):
    """Answer extraction for the non-thinking 4-shot baselines (same as training/src/baseline_runner.py)."""
    if not text:
        return "1"
    if "</think>" in text:
        text = text.split("</think>")[-1]
    text = text.strip()
    ex = re.findall(r"(?:පිළිතුර|Answer|Final Answer|Correct Option|නිවැරදි පිළිතුර|තෝරාගැනීම)\s*(?:is|වන්නේ)?\s*[:\s*#(\[]*(?:Option\s*)?([1-4])\b", text, re.I)
    if ex:
        return ex[-1]
    m = re.match(r"^[(\[]?\s*([1-4])\b", text)
    if m and len(text) <= 8:
        return m.group(1)
    d = re.findall(r"\b([1-4])\b", text)
    if d:
        return d[-1]
    m = re.search(r"[1-4]", text)
    return m.group(0) if m else "1"


# ---------------------------------------------------------------- checkpointed parallel runner
def load_ck(path, required=False):
    res = {}
    path = Path(path)
    if not path.exists():
        if required:
            raise FileNotFoundError(f"Required saved pass missing: {path}")
        return res
    with path.open(encoding="utf-8") as f:
        for number, line in enumerate(f, 1):
            if line.strip():
                r = json.loads(line)
                key = (r["q_no"], r["subject"], r["category"], r["level"])
                answer = str(r["answer"])
                if key in res or answer not in {"1", "2", "3", "4"}:
                    raise ValueError(f"Duplicate key or invalid answer at {path}:{number}")
                res[key] = answer
    return res


def run_stage(name, qs, solver, ck_path, hosts, workers):
    """Run solver(host, q) -> answer-or-None for every question in qs not yet in the checkpoint."""
    done = load_ck(ck_path)
    todo = [q for q in qs if kof(q) not in done]
    print(f"[{name}] targets={len(qs)} done={len(done)} todo={len(todo)}", flush=True)
    t0 = time.time()
    n = 0
    with open(ck_path, "a", encoding="utf-8") as fo, ThreadPoolExecutor(max_workers=workers) as ex:
        futs = {ex.submit(solver, hosts[i % len(hosts)], q): q for i, q in enumerate(todo)}
        for fut in as_completed(futs):
            q = futs[fut]
            ans = fut.result()
            if ans is not None:
                fo.write(json.dumps({"q_no": q["q_no"], "subject": q["subject"], "category": q["category"],
                                     "level": q["_level"], "answer": ans}, ensure_ascii=False) + "\n")
                fo.flush()
                done[kof(q)] = ans
            n += 1
            if n % 20 == 0 or n == len(todo):
                print(f"  [{name}] {n}/{len(todo)}  ({(time.time() - t0) / 60:.1f} min)", flush=True)
    return done


def tries(fn, n=2):
    for _ in range(n):
        try:
            return fn()
        except Exception as exc:
            print(f"Model request failed ({type(exc).__name__}): {exc}", file=sys.stderr, flush=True)
            time.sleep(2)
    return None


# ---------------------------------------------------------------- stages
def stage_baseline(qs, hosts, workers, wd, cot):
    sys.path.insert(0, str(training_src()))
    from few_shot_prompts import build_few_shot_prompt
    system = ("You are an expert examiner for the Sri Lankan school curriculum (Sinhala medium). "
              "For each multiple-choice question, write a concise 1-2 sentence English translation and fact-check ('Analysis:'), "
              "then end on a new line with strictly 'Answer: <digit>' (where <digit> is 1, 2, 3, or 4)."
              if cot else
              "You are an examination AI for the Sri Lankan school curriculum (Sinhala medium). "
              "For every multiple-choice question, analyze the options carefully and answer by replying "
              "with ONLY the single digit of the best option (1, 2, 3, or 4). "
              "Do NOT provide explanations, do NOT include option text, and do NOT write anything other than the single digit.")
    opts = {"temperature": 0.0, "num_predict": 260 if cot else 15, "stop": STOP_NONTHINK}

    def solve(host, q):
        prompt = build_few_shot_prompt(q, num_shots=4, cot=cot, cot_lang="en")
        r = tries(lambda: chat(host, system, prompt, False, opts, 240), n=1)
        return extract_answer_digit(r[0]) if r else "1"

    return run_stage("cot" if cot else "direct", qs, solve, wd / ("baseline_cot_4shot.jsonl" if cot else "baseline_direct_4shot.jsonl"), hosts, workers)


def stage_pass1(qs, hosts, workers, wd, cot, name="pass1", ck="pass1_think_2048_orig.jsonl"):
    opts = {"temperature": 0.0, "num_predict": 2048}

    def solve(host, q):
        r = tries(lambda: chat(host, SYSTEM_THINK, think_user_prompt(q), True, opts, 240), n=1)
        a = extract_think(*r) if r else None
        return a if a else cot.get(kof(q), "1")

    return run_stage(name, qs, solve, wd / ck, hosts, workers)


def stage_rot(qs, hosts, workers, wd, shift, name, ck, p1, cot, fallback=True):
    """Rotated-options pass (4096 tokens, greedy). Predictions are mapped back to the original option numbers."""
    opts = {"temperature": 0.0, "num_predict": 4096}

    def solve(host, q):
        r = tries(lambda: chat(host, SYSTEM_THINK, think_user_prompt(q, shift), True, opts, 300), n=1 if fallback else 2)
        if r is None:
            return (p1.get(kof(q)) or cot.get(kof(q), "1")) if fallback else None
        a = extract_think(*r, digit_fallback=fallback)
        if a is None:
            if not fallback:
                return None
            a = "1"  # Original Pass 2 defaults to rotated option 1, then maps back.
        return str(((int(a) - 1 + shift) % 4) + 1)

    return run_stage(name, qs, solve, wd / ck, hosts, workers)


def stage_pass3(qs, hosts, workers, wd, p1):
    opts = {"temperature": 0.1, "top_p": 0.95, "num_predict": 4096}

    def solve(host, q):
        r = tries(lambda: chat(host, SYSTEM_THINK_HINT, think_user_prompt(q), True, opts, 300), n=3)
        if r is None:
            return p1.get(kof(q), "1")
        return extract_think(*r) or "1"

    return run_stage("pass3", qs, solve, wd / "pass3_think_4096_hint.jsonl", hosts, workers)


# =============== merge: fixed level-aware weighted vote (identical to training/ensemble_v9.py) ===============
def merge_one(k, pos, p1, p2, p3, p4, p5, mac, cot, dr):
    lvl = k[3]
    v = Counter()
    if lvl == "hard":
        v[p1[k]] += 1.15
        if k in mac: v[mac[k]] += 1.10
        if k in p3:  v[p3[k]]  += 1.10
        for m in (p2, p4, p5):
            if k in m and k not in pos: v[m[k]] += 1.10
        if k in cot: v[cot[k]] += 0.35
        if k in dr:  v[dr[k]]  += 0.35
    elif lvl == "medium":
        v[p1[k]] += 1.25
        if k in mac: v[mac[k]] += 1.05
        if k in p3:  v[p3[k]]  += 1.10
        for m in (p2, p4, p5):
            if k in m and k not in pos: v[m[k]] += 1.05
        if k in cot: v[cot[k]] += 0.35
        if k in dr:  v[dr[k]]  += 0.35
    else:  # easy: pass 1 anchored at 2.40 so only a margin >= 2 supermajority can overturn it
        v[p1[k]] += 2.40
        if k in mac: v[mac[k]] += 1.0
        if k in p3:  v[p3[k]]  += 1.0
        for m in (p2, p4, p5):
            if k in m and k not in pos: v[m[k]] += 0.95
        if k in cot: v[cot[k]] += 0.30
        if k in dr:  v[dr[k]]  += 0.30
    return v.most_common(1)[0][0]


def training_src():
    p = HERE / "training" / "src"
    if not (p / "few_shot_prompts.py").exists():
        raise FileNotFoundError("Download the complete team repository, including training/src.")
    return p


def priority_order(rows, cot, dr):
    # Original scheduler: disagreements first, then hard -> medium -> easy.
    canonical = sorted(rows, key=lambda q: (LEVEL_RANK[q["_level"]], q["_file"]))
    return sorted(canonical, key=lambda q: (
        0 if cot.get(kof(q)) != dr.get(kof(q)) else 1,
        {"hard": 0, "medium": 1, "easy": 2}[q["_level"]]))


def stage_mac(qs, hosts, wd, cot):
    # Original local runner was sequential and used its own broader answer parser.
    sys.path.insert(0, str(training_src()))
    from baseline_runner import extract_answer_digit as original_extract
    def solve(host, q):
        r = tries(lambda: chat(host, SYSTEM_THINK, think_user_prompt(q), True,
                              {"temperature": 0.0, "num_predict": 2048}, 240), n=1)
        if r is None:
            return cot.get(kof(q), "1")
        content, thinking = r
        found = ANS_RE.findall(content) or ANS_RE.findall(thinking)
        return found[-1] if found else original_extract(content if content.strip() else thinking)
    return run_stage("mac", qs, solve, wd / "mac_think_2048_orig.jsonl", hosts[:1], 1)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--merge-only", metavar="PASSES_DIR", default=None,
                    help="skip the model: merge the saved raw passes (training/passes) -> reproduces predictions.jsonl")
    ap.add_argument("--local-data", metavar="DIR", default=None, help="read easy/medium/hard/*.json from a local folder instead of Hugging Face")
    ap.add_argument("--output", default=OUTPUT_FILE, help="output path; submitted predictions are protected")
    args = ap.parse_args()
    output = Path(args.output).resolve()
    if output == (HERE / "predictions.jsonl").resolve():
        raise ValueError("Do not overwrite submitted predictions.jsonl. Run from a separate folder or use --output.")

    rows = prepare(load_questions(DATA_REPO, args.local_data))
    limit = int(os.environ.get("LIMIT", "0"))
    if limit:
        keep, cnt = [], Counter()
        for r in rows:
            if cnt[r["_level"]] < limit:
                keep.append(r)
                cnt[r["_level"]] += 1
        rows = keep
    print(len(rows), "questions loaded from", args.local_data or DATA_REPO)
    pos = {kof(q) for q in rows if POS_RE.search(" ".join(map(str, q["choices"])))}

    if args.merge_only:
        d = Path(args.merge_only)
        names = {"p1": "pass1_think_2048_orig.jsonl", "p2": "pass2_think_4096_rot1.jsonl", "p3": "pass3_think_4096_hint.jsonl",
                 "p4": "pass4_think_4096_rot2.jsonl", "p5": "pass5_think_4096_rot3.jsonl", "mac": "mac_think_2048_orig.jsonl",
                 "cot": "baseline_cot_4shot.jsonl", "dr": "baseline_direct_4shot.jsonl"}
        M = {k: load_ck(d / v, required=True) for k, v in names.items()}
        expected = {kof(q) for q in rows}
        if set(M["p1"]) != expected:
            raise ValueError("Saved Pass 1 must match the complete input key set (no LIMIT in merge-only mode).")
        if any(set(v) - expected for v in M.values()):
            raise ValueError("Saved passes contain question keys absent from the supplied dataset.")
    else:
        wd = Path(os.environ.get("WORKDIR", "inference_work"))
        wd.mkdir(parents=True, exist_ok=True)
        stages = set(STAGES_DEFAULT.split(","))
        hosts = default_hosts()
        workers = int(os.environ.get("WORKERS", 2 * len(hosts)))
        ensure_ollama(hosts)
        print("hosts:", hosts, "| workers:", workers, "| model:", MODEL_TAG)

        config = {"dataset_sha256": hashlib.sha256(json.dumps(rows, ensure_ascii=False, sort_keys=True).encode()).hexdigest(),
                  "runner_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                  "model": MODEL_TAG, "hosts": hosts, "workers": workers,
                  "python": platform.python_version(), "seed": None}
        meta = wd / "run_config.json"
        if meta.exists() and json.loads(meta.read_text()) != config:
            raise ValueError("Workdir belongs to a different dataset/configuration. Use a new WORKDIR.")
        meta.write_text(json.dumps(config, indent=2) + "\n")
        print("Fresh-run reconstruction: no guarantee of the historical score; see REPRODUCIBILITY.md.")
        baseline_rows = sorted(rows, key=lambda q: (LEVEL_RANK[q["_level"]], q["_file"]))
        cot = stage_baseline(baseline_rows, hosts[:1], 1, wd, True)
        dr = stage_baseline(baseline_rows, hosts[:1], 1, wd, False)
        priority = priority_order(rows, cot, dr)
        p1 = stage_pass1(priority, hosts, workers, wd, cot)
        rot_ok = [q for q in priority if len(q["choices"]) == 4]

        # pass 2: hard + questions where the two 4-shot baselines disagree
        t2 = [q for q in rot_ok if q["_level"] == "hard" or cot.get(kof(q)) != dr.get(kof(q))]
        p2 = stage_rot(t2, hosts, workers, wd, 1, "pass2", "pass2_think_4096_rot1.jsonl", p1, cot) if "pass2" in stages else {}

        # pass 3 targets: Pass1 != Pass2, or (not in Pass2) Pass1 != CoT / Direct
        t3 = []
        for q in baseline_rows:
            k = kof(q)
            if k in p2:
                if p1[k] != p2[k]:
                    t3.append(q)
            elif p1[k] != cot.get(k) or p1[k] != dr.get(k):
                t3.append(q)
        p3 = stage_pass3(t3, hosts, workers, wd, p1) if "pass3" in stages else {}

        # pass 4/5: same uncertain set, only 4-option questions without positional options; unclear replies are skipped
        t45 = [q for q in t3 if len(q["choices"]) == 4 and kof(q) not in pos]
        p4 = stage_rot(t45, hosts, workers, wd, 2, "pass4", "pass4_think_4096_rot2.jsonl", p1, cot, fallback=False) if "pass45" in stages else {}
        p5 = stage_rot(t45, hosts, workers, wd, 3, "pass5", "pass5_think_4096_rot3.jsonl", p1, cot, fallback=False) if "pass45" in stages else {}

        # The original Mac runner traversed ALL questions in priority order and was stopped.
        # Do not encode a historical test-set count as a new-data selection rule.
        mac = stage_mac(priority, hosts, wd, cot) if "mac" in stages else {}
        M = {"p1": p1, "p2": p2, "p3": p3, "p4": p4, "p5": p5, "mac": mac, "cot": cot, "dr": dr}

    # =============== 3. WRITE PREDICTIONS ===============
    answers = []
    for r in rows:
        k = kof(r)
        if k not in M["p1"]:
            M["p1"][k] = M["cot"].get(k, "1")
        answers.append(merge_one(k, pos, M["p1"], M["p2"], M["p3"], M["p4"], M["p5"], M["mac"], M["cot"], M["dr"]))
    assert len(answers) == len(rows), "one answer per question required"
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("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": r["_level"], "answer": a}, ensure_ascii=False) + "\n")
    print("\nSaved", output)


if __name__ == "__main__":
    main()