"""IOL-AI 2026 submission — v9-refine (same Qwen weights). Two-pass directed self-critique. Pass 1 is v5's exact greedy run (drafts written to submission.csv as a safety net). Pass 2 shows the model its own draft answers alongside the problem and asks it to verify each one against the data — form, language, one-to-one letter use, arithmetic — and re-emit corrected FINAL ANSWERS. The draft is kept whenever the refine pass fails to produce a parseable answer list or time runs short. """ import os os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import json import re import time import pandas as pd import torch from transformers import AutoModelForCausalLM, AutoTokenizer TAG = "v9-refine" # submission tag, printed into the eval log MODEL_ID = "." T_START = time.monotonic() TIME_BUDGET_S = 27 * 60 MAX_NEW_TOKENS = 1280 # Minimum seconds that must remain before starting a refine generation. MIN_S_PER_GEN = 150 SYSTEM_PROMPT = ( "You solve International Linguistics Olympiad problems by reasoning from " "the data you are given. Everything needed is in the problem data. " "Reason step by step first: work out the vocabulary and rules from the " "examples, and double-check each answer against them. " "Each answer must be the same kind of form that appears in that position " "in the examples: if the examples show a phonetic transcription in " "[brackets], give a phonetic transcription in [brackets], not the English " "meaning; if they show a word or sentence in the problem language, answer " "in the problem language, copying its exact spelling conventions. For " "matching items give just the option letter -- the items and options form " "a one-to-one matching, so use each letter exactly once, never repeating " "or inventing letters. For number items give digits or the spelled-out " "number as the query asks. " "Then write a line that says exactly FINAL ANSWERS: and, below it, one " "answer per line in the order the items are asked -- the bare answer " "only, no numbering, no quotes, no extra text." ) REFINE_PROMPT = ( "Below is a linguistics problem and DRAFT answers produced on a first " "attempt. Carefully verify every draft answer against the problem data: " "re-derive the rules from the examples and check each answer follows " "them. Look especially for: answers in the wrong form or language " "compared to what fills that position in the examples; matching answers " "that repeat an option letter or use a letter that is not offered (the " "matching is one-to-one, each letter exactly once); number conversions " "that do not satisfy the number system of the data; and morphology that " "contradicts a pattern in the examples. Keep every draft answer that " "checks out, and correct the ones that do not. " "Then write a line that says exactly FINAL ANSWERS: and, below it, one " "answer per line in the same order -- the bare answer only, no " "numbering, no quotes, no extra text.\n\n" "{context}\n\n{query}\n\nDRAFT ANSWERS:\n{draft}" ) def count_items(query: str): """Count numbered items in a query; None when no numbering is detectable.""" nums = set() for m in re.finditer(r"\((\d{1,3})\s*[-–]\s*(\d{1,3})\)", query): a, b = int(m.group(1)), int(m.group(2)) if a < b and b - a < 40: nums.update(range(a, b + 1)) for m in re.finditer(r"^\s*(\d{1,3})[.)]\s", query, re.M): nums.add(int(m.group(1))) for m in re.finditer(r"\((\d{1,3})\)", query): nums.add(int(m.group(1))) nums = {n for n in nums if 1 <= n <= 999} return len(nums) if nums else None def parse_answers(text: str, n_items): """Lines after the last FINAL ANSWERS: marker; pad to n_items, never truncate.""" marker = None for marker in re.finditer(r"(?i)final answers?\s*:", text): pass if marker is not None: zone = text[marker.end() :] else: zone = text answers = [ln.strip() for ln in zone.splitlines() if ln.strip()] answers = [re.sub(r"^\s*\d{1,3}[.)]\s*", "", x) for x in answers] answers = [ x[1:-1].strip() if len(x) >= 2 and x[0] == x[-1] and x[0] in "'\"" else x for x in answers ] if marker is None and n_items is not None and len(answers) > n_items: answers = answers[-n_items:] if n_items is not None and len(answers) < n_items: answers = answers + [""] * (n_items - len(answers)) return answers def generate(tok, model, messages): chat = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = tok(chat, return_tensors="pt").to(model.device) with torch.no_grad(): out = model.generate( **inputs, max_new_tokens=MAX_NEW_TOKENS, do_sample=False, pad_token_id=tok.pad_token_id or tok.eos_token_id, ) return tok.decode( out[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=True ).strip() def write_submission(df, preds): rows = [ {"id": r["id"], "pred": json.dumps(preds[i], ensure_ascii=False)} for i, (_, r) in enumerate(df.iterrows()) ] pd.DataFrame(rows).to_csv("submission.csv", index=False) def main(): print(f"submission tag: {TAG}", flush=True) print("loading model...", flush=True) tok = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=torch.float16, # T4 has no bfloat16 device_map="auto", attn_implementation="sdpa", ).eval() print(f"model loaded in {time.monotonic() - T_START:.0f}s", flush=True) df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("") n_items_list = [count_items(r["query"]) for _, r in df.iterrows()] # Pass 1: greedy drafts (identical to the proven v5 run). drafts = [] for i, (_, r) in enumerate(df.iterrows()): remaining = TIME_BUDGET_S - (time.monotonic() - T_START) if remaining > 30: text = generate( tok, model, [ {"role": "system", "content": SYSTEM_PROMPT}, { "role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}", }, ], ) drafts.append(parse_answers(text, n_items_list[i])) else: drafts.append([""] * (n_items_list[i] or 0)) print( f"draft {i + 1}/{len(df)} done (elapsed={time.monotonic() - T_START:.0f}s)", flush=True, ) # Safety net: a valid submission exists from here on. write_submission(df, drafts) print("draft submission.csv written", flush=True) # Pass 2: verify-and-correct each draft, time permitting. finals = list(drafts) try: for i, (_, r) in enumerate(df.iterrows()): remaining = TIME_BUDGET_S - (time.monotonic() - T_START) if remaining < MIN_S_PER_GEN: print("time budget reached, keeping remaining drafts", flush=True) break if not any(drafts[i]): continue # nothing to refine draft_lines = "\n".join( f"{j + 1}. {a}" for j, a in enumerate(drafts[i]) if a ) text = generate( tok, model, [ {"role": "system", "content": SYSTEM_PROMPT}, { "role": "user", "content": REFINE_PROMPT.format( context=r["context"].strip(), query=r["query"].strip(), draft=draft_lines, ), }, ], ) refined = parse_answers(text, n_items_list[i]) if refined and any(refined): finals[i] = refined print( f"refine {i + 1}/{len(df)} done " f"(changed={refined != drafts[i]}, " f"elapsed={time.monotonic() - T_START:.0f}s)", flush=True, ) except Exception as e: # never let refinement cost us the submission print(f"refine pass aborted: {e}", flush=True) write_submission(df, finals) print("wrote submission.csv", flush=True) if __name__ == "__main__": main()