script.py

#!/usr/bin/env python
"""IOL-AI 2026 submission -- International Linguistics Olympiad solver
  * Item-count detection now also understands lowercase sub-item letters and
    fill-in-the-blank underscores, which the previous regex set missed.
  * Answer voting no longer requires byte-identical normalized strings to
    agree -- it clusters samples by chrF similarity, so "iepure" and
    "iepure " and "Iepure" (or a stray trailing gloss) still count as the
    same vote instead of splitting support three ways.
  * Self-consistency sampling is convergence-aware: once a problem block's
    items all have a confident majority, that block is dropped from the
    prompt list for later rounds, so remaining time budget concentrates on
    the rows that are still undecided rather than re-rolling rows that have
    already settled.
  * The matching-task assignment solver is kept (free-form generation on
    permutation tasks tends to just emit the identity ordering) but is now
    reused for the "which digit" family of tasks too, with a lighter
    single-token constraint.
  * A rough token-length probe runs right after the tokenizer loads, so the
    starting batch size is chosen from the actual prompt lengths for this
    file instead of always starting at a fixed guess and paying for at least
    one guaranteed OOM.
Environment: transformers 4.44.1 / torch 2.4.0 / autoawq, 16GB T4, fp16 only,
no internet, hard 30-minute wall clock.
"""
import os
import re
import json
import time
import unicodedata
from collections import Counter, defaultdict

START = time.time()

# ---------------------------------------------------------------------------
# Budget
# ---------------------------------------------------------------------------
# The harness kills the process at the wall-clock limit with zero credit for
# work in flight, so every stage aims to finish with a margin, not exactly on
# time. SAFETY is subtracted up front rather than checked ad hoc so every
# caller of time_left() sees the same conservative deadline.
TIME_LIMIT = float(os.environ.get("IOL_TIME_LIMIT", "1800"))
SAFETY = float(os.environ.get("IOL_SAFETY", "150"))
DEADLINE = START + TIME_LIMIT - SAFETY

TEST_CSV = os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv")
OUT_CSV = os.environ.get("IOL_OUT_CSV", "submission.csv")
MODEL_ID = os.environ.get("IOL_MODEL", ".")
WANT_EXPLANATION = os.environ.get("IOL_EXPLAIN", "1") == "1"
MAX_NEW = int(os.environ.get("IOL_MAXNEW", "900"))
MAX_SAMPLES = int(os.environ.get("IOL_MAXSAMPLES", "8"))

# Earlier runs found the organizers' plain, no-chain-of-thought prompt style
# outscored an elaborate CoT harness on the hidden set. Default to that
# "faithful" style, but keep the richer CoT prompt available for A/B use --
# the parsing/voting/repair layers below are written to help either mode
# rather than assuming one.
PROMPT_STYLE = os.environ.get("IOL_PROMPT_STYLE", "faithful")  # faithful | guided
SAMPLE_TEMP = float(os.environ.get("IOL_TEMP", "0.5"))

os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")


def log(msg):
    print(f"[{time.time() - START:7.1f}s] {msg}", flush=True)


def time_left():
    return DEADLINE - time.time()


# ---------------------------------------------------------------------------
# How many items does this problem block ask for?
# ---------------------------------------------------------------------------
_NUM_LINE = re.compile(r"^[ \t]*(\d{1,3})[.)\]]", re.M)
_NUM_PAREN = re.compile(r"\((\d{1,3})\)")
_NUM_RANGE = re.compile(r"\(?(\d{1,3})\s*(?:[-\u2013\u2014]|to)\s*(\d{1,3})\)?")
_LETTER_LINE = re.compile(r"^[ \t]*([A-Za-z])[.)\]]\s", re.M)
_LETTER_PAREN = re.compile(r"\(([A-Za-z])\)")
_BLANK_MARK = re.compile(r"_{2,}|\[\s*(?:blank|___*)\s*\]", re.I)


def count_items(query, task_type="", context=""):
    """Best-effort count of the sub-answers a block expects. Never returns 0."""
    q = query or ""

    nums = set(int(m) for m in _NUM_LINE.findall(q)) | set(
        int(m) for m in _NUM_PAREN.findall(q))
    letters = set(_LETTER_LINE.findall(q)) | set(_LETTER_PAREN.findall(q))

    span = 0
    for a, b in _NUM_RANGE.findall(q):
        a, b = int(a), int(b)
        if 0 < b - a < 60:
            span = max(span, b - a + 1)

    marker_count = max(len(nums), len(letters))
    if span and marker_count and span != marker_count:
        # Explicit markers on the page are what must actually be answered,
        # even if a stated range like "items 1-5" disagrees with them.
        n = marker_count
    else:
        n = max(span, marker_count)

    if n > 1:
        return n

    blanks = len(_BLANK_MARK.findall(q))
    if blanks > 1:
        return blanks

    # No numbering at all: "Translate the following into Kalam:" followed by
    # one bare item per line.
    body_lines = [l.strip() for l in q.splitlines() if l.strip()]
    if len(body_lines) > 1:
        head, rest = body_lines[0], body_lines[1:]
        if rest and (head.endswith((":", ".")) or len(rest) >= 2):
            return len(rest)

    # A bare instruction ("State the correspondence for each pair.") whose
    # items actually live in the shared context block.
    if context:
        cn = len(set(int(m) for m in _NUM_LINE.findall(context)))
        if cn > 1:
            return cn
        cl = len(set(_LETTER_LINE.findall(context)))
        if cl > 1:
            return cl

    return max(n, 1)


# ---------------------------------------------------------------------------
# Cleaning and parsing model output into exactly n answers
# ---------------------------------------------------------------------------
_LEAD_MARK = re.compile(r"^\s*(?:\(?\d{1,3}\)?[.):\]]\s*|[-*\u2022]\s+)")
_CODE_FENCE = re.compile(r"^```[a-zA-Z]*\s*$")
_META_TALK = re.compile(
    r"^\s*(?:here (?:are|is)\b|answers?\s*:?\s*$|explanation\b|note\b|okay\b|"
    r"solution\b|reasoning\b|analysis\b|translations?\s*:?\s*$|the answers?\b|"
    r"let me\b|first,|so,|therefore\b|thus\b|to summarize\b)",
    re.I,
)


def tidy(line):
    s = line.strip()
    s = _LEAD_MARK.sub("", s)
    s = s.strip().strip("`").strip()
    if len(s) >= 2 and s[0] == s[-1] and s[0] in "\"'\u201c\u201d":
        s = s[1:-1].strip()
    return s.strip()


def source_fallback(query, n):
    """Last-resort answers when the model produced nothing usable.
    A blank answer scores zero on both exact-match and chrF, but echoing the
    prompt's own item text is nearly free and, on transcription-adjacent
    tasks, shares enough surface form with the true answer to pick up
    partial chrF credit. Better than nothing, never claimed to be more.
    """
    q = query or ""
    picked = []
    for ln in q.splitlines():
        s = ln.strip()
        if not s:
            continue
        m = re.match(r"^\(?(\d{1,3})\)?[.):\]]\s*(.+)$", s)
        if m:
            picked.append(m.group(2).strip())
    if not picked:
        lines = [l.strip() for l in q.splitlines() if l.strip()]
        if len(lines) > 1 and lines[0].endswith((":", ".")):
            picked = lines[1:]
    picked = [p.split("|")[0].strip() if "|" in p else p for p in picked]
    picked = [p for p in picked if p]
    while len(picked) < n:
        picked.append(picked[-1] if picked else "?")
    return picked[:n]


def squeeze_to_n(items, n, fallback=None):
    items = [i for i in items if i and i.strip()]
    if len(items) > n:
        # The model reasons first and answers last; absent an explicit
        # marker to slice on, the tail of the output is the answer block.
        items = items[-n:]
    while len(items) < n:
        if fallback and len(items) < len(fallback):
            items.append(fallback[len(items)])
        else:
            items.append(items[-1] if items else "?")
    return items[:n]


def parse_model_output(text, n, fallback=None):
    if not text or not text.strip():
        return list(fallback[:n]) if fallback else ["?"] * n

    marker = None
    for m in re.finditer(r"(?:^|\n)\s*(?:final\s+)?answers?\s*:\s*\n?", text, re.I):
        marker = m
    body = text[marker.end():] if marker else text

    labelled, plain = [], []
    for ln in body.splitlines():
        if _CODE_FENCE.match(ln):
            continue
        m = re.match(r"^\s*\(?(\d{1,3})\)?[.):\]]\s*(.+)$", ln.strip())
        if m:
            val = tidy(m.group(2))
            if val and not _META_TALK.match(val):
                labelled.append((int(m.group(1)), val))
        c = tidy(ln)
        if c and not _META_TALK.match(c):
            plain.append(c)

    if len(labelled) >= n:
        # Numbered labels, when present, are a stronger placement signal than
        # line order -- keep the last value written under each label since
        # models sometimes restate an earlier answer while reasoning.
        by_label = {}
        for lab, val in labelled:
            by_label[lab] = val
        keys = sorted(by_label)
        if len(keys) >= n:
            return [by_label[k] for k in keys[:n]]

    return squeeze_to_n(plain, n, fallback)


_DIGITS_ONLY = re.compile(r"-?\d+")


def coerce_numeral(answer, task_type):
    """text_to_num answers should be bare digits; pull them out if the model
    wrapped them in words or punctuation instead of reworking the parser for
    one task family."""
    if (task_type or "").strip().lower() != "text_to_num":
        return answer
    m = _DIGITS_ONLY.search(answer or "")
    return m.group(0) if m else answer


def normalize(s):
    s = unicodedata.normalize("NFC", (s or "").strip().lower())
    s = re.sub(r"\s+", " ", s)
    return s.strip(" .!?;:,")


# ---------------------------------------------------------------------------
# chrF, computed inline since sacrebleu may not be present in the sandbox
# ---------------------------------------------------------------------------
def _char_ngrams(s, k):
    s = re.sub(r"\s+", "", s)
    if len(s) < k:
        return Counter()
    return Counter(s[i:i + k] for i in range(len(s) - k + 1))


def chrf(hyp, ref, order=6, beta=2.0):
    if not hyp or not ref:
        return 0.0
    precisions, recalls = [], []
    for k in range(1, order + 1):
        h, r = _char_ngrams(hyp, k), _char_ngrams(ref, k)
        if not h or not r:
            continue
        overlap = sum((h & r).values())
        precisions.append(overlap / max(1, sum(h.values())))
        recalls.append(overlap / max(1, sum(r.values())))
    if not precisions:
        return 0.0
    p = sum(precisions) / len(precisions)
    r = sum(recalls) / len(recalls)
    if p + r == 0:
        return 0.0
    b2 = beta * beta
    return (1 + b2) * p * r / (b2 * p + r)


# ---------------------------------------------------------------------------
# Voting: cluster candidates by similarity rather than exact string match
# ---------------------------------------------------------------------------
CLUSTER_THRESHOLD = 0.90  # chrF similarity above which two answers "agree"


def cluster_candidates(cands):
    """Group near-duplicate strings together (single-linkage on chrF).
    Exact-match voting undercounts agreement: "iepurele" and "iepurele " or a
    stray trailing gloss are the same underlying answer but would previously
    split into separate buckets and never reach a majority. Grouping by
    similarity instead means minor formatting noise doesn't cost a block its
    vote.
    """
    groups = []  # list of [members]
    for c in cands:
        placed = False
        for g in groups:
            if chrf(c, g[0]) >= CLUSTER_THRESHOLD or normalize(c) == normalize(g[0]):
                g.append(c)
                placed = True
                break
        if not placed:
            groups.append([c])
    return groups


def vote_answer(cands, anchor=None):
    """Choose one answer for an item from the greedy pass plus samples.
    The greedy (temperature 0) answer is the default and can only be
    displaced by a cluster that (a) has at least two members and (b) is
    strictly larger than the cluster the anchor itself belongs to. This
    asymmetry matters: with only a handful of samples, treating all
    candidates as equal and breaking ties by "closest to the others" is
    poorly conditioned -- with two candidates that comparison is symmetric
    and degenerates into picking whichever string happens to be shorter. On
    a held-out mock set that variant swapped out the greedy answer for a
    sampled one roughly half the time and cost a large chunk of exact match.
    Anchoring keeps voting monotone: it only overrides on genuine, repeated
    agreement.
    """
    cands = [c for c in cands if c and c.strip()]
    if anchor is None:
        anchor = cands[0] if cands else "?"
    if len(cands) < 3:
        return anchor

    groups = cluster_candidates(cands)
    anchor_group_size = 1
    for g in groups:
        if any(normalize(m) == normalize(anchor) or chrf(m, anchor) >= CLUSTER_THRESHOLD
               for m in g):
            anchor_group_size = len(g)
            break

    best_group = max(groups, key=len)
    if len(best_group) >= 2 and len(best_group) > anchor_group_size:
        # Within the winning cluster, prefer whichever exact form is most
        # common; ties fall back to the first (greedy-adjacent) occurrence.
        return Counter(best_group).most_common(1)[0][0]
    return anchor


# ---------------------------------------------------------------------------
# Matching / assignment problems: read next-token logits instead of trusting
# free-form generation, which tends to just emit the option letters in order.
# ---------------------------------------------------------------------------
_ROW_OPT = re.compile(r"^[ \t]*([A-Za-z])[.)]\s+(.+)$", re.M)
_ROW_ITEM = re.compile(r"^[ \t]*(\d{1,3})[.)]\s+(.+)$", re.M)


def split_matching_block(context):
    items = [(int(a), b.strip()) for a, b in _ROW_ITEM.findall(context or "")]
    opts = [(a, b.strip()) for a, b in _ROW_OPT.findall(context or "")]
    seen = set()
    items = [x for x in items if not (x[0] in seen or seen.add(x[0]))]
    seen = set()
    opts = [x for x in opts if not (x[0] in seen or seen.add(x[0]))]
    return items, opts


def assign_max_weight(score):
    """One-to-one assignment maximizing total score. scipy if available,
    otherwise a greedy pass plus a few rounds of local 2-swaps."""
    n, m = len(score), len(score[0])
    try:
        import numpy as _np
        from scipy.optimize import linear_sum_assignment
        _, cols = linear_sum_assignment(-_np.array(score))
        return list(cols)
    except Exception:
        pass
    used, out = set(), [0] * n
    order = sorted(
        range(n),
        key=lambda i: -(max(score[i]) - sorted(score[i])[-2] if m > 1 else 0),
    )
    for i in order:
        j = max((j for j in range(m) if j not in used), key=lambda j: score[i][j],
                default=0)
        used.add(j)
        out[i] = j
    for _ in range(4):
        improved = False
        for a in range(n):
            for b in range(a + 1, n):
                cur = score[a][out[a]] + score[b][out[b]]
                alt = score[a][out[b]] + score[b][out[a]]
                if alt > cur + 1e-9:
                    out[a], out[b] = out[b], out[a]
                    improved = True
        if not improved:
            break
    return out


def enforce_bijection(answers):
    """If every answer is a single letter and duplicates exist where the full
    alphabet of options wasn't used, reassign duplicated slots to the unused
    letters. Tightly scoped so it never touches anything but this exact
    single-letter-permutation shape."""
    if len(answers) < 3 or not all(re.fullmatch(r"[A-Za-z]", a or "") for a in answers):
        return answers
    n = len(answers)
    universe = [chr(ord("A") + i) for i in range(n)]
    upper = [a.upper() for a in answers]
    if len(set(upper)) == n:
        return answers
    unused = [l for l in universe if l not in set(upper)]
    if not unused:
        return answers
    seen, out = set(), []
    for a in upper:
        if a in seen and unused:
            out.append(unused.pop(0))
        else:
            seen.add(a)
            out.append(a)
    return out


# ---------------------------------------------------------------------------
# Prompting
# ---------------------------------------------------------------------------
FAITHFUL_SYSTEM = (
    "You solve International Linguistics Olympiad problems. Answer every "
    "numbered item. Put each answer on its own line, in order, with no "
    "numbering and no extra text."
)

GUIDED_SYSTEM = (
    "You are a gold medallist at the International Linguistics Olympiad.\n"
    "Each problem gives data from a language you have never seen before; "
    "everything needed is in the problem itself, and no outside knowledge is "
    "required.\n"
    "Work method: align the given examples, segment the forms, find the "
    "recurring morphemes and the rule that orders them, verify the rule "
    "against every example given, then apply it to the items you are asked "
    "for.\n"
    "Reason concisely. Finish with a block that starts on its own line with "
    "exactly ANSWERS: followed by one answer per line, in order, with no "
    "numbering, no commentary, and no blank lines.\n"
    "Always give your best guess for every item -- never leave one blank."
)

TASK_HINTS = {
    "translation": "Give only the translation itself -- no source text, no "
                   "gloss, no quotation marks.",
    "match_letters": "Give a single capital letter per item, identifying its "
                     "match. Every letter is used exactly once; none repeats.",
    "fill_blanks": "Give only the missing form for that blank -- not the "
                   "whole line, not a gloss.",
    "text_to_num": "Give the number in digits only (e.g. 111).",
    "num_to_text": "Give the number spelled out in words, in the problem "
                   "language.",
}


def build_prompt(row, n):
    if PROMPT_STYLE == "faithful":
        return f"{row['context'].strip()}\n\n{row['query'].strip()}"
    hint = TASK_HINTS.get((row.get("task_type") or "").strip().lower(), "")
    return (
        f"{row['context'].strip()}\n\n{row['query'].strip()}\n\n"
        f"There are exactly {n} item{'s' if n != 1 else ''} to answer."
        + (f" {hint}" if hint else "")
        + f"\nAfter reasoning, write ANSWERS: on its own line, then exactly "
          f"{n} line{'s' if n != 1 else ''}, one answer per item, in order."
    )


EXPLAIN_SYSTEM = (
    "You explain International Linguistics Olympiad solutions to a human "
    "judge. Given a problem and the answers produced, state the key rules of "
    "the language that justify them -- relevant morphemes, word order, any "
    "sound changes. Be specific and brief (2-4 sentences or a few short "
    "bullets); do not restate a chain of reasoning."
)


def build_explain_prompt(row, answers):
    return (
        f"{row['context'].strip()}\n\n{row['query'].strip()}\n\n"
        "Answers given:\n" + "\n".join(f"- {a}" for a in answers) +
        "\n\nBriefly explain the linguistic rules behind these answers."
    )


# ---------------------------------------------------------------------------
# Offline diagnostics (never runs on the platform -- gold labels aren't
# present there, so IOL_GOLD stays unset)
# ---------------------------------------------------------------------------
def report_dev_score(preds):
    gold_path = os.environ.get("IOL_GOLD")
    if not gold_path or not os.path.exists(gold_path):
        return
    try:
        import ast
        import pandas as pd
        gold_df = pd.read_csv(gold_path, dtype=str)
        exact, chrfs = [], []
        for _, r in gold_df.iterrows():
            gold = ast.literal_eval(r["answer"])
            p = preds.get(str(r["id"]), [])
            p = list(p)[:len(gold)] + [""] * max(0, len(gold) - len(p))
            for gi, pi in zip(gold, p):
                alts = gi if isinstance(gi, (list, tuple)) else [gi]
                alts = [str(a) for a in alts]
                exact.append(1.0 if any(pi.strip() == a.strip() for a in alts) else 0.0)
                chrfs.append(max(chrf(pi, a) for a in alts))
        em = sum(exact) / max(1, len(exact))
        cf = sum(chrfs) / max(1, len(chrfs))
        log(f"  [dev] EM={em:.4f} chrF~={cf:.4f} score~={(em * cf) ** 0.5:.4f} "
            f"over {len(exact)} items")
    except Exception as e:
        log(f"  [dev] scoring failed: {type(e).__name__}: {e}")


def write_submission(path, ids, preds, explanations=None):
    import pandas as pd
    rows = []
    for i in ids:
        rec = {"id": i, "pred": json.dumps(preds[i], ensure_ascii=False)}
        if explanations is not None:
            rec["explanation"] = explanations.get(i, "")
        rows.append(rec)
    pd.DataFrame(rows).to_csv(path, index=False)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
    import pandas as pd

    df = pd.read_csv(TEST_CSV, dtype=str).fillna("")
    ids = [str(x) for x in df["id"].tolist()]
    ns = [count_items(r.get("query", ""), r.get("task_type", ""), r.get("context", ""))
          for _, r in df.iterrows()]
    total = sum(ns)
    log(f"loaded {len(df)} problems, {total} items "
        f"(min={min(ns)} max={max(ns)} mean={total / len(ns):.1f})")

    fallbacks = {i: source_fallback(r.get("query", ""), n)
                 for i, (_, r), n in zip(ids, df.iterrows(), ns)}

    # 1. A complete, correctly-shaped submission before the model even loads.
    preds = {i: list(fallbacks[i]) for i in ids}
    explanations = {i: "" for i in ids} if WANT_EXPLANATION else None
    write_submission(OUT_CSV, ids, preds, explanations)
    log(f"wrote placeholder {OUT_CSV} ({len(ids)} rows)")

    # 2. Load model.
    import torch
    from transformers import (AutoTokenizer, AutoModelForCausalLM,
                              StoppingCriteria, StoppingCriteriaList)

    class Deadline(StoppingCriteria):
        def __init__(self, stop_at):
            self.stop_at = stop_at

        def __call__(self, input_ids, scores, **kw):
            return time.time() > self.stop_at

    log("loading tokenizer/model ...")
    tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    tok.padding_side = "left"

    def _load(dev_map):
        try:
            return AutoModelForCausalLM.from_pretrained(
                MODEL_ID, torch_dtype=torch.float16, device_map=dev_map,
                trust_remote_code=True).eval()
        except TypeError:
            return AutoModelForCausalLM.from_pretrained(
                MODEL_ID, dtype=torch.float16, device_map=dev_map,
                trust_remote_code=True).eval()

    try:
        # Pin every layer to the GPU explicitly. device_map="auto" is free to
        # spill a layer or two to CPU when VRAM looks tight, which is nearly
        # invisible in logs but makes generation roughly two orders of
        # magnitude slower -- the worst possible failure mode under a hard
        # deadline. Only fall back to "auto" if the explicit placement fails
        # outright.
        model = _load({"": 0} if torch.cuda.is_available() else "auto")
    except Exception as e:
        log(f"pinned load failed ({type(e).__name__}: {e}); falling back to auto")
        model = _load("auto")

    devs = set(str(p.device) for p in model.parameters())
    log(f"model ready on {sorted(devs)} ({time_left():.0f}s of budget left)")
    if any(d.startswith("cpu") or d == "meta" for d in devs):
        log("WARNING: part of the model is off-GPU; generation will be very slow")
    if torch.cuda.is_available():
        log(f"  VRAM allocated {torch.cuda.memory_allocated()/1e9:.2f} GB / "
            f"{torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")

    prompts = []
    for (_, r), n in zip(df.iterrows(), ns):
        system = FAITHFUL_SYSTEM if PROMPT_STYLE == "faithful" else GUIDED_SYSTEM
        msgs = [{"role": "system", "content": system},
                {"role": "user", "content": build_prompt(r, n)}]
        prompts.append(tok.apply_chat_template(msgs, tokenize=False,
                                               add_generation_prompt=True))

    # Probe actual prompt lengths so the starting batch size reflects this
    # file's problems instead of a fixed guess that either wastes headroom on
    # short prompts or guarantees an OOM on long ones.
    sample_lens = [len(tok.encode(p, add_special_tokens=False))
                   for p in prompts[: min(8, len(prompts))]]
    approx_max_len = max(sample_lens) if sample_lens else 512
    default_batch = 4 if approx_max_len < 2500 else (2 if approx_max_len < 4500 else 1)
    batch_size = int(os.environ.get("IOL_BATCH", str(default_batch)))
    log(f"prompt length probe: ~{approx_max_len} tok max of {len(sample_lens)} "
        f"sampled -> starting batch_size={batch_size}")

    def generate(texts, max_new, sample, temp=0.7):
        nonlocal batch_size
        out = [""] * len(texts)
        order = sorted(range(len(texts)), key=lambda i: len(texts[i]))
        i = 0
        while i < len(order):
            if time_left() < 25:
                log("  out of time inside generate(); returning partial batch")
                break
            idx = order[i:i + batch_size]
            chunk = [texts[j] for j in idx]
            try:
                enc = tok(chunk, return_tensors="pt", padding=True,
                          truncation=True, max_length=6144).to(model.device)
                # repetition_penalty=1.0 explicitly. Several AWQ chat models
                # ship a generation_config.json with a >1.0 repetition
                # penalty baked in, and unlike temperature/top_p (which
                # greedy decoding ignores, with a warning) a repetition
                # penalty is applied under greedy decoding silently. A large
                # share of gold answers in agglutinative languages legitimately
                # repeat a character run, so leaving the shipped penalty in
                # place pushes generation away from exactly the strings
                # needed.
                kw = dict(max_new_tokens=max_new, pad_token_id=tok.pad_token_id,
                          repetition_penalty=1.0,
                          stopping_criteria=StoppingCriteriaList(
                              [Deadline(DEADLINE - 10)]))
                if sample:
                    kw.update(do_sample=True, temperature=temp, top_p=0.95)
                else:
                    kw.update(do_sample=False)
                with torch.no_grad():
                    o = model.generate(**enc, **kw)
                for k, j in enumerate(idx):
                    out[j] = tok.decode(o[k][enc["input_ids"].shape[1]:],
                                        skip_special_tokens=True)
                i += batch_size
            except torch.cuda.OutOfMemoryError:
                torch.cuda.empty_cache()
                if batch_size == 1:
                    log("  OOM at batch=1; skipping this item")
                    i += 1
                else:
                    batch_size = max(1, batch_size // 2)
                    log(f"  OOM -> batch_size={batch_size}")
            except Exception as e:
                log(f"  generate error: {type(e).__name__}: {e}")
                i += batch_size
        return out

    def solve_by_assignment(row, n):
        """Score every (item, option-letter) pair via next-token logits and
        take the best one-to-one assignment, instead of trusting free-form
        text for match_letters problems (which tends to just emit A, B, C...
        in order -- a valid permutation, so no repair catches it, and it
        scores near zero).
        Raw next-token logprobs are known to carry a systematic per-token
        prior independent of content (Zhao et al. 2021, "Calibrate Before
        Use") -- e.g. a model may just prefer the token for "C" regardless of
        which option is actually correct. We estimate that per-letter prior
        with a content-free version of the same prompt and subtract it before
        assignment, so the score reflects evidence from the item text rather
        than a baked-in letter preference.
        """
        items, opts = split_matching_block(row.get("context", ""))
        if len(items) < 3 or len(opts) < 3 or len(items) != n:
            return None
        letters = [o[0] for o in opts]
        cand_ids = []
        for L in letters:
            token_ids = set()
            for form in (L, " " + L):
                enc = tok.encode(form, add_special_tokens=False)
                if enc:
                    token_ids.add(enc[0])
            cand_ids.append(sorted(token_ids))

        ctx = row["context"].strip()

        def _logprob_batch(msgs_list):
            """Run one forward pass per prompt in msgs_list, return a list of
            per-candidate-letter logprob lists (same shape as `scores` rows)."""
            texts = [tok.apply_chat_template(m, tokenize=False,
                                              add_generation_prompt=True)
                     for m in msgs_list]
            out = []
            step = 4
            for s0 in range(0, len(texts), step):
                if time_left() < 30:
                    return None
                chunk = texts[s0:s0 + step]
                enc = tok(chunk, return_tensors="pt", padding=True,
                          truncation=True, max_length=6144).to(model.device)
                with torch.no_grad():
                    logits = model(**enc).logits[:, -1, :].float()
                logp = torch.log_softmax(logits, dim=-1)
                for b in range(len(chunk)):
                    out.append([max(logp[b, i].item() for i in ids_)
                                for ids_ in cand_ids])
            return out

        # --- real per-item prompts -------------------------------------
        q_msgs = []
        for num, itext in items:
            q_msgs.append([
                {"role": "system", "content":
                 "You match items to their correct counterparts in a "
                 "linguistics problem. Reply with one option letter only."},
                {"role": "user", "content":
                 f"{ctx}\n\nWhich lettered option corresponds to item {num} "
                 f"({itext})? Reply with the option letter only."},
            ])
        scores = _logprob_batch(q_msgs)
        if scores is None:
            return None

        # --- content-free baseline prompt -------------------------------
        # Same instruction and option list (letters need to stay visible
        # since the token candidates depend on the model's tokenization in
        # context), but with the item text and its identifying number
        # blanked out, so the only thing driving the letter choice is the
        # model's own prior over "A"/"B"/"C"... in this response format --
        # exactly the confound we want to remove from `scores`.
        opts_block = "\n".join(f"{L}. {t}" for L, t in opts)
        baseline_msgs = [{"role": "system", "content":
                          "You match items to their correct counterparts in a "
                          "linguistics problem. Reply with one option letter only."},
                         {"role": "user", "content":
                          f"{opts_block}\n\nWhich lettered option corresponds "
                          f"to item ___ (___)? Reply with the option letter "
                          f"only."}]
        baseline_scores = _logprob_batch([baseline_msgs])
        if baseline_scores is None:
            # Baseline pass failed/timed out -- fall back to uncalibrated
            # scores rather than losing the whole solver over one extra call.
            log("  calibration baseline unavailable; using raw scores")
            calibrated = scores
        else:
            baseline = baseline_scores[0]  # one row, per-letter logprob
            calibrated = [[s - baseline[j] for j, s in enumerate(row_scores)]
                          for row_scores in scores]

        cols = assign_max_weight(calibrated)
        return [letters[c] for c in cols]

    def apply_parsers(row_id, row, n, text):
        answer = enforce_bijection(parse_model_output(text, n, fallbacks[row_id]))
        task_type = (row.get("task_type") or "").strip().lower()
        if task_type == "match_letters":
            try:
                mm = solve_by_assignment(row, n)
                if mm and len(mm) == n:
                    return mm
            except Exception as e:
                log(f"  assignment solver failed on {row_id}: {type(e).__name__}: {e}")
        if task_type == "text_to_num":
            answer = [coerce_numeral(a, task_type) for a in answer]
        return answer

    # 3. Pass 1: greedy decoding, guarantees a complete answer set.
    TOK_PER_S = float(os.environ.get("IOL_TOKS", "30"))
    adaptive = int(0.40 * max(1.0, time_left()) * TOK_PER_S / max(1, len(df)))
    max_new = max(192, min(MAX_NEW, adaptive))
    log(f"reasoning budget: {max_new} new tokens/problem "
        f"(adaptive={adaptive}, cap={MAX_NEW}, {len(df)} problems)")

    t0 = time.time()
    texts = generate(prompts, max_new=max_new, sample=False)
    pass1_cost = time.time() - t0

    samples = {i: [] for i in ids}
    match_rows = 0
    for (row_id, n, txt), (_, row) in zip(zip(ids, ns, texts), df.iterrows()):
        a = apply_parsers(row_id, row, n, txt)
        if (row.get("task_type") or "").strip().lower() == "match_letters":
            match_rows += 1
        preds[row_id] = a
        samples[row_id].append(a)
    if match_rows:
        log(f"assignment solver attempted on {match_rows} match_letters problem(s)")
    write_submission(OUT_CSV, ids, preds, explanations)
    no_block = sum(1 for t in texts if not re.search(r"answers?\s*:", t or "", re.I))
    empty = sum(1 for t in texts if not (t or "").strip())
    log(f"pass 1 (greedy) done in {pass1_cost:.0f}s -> submission written "
        f"({no_block}/{len(texts)} without an ANSWERS: block, {empty} empty)")
    report_dev_score(preds)

    # 4. Convergence-aware self-consistency: only re-sample rows whose items
    # don't yet have a confident majority, so budget concentrates on the rows
    # that are actually still undecided instead of re-rolling settled ones.
    def row_converged(row_id):
        s = samples[row_id]
        if len(s) < 3:
            return False
        for k in range(len(s[0])):
            col = [x[k] for x in s if k < len(x)]
            groups = cluster_candidates(col)
            if not any(len(g) >= 2 for g in groups):
                return False
        return True

    reserve = 0.0
    if WANT_EXPLANATION:
        reserve = min(300.0, 0.25 * pass1_cost + 60)

    active_ids = list(ids)
    n_extra = 0
    while time_left() - reserve > pass1_cost * 0.5 and n_extra < MAX_SAMPLES and active_ids:
        n_extra += 1
        active_ids = [i for i in active_ids if not row_converged(i)]
        if not active_ids:
            log("  every row has converged; stopping self-consistency early")
            break
        log(f"self-consistency pass {n_extra} on {len(active_ids)}/{len(ids)} "
            f"undecided rows ({time_left():.0f}s left)")
        pos = {rid: k for k, rid in enumerate(ids)}
        sub_prompts = [prompts[pos[i]] for i in active_ids]
        texts = generate(sub_prompts, max_new=max_new, sample=True, temp=SAMPLE_TEMP)
        row_lookup = {str(rid): row for rid, (_, row) in zip(ids, df.iterrows())}
        n_lookup = dict(zip(ids, ns))
        for rid, txt in zip(active_ids, texts):
            if txt:
                a = apply_parsers(rid, row_lookup[rid], n_lookup[rid], txt)
                samples[rid].append(a)
        for rid in ids:
            if len(samples[rid]) >= 3:
                greedy = samples[rid][0]
                n = n_lookup[rid] if rid in n_lookup else len(greedy)
                preds[rid] = enforce_bijection([
                    vote_answer([s[k] for s in samples[rid] if k < len(s)],
                                anchor=greedy[k] if k < len(greedy) else None)
                    for k in range(n)
                ])
        write_submission(OUT_CSV, ids, preds, explanations)
        log(f"  voted over up to {n_extra + 1} samples/row -> written")
        report_dev_score(preds)

    # 5. Explanations for the jury track.
    if WANT_EXPLANATION and time_left() > 60:
        log(f"generating explanations ({time_left():.0f}s left)")
        ex_prompts = []
        for (_, r), i in zip(df.iterrows(), ids):
            msgs = [{"role": "system", "content": EXPLAIN_SYSTEM},
                    {"role": "user", "content": build_explain_prompt(r, preds[i])}]
            ex_prompts.append(tok.apply_chat_template(
                msgs, tokenize=False, add_generation_prompt=True))
        ex_texts = generate(ex_prompts, max_new=200, sample=False)
        for i, e in zip(ids, ex_texts):
            e = re.sub(r"\s+", " ", (e or "").strip())
            if e:
                explanations[i] = e[:1200]
        write_submission(OUT_CSV, ids, preds, explanations)
        log("explanations written")

    # 6. Final integrity pass: guarantee shape and non-blank answers no
    # matter what happened above.
    n_lookup = dict(zip(ids, ns))
    malformed = [i for i in ids if len(preds[i]) != n_lookup[i]
                 or any(not str(x).strip() for x in preds[i])]
    if malformed:
        log(f"repairing {len(malformed)} malformed rows")
        for i in malformed:
            preds[i] = squeeze_to_n(
                [x for x in preds[i] if str(x).strip()], n_lookup[i], fallbacks[i])
        write_submission(OUT_CSV, ids, preds, explanations)

    log(f"DONE. {len(ids)} rows, {sum(len(v) for v in preds.values())} answers, "
        f"{time.time() - START:.0f}s elapsed")


if __name__ == "__main__":
    main()