| |
| 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")) |
| VOTE_CAP = int(os.environ.get("IOL_MAXSAMPLES", "12")) |
| 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")) |
| 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() |
|
|
| |
| _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] |
|
|
| |
| 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.") |
|
|
| |
| _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() |
|
|
