"""IOL-AI 2026 submission script. Runs fully offline against /tmp/data/test.csv and writes submission.csv with columns: id, pred (JSON list of answer strings), explanation (omitted). """ import os os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" import time T_START = time.monotonic() import json import re from collections import Counter import pandas as pd import torch from transformers import AutoModelForCausalLM, AutoTokenizer # --- Configuration ----------------------------------------------------- # Production always loads from "." (weights shipped alongside this script in # the HF model repo). This constant only documents which Hub model those # weights came from, for keeping the Colab test notebook in sync. MODEL_ID_FOR_LOCAL_TEST = "Qwen/Qwen2.5-7B-Instruct-AWQ" # Fallback if Colab timing shows 7B has slack to spare: swap MODEL_ID_FOR_LOCAL_TEST # to "Qwen/Qwen2.5-14B-Instruct-AWQ" and re-run the notebook before re-uploading. MODEL_PATH = "." DATA_PATH = "/tmp/data/test.csv" OUTPUT_PATH = "submission.csv" TIME_BUDGET_SEC = 27 * 60 # 3-min safety margin under the 30-min hard limit BATCH_SIZE = 4 MIN_NEW_TOKENS = 128 MAX_NEW_TOKENS = 1024 # Self-consistency: sample k >= 1 completions per row and majority-vote per # answer position when the time budget has slack. k=1 is plain greedy decode # (current default, deterministic). See adaptive_k(). MIN_K = 1 MAX_K = 1 SAMPLE_TEMPERATURE = 0.7 # Which SYSTEM prompt variant to use — "full_cot" or "plain_io". Published # evidence on whether heavy chain-of-thought helps or hurts a 7B-class model # on this task type is mixed (modeLing, arXiv:2406.17038), so both variants # are kept and should be A/B tested on the Colab harness rather than assumed. PROMPT_VARIANT = "plain_io" # Worked example shared by both prompt variants: a made-up toy language (not # a real one) demonstrating (a) the answer must be a word-form in the target # language, never an English gloss, and (b) subject/object roles must be # re-derived per query, not copied from a similar-looking example. Directly # targets the two failure modes observed in real Colab runs (language # confusion — arXiv:2406.20052 shows few-shot examples largely eliminate it; # and swapped argument roles). _WORKED_EXAMPLE = """Worked example (a made-up toy language, unrelated to your actual task — the words and rules below apply ONLY to this example; never reuse them): CONTEXT: mota = "dog" mota-ta = "dog" (as object) suno = "cat" suno-ta = "cat" (as object) kire = "sees" mota suno-ta kire = "the dog sees the cat" suno mota-ta kire = "the cat sees the dog" QUERY: Translate into the made-up language: "the cat sees the dog" IDENTIFY PATTERNS: word order is SUBJECT OBJECT-ta VERB; subject is unmarked, object takes suffix "-ta", verb is always last. RULE TABLE: - subject noun -> bare noun, placed first - object noun -> noun + "-ta", placed second - verb -> placed last, unchanged TEST: "mota suno-ta kire" = dog(subj) cat-ta(obj) sees = "the dog sees the cat" — matches. "suno mota-ta kire" = cat(subj) dog-ta(obj) sees = "the cat sees the dog" — matches. APPLY: query is "the cat sees the dog" -> subject=cat=suno, object=dog=mota-ta, verb=kire. VERIFY: answer is a word-form in the made-up language, not an English gloss. Roles checked: cat is subject (first, unmarked), dog is object (second, "-ta" suffix) — matches "the cat sees the dog", not reversed. FINAL ANSWERS: suno mota-ta kire --- end worked example. Now solve the actual task below using ONLY the data in its own CONTEXT and QUERY — never reuse the words or rules above. ---""" _ANSWER_FORMAT_RULES = """Answer formatting rules by task_type: - translation: give the full translated phrase or sentence for each numbered item, written as a natural fluent sentence exactly like the answer column in the given examples — never as a morpheme-by-morpheme gloss with parentheses like "you(shuddered)" or "we(spat(on him))". - fill_blanks: give only the missing word(s)/form for each numbered blank. - match_letters: give the matching letter or number for each item. - text_to_num: give only the numeral for each item. - num_to_text: give only the word(s) for each item. Output a line that says exactly: FINAL ANSWERS: followed by one answer per line, in the same order and count as the numbered items in the query, with no numbering, quotes, or extra commentary — just the bare answer text for each line.""" SYSTEM_FULL_COT = f"""You are an expert linguist solving International Linguistics Olympiad problems. You are given a self-contained set of data in a language you have never seen before, plus a query asking you to apply what that data teaches you. You must reason only from the data given — never from memorized knowledge of real-world languages. Work in stages, showing your work: 1. IDENTIFY PATTERNS: look for recurring forms, affixes, word order, sound correspondences, or structural regularities in the given data. 2. RULE TABLE: write your hypotheses as short "TRIGGER -> TRANSFORMATION" bullet lines (e.g. "subject noun -> bare noun, placed first"), not prose — this keeps rules mechanical and easy to re-check, rather than vague descriptions that drift away from the actual data. 3. TEST HYPOTHESES: check each rule-table line against every example in the given data. Discard or refine any rule that doesn't hold up on every example. 4. APPLY: use only the rules that survived testing to answer the query. 5. VERIFY: before writing FINAL ANSWERS, re-check every draft answer against two common mistakes: - LANGUAGE: look at the answer column in the given examples, not the instructions, to see what language/form your answer must be in. If the examples' answers are word-forms in the unfamiliar language, your answer must also be a word-form in that language — never substitute an English gloss or description of the meaning, even if you're unsure of the exact form; give your best-guess constructed form instead. - ROLES: for anything involving "who did what to whom" (subject, object, possessor, giver/receiver), re-read the query and re-derive which argument fills which role from scratch. Do not assume the same word order or role assignment as a similar-looking training example — verify it against the actual affixes/markers in that example. {_WORKED_EXAMPLE} {_ANSWER_FORMAT_RULES}""" SYSTEM_PLAIN_IO = f"""You are an expert linguist solving International Linguistics Olympiad problems. You are given a self-contained set of data in a language you have never seen before, plus a query asking you to apply what that data teaches you. Reason only from the data given — never from memorized knowledge of real-world languages. Give your best-guess answer in the same language/form as the examples' answers — never substitute an English gloss. Double-check which argument is subject vs object before answering. {_WORKED_EXAMPLE} {_ANSWER_FORMAT_RULES}""" _SYSTEM_VARIANTS = {"full_cot": SYSTEM_FULL_COT, "plain_io": SYSTEM_PLAIN_IO} SYSTEM = _SYSTEM_VARIANTS[PROMPT_VARIANT] USER_TEMPLATE = """Task type: {task_type} Eval type: {eval_type} Working language -> Task language: {work_lang} -> {task_lang} CONTEXT: {context} QUERY: {query}""" # --- Query parsing helpers --------------------------------------------- _RANGE_RE = re.compile(r"\((\d+)\s*[-–]\s*(\d+)\)") _LEADING_NUM_RE = re.compile(r"^\s*(\d+)[.)]", re.MULTILINE) _INLINE_PAREN_NUM_RE = re.compile(r"\((\d+)\)") def _is_contiguous(nums: list[int]) -> bool: if not nums: return False uniq = sorted(set(nums)) return uniq == list(range(uniq[0], uniq[-1] + 1)) def count_expected_items(query: str) -> int: """Count how many answers the query expects, robust to numbering style.""" range_match = _RANGE_RE.search(query) if range_match: start, end = int(range_match.group(1)), int(range_match.group(2)) if end >= start: return end - start + 1 leading_nums = [int(n) for n in _LEADING_NUM_RE.findall(query)] if _is_contiguous(leading_nums): return len(set(leading_nums)) inline_nums = [int(n) for n in _INLINE_PAREN_NUM_RE.findall(query)] if _is_contiguous(inline_nums): return len(set(inline_nums)) non_empty_lines = [ln for ln in query.splitlines() if ln.strip()] return max(len(non_empty_lines) - 1, 1) # --- Answer parsing helpers ---------------------------------------------- _MARKER_RE = re.compile(r"FINAL ANSWERS:\s*", re.IGNORECASE) _LINE_NUM_PREFIX_RE = re.compile(r"^\s*\(?\d+\)?[.):]?\s*") _QUOTE_STRIP_RE = re.compile(r'^["\'](.*)["\']$') def _strip_line(line: str) -> str: line = line.strip() line = _LINE_NUM_PREFIX_RE.sub("", line, count=1) m = _QUOTE_STRIP_RE.match(line) if m: line = m.group(1) return line.strip() def parse_answers(text: str, expected_n: int) -> list[str]: """Extract exactly expected_n answer strings from raw model output.""" parts = _MARKER_RE.split(text, maxsplit=1) tail = parts[1] if len(parts) > 1 else text lines = [_strip_line(ln) for ln in tail.splitlines()] answers = [ln for ln in lines if ln] # Model sometimes emits all answers on one "|"-delimited line instead of # one per line; only fall back to splitting on "|" when the plain # line-split came up short, so genuine single-line answers containing "|" # aren't mangled. if len(answers) < expected_n and any("|" in a for a in answers): expanded = [] for a in answers: if "|" in a: expanded.extend(_strip_line(p) for p in a.split("|")) else: expanded.append(a) expanded = [a for a in expanded if a] if len(expanded) > len(answers): answers = expanded if len(answers) < expected_n: answers = answers + [""] * (expected_n - len(answers)) elif len(answers) > expected_n: answers = answers[:expected_n] return answers # --- Time-budget-aware batched generation -------------------------------- def elapsed() -> float: return time.monotonic() - T_START def remaining_budget() -> float: return TIME_BUDGET_SEC - elapsed() def build_prompt(tokenizer, row) -> str: user_msg = USER_TEMPLATE.format( task_type=row["task_type"], eval_type=row["eval_type"], work_lang=row.get("work_lang", ""), task_lang=row.get("task_lang", ""), context=row["context"], query=row["query"], ) messages = [ {"role": "system", "content": SYSTEM}, {"role": "user", "content": user_msg}, ] return tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) def adaptive_max_new_tokens(rows_left: int) -> int: if rows_left <= 0: return MIN_NEW_TOKENS per_row_budget = remaining_budget() / rows_left # Rough T4/AWQ-7B throughput assumption used only to size generation length; # errs conservative so we degrade gracefully rather than time out. tokens_per_sec_estimate = 20.0 budget_tokens = int(per_row_budget * tokens_per_sec_estimate) return max(MIN_NEW_TOKENS, min(MAX_NEW_TOKENS, budget_tokens)) def adaptive_k(rows_left: int, max_new_tokens: int) -> int: """How many sampled completions per row the remaining budget affords. Only returns >1 once max_new_tokens has already saturated at MAX_NEW_TOKENS (i.e. there's more budget than a single generation pass needs) — never trades away reasoning length for extra samples. """ if rows_left <= 0 or max_new_tokens <= 0: return MIN_K per_row_budget = remaining_budget() / rows_left tokens_per_sec_estimate = 20.0 budget_tokens = int(per_row_budget * tokens_per_sec_estimate) k = budget_tokens // max_new_tokens return max(MIN_K, min(MAX_K, k)) def majority_vote_answers(sampled_answer_lists: list[list[str]]) -> list[str]: """Collapse k parsed answer lists (same length, one per sample) into one by majority vote per position. Prefers non-empty answers over blanks when both appear, since a wrong guess still earns partial chrF credit and a blank never does. Ties break toward the first sample's value. """ if not sampled_answer_lists: return [] n = len(sampled_answer_lists[0]) result = [] for i in range(n): votes = [sample[i] for sample in sampled_answer_lists if i < len(sample)] pool = [v for v in votes if v] or votes if not pool: result.append("") continue counts = Counter(pool) max_count = max(counts.values()) winner = next(v for v in pool if counts[v] == max_count) result.append(winner) return result def pad_row(expected_n: int) -> str: return json.dumps([""] * expected_n) def main() -> None: df = pd.read_csv(DATA_PATH, dtype=str).fillna("") tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "left" model = AutoModelForCausalLM.from_pretrained( MODEL_PATH, torch_dtype=torch.float16, device_map="auto" ) model.eval() results: dict[str, str] = {} rows = df.to_dict("records") i = 0 while i < len(rows): rows_left = len(rows) - i if remaining_budget() < 30: # not enough time left to safely attempt a batch for row in rows[i:]: expected_n = count_expected_items(row["query"]) results[row["id"]] = pad_row(expected_n) break batch = rows[i : i + BATCH_SIZE] expected_counts = [count_expected_items(r["query"]) for r in batch] prompts = [build_prompt(tokenizer, r) for r in batch] max_new_tokens = adaptive_max_new_tokens(rows_left) k = adaptive_k(rows_left, max_new_tokens) inputs = tokenizer( prompts, return_tensors="pt", padding=True, truncation=True ).to(model.device) sampled_decoded: list[list[str]] = [] for sample_idx in range(k): gen_kwargs = dict( max_new_tokens=max_new_tokens, pad_token_id=tokenizer.pad_token_id ) if k > 1: gen_kwargs.update(do_sample=True, temperature=SAMPLE_TEMPERATURE, top_p=0.9) else: gen_kwargs.update(do_sample=False) try: with torch.no_grad(): output_ids = model.generate(**inputs, **gen_kwargs) input_len = inputs["input_ids"].shape[1] decoded = tokenizer.batch_decode( output_ids[:, input_len:], skip_special_tokens=True ) except Exception: decoded = [""] * len(batch) sampled_decoded.append(decoded) for row_idx, (row, expected_n) in enumerate(zip(batch, expected_counts)): sampled_answers = [ parse_answers(sampled_decoded[s][row_idx], expected_n) for s in range(k) ] answers = majority_vote_answers(sampled_answers) results[row["id"]] = json.dumps(answers) i += len(batch) out_df = pd.DataFrame( {"id": df["id"], "pred": [results[i] for i in df["id"]]} ) out_df.to_csv(OUTPUT_PATH, index=False) if __name__ == "__main__": main()