"""IOL-AI 2026 submission — v5 hybrid (solver + budget-managed gemma4:12b). Everything here was validated on a local dev set (see README): pass 0 symbolic numeral solver (numeral_solver.py, shipped in this repo); answers text_to_num / num_to_text exactly when the system fits the searched grammar family, returns None -> LLM fallback otherwise pass 1 fast low-token answer for every remaining problem, submission.csv atomically checkpointed after every row (a kill never leaves an invalid/partial file) pass 2 per-row time-sliced re-solve with reasoning; overwrites baseline only if the result parses to the right number of answers; translation reasoning is deliberately capped LOW (truncated reasoning + a forced short answer scored higher than completed reasoning on gemma4:12b) Repo layout expected: this file as script.py, numeral_solver.py beside it, gemma4:12b weights in the repo root (load from "."). fp16 12B does not fit a 16 GB T4, so weights load 4-bit via bitsandbytes (the organizer-reference recipe). """ import os import subprocess import sys import time T0 = time.monotonic() TIME_LIMIT = 30 * 60 DEADLINE = T0 + TIME_LIMIT - 3 * 60 # 3-min reserve for writes/exit # Smoke mode (local CPU shakeout, mirrors the leader's IOL_DUMMY pattern): # IOL_SMOKE=1 skip bitsandbytes (no CUDA), load fp32 on CPU # IOL_MODEL_ID=... substitute a tiny model for the shipped weights # IOL_INPUT=... read a local CSV instead of /tmp/data/test.csv SMOKE = os.environ.get("IOL_SMOKE", "0") == "1" deps = ["transformers>=4.51", "accelerate>=0.30", "torch>=2.2", "pandas"] if not SMOKE: deps.append("bitsandbytes") subprocess.run([sys.executable, "-m", "pip", "install", "-q", *deps], check=True) import json import re import pandas as pd import torch from transformers import AutoModelForCausalLM, AutoTokenizer import numeral_solver MODEL_ID = os.environ.get("IOL_MODEL_ID", ".") INPUT_CSV = os.environ.get("IOL_INPUT", "/tmp/data/test.csv") MARKER = "FINAL ANSWERS" TRANSLATION_CAP = 3000 # tokens; see module docstring GLOBAL_CAP = 8192 BASE_RULES = ( f"You solve International Linguistics Olympiad problems. Your reasoning budget " f"is limited, so be systematic and compact — do not second-guess a hypothesis " f"that fits all the data. Verify against every given example once, then commit. " f"End with the line '{MARKER}:' followed by one answer per line, in item order, " f"with no numbering and no extra text." ) NUMERAL_CORE = ( "1) For EVERY example, write one arithmetic equation showing exactly how its " "words produce its value. 2) Determine the base from the single-word values. " "3) CRITICAL: find pairs of examples using the same words in different orders " "with different values — decide which order means multiplication and which " "means addition. 4) Only after your equations reproduce ALL examples, " ) STRATEGIES = { "translation": ( "Method: 1) Align the given sentence pairs and segment every word into " "morphemes by comparing entries that share meaning components. 2) Write a " "compact table: each root, prefix, and suffix with its meaning, plus the " "morpheme order. 3) Compose each requested item from the table. Mind the " "direction of translation asked for." ), "text_to_num": "Method: " + NUMERAL_CORE + "convert each item to digits.", "num_to_text": "Method: " + NUMERAL_CORE + "compose each requested number in " "the puzzle language.", "fill_blanks": ( "Method: deduce the paradigm from the completed cells, state the rule for " "each row/column, then fill each blank consistently with it." ), "match_letters": ( "Method: find anchor items you can pair with certainty first (word length, " "repeated letters), then use elimination. Every item gets exactly one match." ), } SALVAGE = ( "Time is up. Based on the partial analysis below, give your best answer for " "all {n} items RIGHT NOW: one answer per line, in order, no numbering, no " "other text, no further reasoning.\n\nThe items to answer:\n{query}\n\n" "Partial analysis:\n{tail}" ) def remaining(): return DEADLINE - time.monotonic() tok = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=not SMOKE) if SMOKE: model = AutoModelForCausalLM.from_pretrained(MODEL_ID).eval() else: from transformers import BitsAndBytesConfig model = AutoModelForCausalLM.from_pretrained( MODEL_ID, quantization_config=BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16), device_map="auto", local_files_only=True, ).eval() print(f"[model] loaded at {time.monotonic()-T0:.0f}s", flush=True) def generate(messages, max_new, max_time=None): """Greedy generation. Returns (text, truncated, tokens_per_sec).""" enc = tok.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", ) if hasattr(enc, "keys"): # BatchEncoding on newer transformers inputs = {k: v.to(model.device) for k, v in enc.items()} else: # bare tensor on older versions inputs = {"input_ids": enc.to(model.device)} prompt_len = inputs["input_ids"].shape[-1] kw = {"max_new_tokens": max_new, "do_sample": False, "pad_token_id": tok.eos_token_id} if max_time: kw["max_time"] = max_time t0 = time.monotonic() with torch.no_grad(): out = model.generate(**inputs, **kw) n_new = out.shape[-1] - prompt_len tps = n_new / max(time.monotonic() - t0, 1e-6) text = tok.decode(out[0][prompt_len:], skip_special_tokens=True).strip() return text, n_new >= max_new, tps def item_count(query): return len(re.findall(r"^\s*\d+[.)]", query, re.M)) or 1 def extract_answers(text, n): m = re.search(rf"{MARKER}\s*:?", text, re.I) block = text[m.end():] if m else text lines = [ln.strip() for ln in block.splitlines() if ln.strip()] if not m: lines = lines[-n:] lines = [re.sub(r"^\d+[.)]\s*", "", ln) for ln in lines] return (lines + [""] * n)[:n] def valid(answers, n): return len(answers) == n and all(answers) def checkpoint(rows): tmp = ".submission.csv.tmp" pd.DataFrame(rows, columns=["id", "pred"]).to_csv(tmp, index=False) os.replace(tmp, "submission.csv") def build_messages(r): system = BASE_RULES + "\n\n" + STRATEGIES.get(r["task_type"], "") return [ {"role": "system", "content": system}, {"role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}"}, ] df = pd.read_csv(INPUT_CSV, dtype=str).fillna("") problems = [r for _, r in df.iterrows()] counts = [item_count(r["query"]) for r in problems] rows = [{"id": r["id"], "pred": json.dumps([""] * counts[i])} for i, r in enumerate(problems)] checkpoint(rows) # ── pass 0: symbolic numeral solver ────────────────────────────────────────── solved = [False] * len(problems) for i, r in enumerate(problems): if r["task_type"] in ("text_to_num", "num_to_text"): try: ans = numeral_solver.answer(r["context"], r["query"], r["task_type"]) except Exception as exc: print(f"[solver] {r['id']} error: {exc}", flush=True) ans = None if ans is not None: rows[i]["pred"] = json.dumps(ans, ensure_ascii=False) solved[i] = True print(f"[solver] {r['id']} solved", flush=True) checkpoint(rows) # ── pass 1: fast complete baseline ─────────────────────────────────────────── tps_est = 20.0 for i, r in enumerate(problems): if solved[i] or remaining() < 60: continue text, _, tps_est = generate(build_messages(r), max_new=48 * counts[i] + 128, max_time=min(90.0, remaining() / 4)) rows[i]["pred"] = json.dumps(extract_answers(text, counts[i]), ensure_ascii=False) checkpoint(rows) print(f"[base {i+1}/{len(problems)}] done at {time.monotonic()-T0:.0f}s", flush=True) # ── pass 2: time-sliced reasoning upgrades ─────────────────────────────────── for i, r in enumerate(problems): if solved[i]: continue n = counts[i] if remaining() < 30: print("[deadline] stopping upgrades", flush=True) break todo = sum(1 for j in range(i, len(problems)) if not solved[j]) slice_s = max(remaining() / max(todo, 1), 10.0) max_new = int(min(max(slice_s * 0.8 * tps_est, 512), GLOBAL_CAP)) if r["task_type"] == "translation": max_new = min(max_new, TRANSLATION_CAP) text, truncated, tps_est = generate(build_messages(r), max_new, max_time=slice_s * 0.8) upgraded = extract_answers(text, n) if text else [] if not valid(upgraded, n) and remaining() > 20: msgs = [ {"role": "system", "content": "Output only the answers, one per line."}, {"role": "user", "content": SALVAGE.format( n=n, query=r["query"].strip(), tail=text[-3000:])}, ] text, _, _ = generate(msgs, max_new=48 * n + 64, max_time=min(60.0, remaining())) upgraded = extract_answers(text, n) if valid(upgraded, n): rows[i]["pred"] = json.dumps(upgraded, ensure_ascii=False) checkpoint(rows) print(f"[think {i+1}/{len(problems)}] upgraded", flush=True) else: print(f"[think {i+1}/{len(problems)}] kept baseline", flush=True) checkpoint(rows) print(f"wrote submission.csv at {time.monotonic()-T0:.0f}s", flush=True)