script.py

#!/usr/bin/env python
import os, re, json, time, unicodedata
from collections import Counter, defaultdict

T0 = time.time()
TIME_LIMIT = float(os.environ.get("IOL_TIME_LIMIT", "1800"))
SAFETY     = float(os.environ.get("IOL_SAFETY", "150"))
DEADLINE   = T0 + 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_EXPL  = os.environ.get("IOL_EXPLAIN", "1") == "1"
REASON_TOK = int(os.environ.get("IOL_MAXNEW", "900"))        # reasoning budget/item
VOTE_CAP   = int(os.environ.get("IOL_MAXSAMPLES", "12"))     # Phi-4 is fast -> vote a lot
VOTE_TEMP  = float(os.environ.get("IOL_TEMP", "0.5"))
BATCH      = int(os.environ.get("IOL_BATCH", "4"))
TOK_PER_S  = float(os.environ.get("IOL_TOKS", "40"))         # Phi-4 throughput est.
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(m): print(f"[{time.time()-T0:7.1f}s] {m}", flush=True)
def left(): return DEADLINE - time.time()

# --- how many items a problem asks for ---
_LN = re.compile(r"^[ \t]*(\d{1,3})[.)\]]", re.M)
_PN = re.compile(r"\((\d{1,3})\)")
_RG = re.compile(r"\(?(\d{1,3})\s*(?:[-–—]|to)\s*(\d{1,3})\)?")
_LL = re.compile(r"^[ \t]*([A-Z])[.)\]]\s", re.M)
def count_items(query, ctx=""):
    q = query or ""; rng = 0
    for a, b in _RG.findall(q):
        a, b = int(a), int(b)
        if 0 < b - a < 60: rng = max(rng, b - a + 1)
    cand = max(len(set(_LN.findall(q))), len(set(_PN.findall(q))), len(set(_LL.findall(q))))
    n = max(rng, cand)
    if n > 1: return n
    lines = [l.strip() for l in q.splitlines() if l.strip()]
    if len(lines) > 1:
        body = lines[1:] if lines[0].endswith((":", ".")) else lines
        if body: return len(body)
    if ctx:
        cn = len(set(_LN.findall(ctx)))
        if cn > 1: return cn
        cl = len(set(_LL.findall(ctx)))
        if cl > 1: return cl
    return max(n, 1)

def item_sources(query, n, task=""):
    if (task or "").strip().lower() == "match_letters": return ["A"] * n
    q = query or ""; out = []
    for ln in q.splitlines():
        m = re.match(r"^\(?(\d{1,3})\)?[.):\]]\s*(.+)$", ln.strip())
        if m: out.append(m.group(2).strip())
    if not out:
        lines = [l.strip() for l in q.splitlines() if l.strip()]
        if len(lines) > 1 and lines[0].endswith((":", ".")): out = lines[1:]
    out = [o.split("|")[0].strip() if "|" in o else o for o in out]
    out = [o for o in out if o]
    while len(out) < n: out.append(out[-1] if out else "?")
    return out[:n]

# --- prompt: reason, then a clean ANSWERS: block, shaped by a per-task hint ---
SYS = ("You are a top competitor at the International Linguistics Olympiad. Each problem "
       "gives data from a language you have never seen; everything you need is inside the "
       "problem. Line up the examples, segment the words, find the recurring morphemes and "
       "the rules ordering them, and check them against every example. Reason briefly, then "
       "write a line that is exactly ANSWERS: and, below it, one answer per item in the order "
       "asked — no numbering, no glosses, no commentary, no blank lines. Never leave one blank.")
HINT = {
    "translation":   "Each answer is the translated form alone — no source word, no gloss, no quotes.",
    "match_letters": "Each answer is one capital letter; every letter is used exactly once, none repeats.",
    "fill_blanks":   "Each answer is only the missing form for that blank — not the whole line, not the gloss.",
    "text_to_num":   "Each answer is digits only (for example 111).",
    "num_to_text":   "Each answer is the number written out in the problem's language, words only.",
}
def user_prompt(row, n):
    h = HINT.get((row.get("task_type") or "").strip().lower(), "")
    return (f"{str(row['context']).strip()}\n\n{str(row['query']).strip()}\n\n"
            f"There are exactly {n} item{'s' if n != 1 else ''} to answer."
            + (f" {h}" if h else "") +
            f"\nAfter reasoning, write ANSWERS: on its own line then exactly {n} "
            f"line{'s' if n != 1 else ''}, one answer per item, in order.")

# --- parsing the ANSWERS block ---
_PFX  = re.compile(r"^\s*(?:\(?\d{1,3}\)?[.):\]]\s*|[-*•]\s+)")
_FENCE = re.compile(r"^```")
_PROSE = re.compile(r"^\s*(?:here (?:are|is)\b|answers?\s*:?\s*$|explanation\b|note\b|okay\b|"
    r"solution\b|reasoning\b|analysis\b|the answers?\b|let me\b|first,|so,|therefore\b|thus\b)", re.I)
_BRK = re.compile(r"\[[^\[\]\n]*\]")
def clean(s):
    s = _PFX.sub("", s.strip()).strip().strip("`").strip()
    if len(s) >= 2 and s[0] == s[-1] and s[0] in "\"'“”": s = s[1:-1].strip()
    return s.strip()
def shape(val, task):
    t = (task or "").strip().lower()
    if t == "match_letters":
        m = re.search(r"(?<![A-Za-z])([A-Z])(?![A-Za-z])", val)
        return m.group(1) if m else val
    if t == "fill_blanks":
        m = _BRK.search(val)
        if m: return m.group(0)
    return val.split("|")[0].strip() if "|" in val else val
def fit_to_n(items, n, fb=None):
    items = [i for i in items if i and i.strip()]
    if len(items) > n: items = items[-n:]
    while len(items) < n:
        items.append(fb[len(items)] if fb and len(items) < len(fb) else (items[-1] if items else "?"))
    return items[:n]
def parse(text, n, task, fb=None):
    if not text: return list(fb[:n]) if fb else ["?"] * n
    m = None
    for m2 in re.finditer(r"(?:^|\n)\s*(?:final\s+)?answers?\s*:\s*\n?", text, re.I): m = m2
    body = text[m.end():] if m else text
    numbered, raw = [], []
    for ln in body.splitlines():
        if _FENCE.match(ln): continue
        mm = re.match(r"^\s*\(?(\d{1,3})\)?[.):\]]\s*(.+)$", ln.strip())
        if mm:
            v = clean(mm.group(2))
            if v and not _PROSE.match(v): numbered.append((int(mm.group(1)), v))
        c = clean(ln)
        if c and not _PROSE.match(c): raw.append(c)
    if len(numbered) >= n:
        by = {}
        for lab, v in numbered: by[lab] = v
        labs = sorted(by)
        if len(labs) >= n: raw = [by[l] for l in labs[:n]]
    vals = [shape(v, task) for v in fit_to_n(raw, n, fb)]
    return fit_to_n(vals, n, fb)

def norm(s):
    s = unicodedata.normalize("NFC", (s or "").strip().lower()); s = re.sub(r"\s+", " ", s)
    return s.strip(" .!?;:,")
def vote(cands, anchor):
    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
    g = defaultdict(list)
    for c in cands: g[norm(c)].append(c)
    a_sup = len(g.get(norm(anchor), [])); bk, bn = None, 0
    for k, v in g.items():
        if len(v) > bn: bk, bn = k, len(v)
    return Counter(g[bk]).most_common(1)[0][0] if (bk and bn >= 2 and bn > a_sup) else anchor

EXPL_SYS = ("You explain International Linguistics Olympiad solutions to a human judge: state "
            "the key rules — morphemes, word order, sound changes — concisely (2-4 sentences). "
            "Do not restate the reasoning.")

def write_out(path, ids, preds, expl):
    import pandas as pd
    rows = [{"id": i, "pred": json.dumps(preds[i], ensure_ascii=False),
             "explanation": (expl or {}).get(i, "")} for i in ids]
    pd.DataFrame(rows).to_csv(path, index=False)

def main():
    import pandas as pd, torch
    from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria, StoppingCriteriaList
    torch.backends.cuda.matmul.allow_tf32 = True; torch.backends.cudnn.allow_tf32 = True

    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("context", "")) for _, r in df.iterrows()]
    tasks = {i: (r.get("task_type", "") or "") for i, (_, r) in zip(ids, df.iterrows())}
    srcs = {i: item_sources(r.get("query", ""), n, r.get("task_type", "")) for i, (_, r), n in zip(ids, df.iterrows(), ns)}
    log(f"loaded {len(df)} problems, {sum(ns)} items")

    preds = {i: list(srcs[i]) for i in ids}; expl = {i: "" for i in ids} if WANT_EXPL else None
    write_out(OUT_CSV, ids, preds, expl); log("placeholder written")

    class Deadline(StoppingCriteria):
        def __init__(self, t): self.t = t
        def __call__(self, i, s, **k): return time.time() > self.t

    log("loading 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(dm):
        try: return AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float16, device_map=dm, trust_remote_code=True).eval()
        except TypeError: return AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float16, device_map=dm, trust_remote_code=True).eval()
    try: model = _load({"": 0} if torch.cuda.is_available() else "auto")
    except Exception as e: log(f"pinned load failed ({e}); auto"); model = _load("auto")
    log(f"model ready ({left():.0f}s left)")

    prompts = [tok.apply_chat_template(
        [{"role": "system", "content": SYS}, {"role": "user", "content": user_prompt(r, n)}],
        tokenize=False, add_generation_prompt=True) for (_, r), n in zip(df.iterrows(), ns)]

    bs = BATCH
    def gen(texts, max_new, sample):
        nonlocal bs
        out = [""] * len(texts); order = sorted(range(len(texts)), key=lambda i: len(texts[i])); i = 0
        while i < len(order):
            if left() < 25: break
            idx = order[i:i+bs]; chunk = [texts[j] for j in idx]
            try:
                enc = tok(chunk, return_tensors="pt", padding=True, truncation=True, max_length=6144).to(model.device)
                kw = dict(max_new_tokens=max_new, pad_token_id=tok.pad_token_id, repetition_penalty=1.0,
                          stopping_criteria=StoppingCriteriaList([Deadline(DEADLINE - 10)]))
                kw.update(dict(do_sample=True, temperature=VOTE_TEMP, top_p=0.95) if sample else dict(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 += bs
            except torch.cuda.OutOfMemoryError:
                torch.cuda.empty_cache(); bs = max(1, bs // 2) if bs > 1 else 1
                if bs == 1: i += 0
            except Exception as e: log(f"gen error: {type(e).__name__}: {e}"); i += bs
        return out

    adaptive = int(0.40 * max(1.0, left()) * TOK_PER_S / max(1, len(df)))
    max_new = max(192, min(REASON_TOK, adaptive))
    log(f"Pass 1 greedy, {max_new} tok/item")
    t = time.time(); texts = gen(prompts, max_new, sample=False); c1 = time.time() - t
    samples = {i: [] for i in ids}
    for i, n, txt in zip(ids, ns, texts):
        preds[i] = parse(txt, n, tasks[i], srcs[i]); samples[i].append(preds[i])
    write_out(OUT_CSV, ids, preds, expl)
    no_block = sum(1 for t2 in texts if not re.search(r"answers?\s*:", t2 or "", re.I))
    log(f"Pass 1 done in {c1:.0f}s ({no_block}/{len(texts)} had no ANSWERS: block)")

    reserve = min(300.0, 0.25 * c1 + 60) if WANT_EXPL else 30.0
    ne = 0
    while left() - reserve > c1 * 1.25 and ne < VOTE_CAP:
        ne += 1; log(f"vote pass {ne} ({left():.0f}s left)")
        texts = gen(prompts, max_new, sample=True)
        for i, n, txt in zip(ids, ns, texts):
            if txt: samples[i].append(parse(txt, n, tasks[i], srcs[i]))
        for i, n in zip(ids, ns):
            if len(samples[i]) >= 3:
                g = samples[i][0]
                preds[i] = fit_to_n([vote([s[k] for s in samples[i] if k < len(s)],
                                          g[k] if k < len(g) else None) for k in range(n)], n, srcs[i])
        write_out(OUT_CSV, ids, preds, expl); log(f"voted over {ne+1} samples")
    log(f"self-consistency: {ne} vote pass(es)")

    if WANT_EXPL and left() > 60:
        ep = [tok.apply_chat_template([{"role": "system", "content": EXPL_SYS},
            {"role": "user", "content": f"{str(r['context']).strip()}\n\n{str(r['query']).strip()}\n\n"
             f"Answers:\n" + "\n".join(f"- {a}" for a in preds[str(r['id'])]) + "\n\nExplain the rules briefly."}],
            tokenize=False, add_generation_prompt=True) for _, r in df.iterrows()]
        for i, e in zip(ids, gen(ep, 200, sample=False)):
            e = re.sub(r"\s+", " ", (e or "").strip())
            if e: expl[i] = e[:1200]
        write_out(OUT_CSV, ids, preds, expl); log("explanations written")

    for i, n in zip(ids, ns):
        if len(preds[i]) != n or any(not str(x).strip() for x in preds[i]):
            preds[i] = fit_to_n([x for x in preds[i] if str(x).strip()], n, srcs[i])
    write_out(OUT_CSV, ids, preds, expl)
    log(f"DONE. {len(ids)} rows, {time.time()-T0:.0f}s.")

if __name__ == "__main__":
    main()