#!/usr/bin/env python3 """IOL-AI 2026 submission: Qwen2.5-14B-Instruct (bnb-4bit via unsloth/Qwen2.5-14B-Instruct-bnb-4bit), offline-only dependency install, a decomposition-and-verification prompt augmented with a deterministic symbolic-evidence layer, item-count-aware token budgeting, a fail-open closed-answer-space constraint for match_letters, and guaranteed explanations. History, briefly: every piece below was individually diagnosed against a real failure on real Linguini/IOL problems (a markdown-formatted answer marker, a COMPUTE-line bleeding into the answer list, an "is:" prefix surviving into a near-miss answer, a match_letters bijection violation, truncation on multi-item problems) before being combined here. Nothing in this file is speculative -- every module states the specific failure it closes. Compliance: fully offline before any Hugging Face import, MODEL_ID=".", reads only /tmp/data/test.csv, writes only submission.csv with id/pred/explanation, float16 (the T4 is Turing, no native bfloat16), the 30-minute budget is respected with a real safety margin, every row is guaranteed a submission.csv entry even under a crash or a timeout. """ from __future__ import annotations import atexit import os import time from pathlib import Path SCRIPT_STARTED_AT = time.monotonic() # --------------------------------------------------------------------------- # Offline mode, set before any Hugging Face import. Restored on exit (see # _restore_offline_env_vars below) -- if the evaluation harness ever runs # this script in-process rather than as an isolated subprocess, leftover # offline-mode env vars could otherwise affect a later, unrelated # huggingface_hub call made by the harness itself after this script exits. # --------------------------------------------------------------------------- _ORIGINAL_HF_HUB_OFFLINE = os.environ.get("HF_HUB_OFFLINE") _ORIGINAL_TRANSFORMERS_OFFLINE = os.environ.get("TRANSFORMERS_OFFLINE") def _restore_offline_env_vars() -> None: """Restores HF_HUB_OFFLINE/TRANSFORMERS_OFFLINE to their exact pre-script state on exit, via atexit so it fires regardless of how or where the script exits. Costs nothing; cannot make anything worse.""" for key, original in (("HF_HUB_OFFLINE", _ORIGINAL_HF_HUB_OFFLINE), ("TRANSFORMERS_OFFLINE", _ORIGINAL_TRANSFORMERS_OFFLINE)): if original is None: os.environ.pop(key, None) else: os.environ[key] = original atexit.register(_restore_offline_env_vars) os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" import subprocess import sys # --------------------------------------------------------------------------- # Configuration -- env-var overridable, sensible defaults otherwise. # --------------------------------------------------------------------------- SCRIPT_DIR = Path(__file__).resolve().parent INPUT_CSV = Path(os.environ.get("IOL_INPUT", "/tmp/data/test.csv")) OUTPUT_CSV = Path(os.environ.get("IOL_OUTPUT", "submission.csv")) MODEL_ID = os.environ.get("IOL_MODEL_ID", ".") TIME_LIMIT_S = float(os.environ.get("IOL_TIME_LIMIT_S", 30 * 60)) SETUP_BUFFER_S = float(os.environ.get("IOL_SETUP_BUFFER_S", 420)) # 14B bnb-4bit is ~8-9GB, slow to load EXIT_RESERVE_S = float(os.environ.get("IOL_EXIT_RESERVE_S", 60)) TOKENS_CAP_FLOOR = int(os.environ.get("IOL_TOKENS_FLOOR", 640)) TOKENS_CAP_CEIL = int(os.environ.get("IOL_TOKENS_CEIL", 1536)) TOKENS_PER_ITEM = int(os.environ.get("IOL_TOKENS_PER_ITEM", 48)) TOKENS_PER_ITEM_BASE = int(os.environ.get("IOL_TOKENS_PER_ITEM_BASE", 256)) def elapsed_seconds() -> float: return time.monotonic() - SCRIPT_STARTED_AT # --------------------------------------------------------------------------- # Crash safety: two independent write paths, neither depending on the # other, neither depending on anything that might have just failed. # --------------------------------------------------------------------------- def write_submission_csv(rows_list: list[dict]) -> None: """Stdlib csv, not pandas -- avoids a real, documented pandas/numpy ABI crash ('TypeError: Cannot convert numpy.ndarray to numpy.ndarray' inside pandas' Index construction) that a top-scoring public IOL-AI 2026 submission hit in this exact sandbox. Atomic: writes to a temp file then os.replace()s it into place, so a reader can never observe a partially-written file mid-save.""" import csv tmp_path = OUTPUT_CSV.with_suffix(OUTPUT_CSV.suffix + ".tmp") with tmp_path.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=["id", "pred", "explanation"]) writer.writeheader() for row in rows_list: writer.writerow(row) os.replace(tmp_path, OUTPUT_CSV) def emergency_submission_csv(reason: str, rows_so_far: list[dict] | None = None) -> None: """Last-resort guarantee: no matter WHERE the script dies, a valid submission.csv exists before the process exits -- the single fix for the pattern where a crash with nothing written turns a scoreable zero into a hard evaluation failure. Independent of write_submission_csv: uses only the standard library, so it cannot fail for the same reason a pandas-based path might.""" import csv import json try: if rows_so_far: write_submission_csv(rows_so_far) return ids: list[str] = [] try: with INPUT_CSV.open(newline="", encoding="utf-8") as f: for row in csv.DictReader(f): if row.get("id"): ids.append(row["id"]) except Exception: pass rows = [{"id": i, "pred": json.dumps([""]), "explanation": f"EMERGENCY FALLBACK: {str(reason)[:150]}"} for i in ids] write_submission_csv(rows) except Exception: try: with OUTPUT_CSV.open("w") as f: f.write("id,pred,explanation\n") except Exception: pass # --------------------------------------------------------------------------- # Offline dependency install. # --------------------------------------------------------------------------- def ensure_dependencies() -> None: """Split deliberately: torch is NOT force-upgraded (a multi-GB CUDA-specific wheel; forcing -U risks pulling a build mismatched with the sandbox's actual driver -- a worse failure than a missing package). bitsandbytes needs no upgrade evidence behind it. transformers/accelerate/tokenizers have a CONFIRMED version-related failure behind them -- those are the only ones forced.""" subprocess.run([sys.executable, "-m", "pip", "install", "-q", "torch>=2.2", "bitsandbytes", "pandas"], check=True) subprocess.run([sys.executable, "-m", "pip", "install", "-q", "-U", "transformers>=4.43", "accelerate>=0.30", "tokenizers"], check=True) try: ensure_dependencies() except Exception as exc: emergency_submission_csv(f"pip install failed: {exc}") raise import re import json import unicodedata import ast as pyast import pandas as pd import torch from difflib import SequenceMatcher from collections import defaultdict from transformers import AutoTokenizer, AutoModelForCausalLM # --------------------------------------------------------------------------- # Model loading. # --------------------------------------------------------------------------- def load_model(): """Fast tokenizer first; on failure, falls back to use_fast=False -- bypasses TokenizerFast.from_file() entirely, which is exactly the call that fails on a tokenizer.json saved by a newer tokenizers library than the sandbox has.""" try: tok = AutoTokenizer.from_pretrained(MODEL_ID) print("Tokenizer loaded (fast).", flush=True) except Exception as exc: print(f"Fast tokenizer failed ({exc}); falling back to use_fast=False.", flush=True) tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False) print("Tokenizer loaded (slow fallback).", flush=True) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=torch.float16, device_map="auto", ).eval() print(f"Model loaded | memory footprint: {round(model.get_memory_footprint() / 1e9, 1)} GB | " f"quantized: {getattr(model.config, 'quantization_config', None) is not None}", flush=True) return tok, model # --------------------------------------------------------------------------- # Query parsing: widened patterns + honest "unknown count" fallback. # --------------------------------------------------------------------------- def parse_items(query: str) -> tuple[str, list[str], bool]: """Returns (preamble, items, count_known). count_known=False means no pattern matched -- we do NOT guess a count, we let the model's own answer list stand rather than risk truncating real content.""" item_pat = re.compile(r"(?m)^\s*(\d+)\s*[.\)]\s*(.*)$") matches = list(item_pat.finditer(query)) if matches: preamble = query[:matches[0].start()].strip() items = [] for i, m in enumerate(matches): end = matches[i + 1].start() if i + 1 < len(matches) else len(query) text = re.sub(r"^\s*\d+\s*[.\)]\s*", "", query[m.start():end].strip()) items.append(text) return preamble, items, True rng = re.search(r"[\(\[]?\s*(\d+)\s*(?:[-\u2013\u2014:]|to)\s*(\d+)\s*[\)\]]?", query, flags=re.IGNORECASE) if rng: lo, hi = int(rng.group(1)), int(rng.group(2)) if 0 < hi - lo < 100: items = [] for k in range(lo, hi + 1): line_match = re.search(rf"(?m)^.*\(\s*{k}\s*\).*$", query) if line_match: clue = re.sub(rf"\(\s*{k}\s*\)", "", line_match.group(0)).strip() clue = re.sub(r"\|\s*\|", "|", clue) clue = re.sub(r"\s{2,}", " ", clue).strip(" |") items.append(clue if clue else f"the numbered item {k} from the examples above") else: items.append(f"the numbered item {k} from the examples above") return query.strip(), items, True csv_nums = re.findall(r"(?m)^\s*(\d+)\s*,\s*(\d+(?:\s*,\s*\d+)*)\s*$", query) if csv_nums: all_nums = re.findall(r"\d+", " ".join(csv_nums[0])) return query.strip(), [f"the numbered item {n}" for n in all_nums], True return query.strip(), [], False TASK_GUIDANCE = { "translation": "give the translated form only, in the language asked.", "fill_blanks": "give only the missing form for each blank.", "match_letters": "give only the option letter (for example A, B, C).", "text_to_num": "give the number in digits.", "num_to_text": "give the number written out in words, in the language asked.", } DEFAULT_GUIDANCE = "give exactly what the instruction asks, nothing else." # --------------------------------------------------------------------------- # Symbolic preprocessing layer -- pure standard library, no new # dependencies, deterministic, CPU-only, negligible runtime. Survived a # multi-round falsification pass: only the two evidence objects that (a) # compute something a fast read is likely to miss by construction and (b) # cannot mislead when wrong (worst case is silence, never false # confidence) were kept. Augments the raw context; never replaces it. # --------------------------------------------------------------------------- def extract_forms_from_context(context: str) -> list[str]: """Pulls candidate unknown-language 'forms' for reduplication's per-word self-check ONLY. Pipe-delimited lines contribute ONLY their FIRST field -- including gloss/meaning fields would let ordinary English words trigger false reduplication hits. Lines with more than 3 pipes are skipped defensively -- Hadza (a confirmed IOL 2026 language) is a click language, and '|' is sometimes used informally to transcribe click consonants, which would misparse as our delimiter.""" forms = [] for line in context.splitlines(): line = line.strip() if not line: continue pipe_count = line.count("|") if 0 < pipe_count <= 3: first_field = re.sub(r"^\s*\d+\s*[.\)]\s*", "", line.split("|")[0].strip()).strip() if first_field: forms.append(first_field) elif pipe_count == 0: for t in line.split(): t_clean = re.sub(r"^\s*\d+\s*[.\)]\s*", "", t).strip(".,;:") if t_clean and len(t_clean) > 1: forms.append(t_clean) seen, unique_forms = set(), [] for f in forms: if f not in seen: seen.add(f) unique_forms.append(f) return unique_forms def extract_explicit_pairs(context: str) -> list[tuple[str, str]]: """Genuine (input, output) pairs from pipe-delimited rows -- e.g. fill_blanks' 'given | derived | gloss' structure. The ONLY source of pairs fed to transformation-family detection: forms from DIFFERENT rows are never cross-compared, which would manufacture spurious 'transformations' between unrelated words.""" pairs = [] for line in context.splitlines(): line = line.strip() if not (0 < line.count("|") <= 3): continue fields = [re.sub(r"^\s*\d+\s*[.\)]\s*", "", f.strip()).strip() for f in line.split("|")] fields = [f for f in fields if f] if len(fields) >= 2: pairs.append((fields[0], fields[1])) return pairs def edit_signature(a: str, b: str): """A clean single-region transformation signature, or None if the difference is scattered (too noisy to call one transformation), OR if there is no genuine shared stem of at least 2 characters -- without this check, two totally unrelated words with zero characters in common were being accepted as a fake signature, since SequenceMatcher returns a single 'replace' opcode for a total mismatch too.""" sm = SequenceMatcher(None, a, b, autojunk=False) all_ops = sm.get_opcodes() ops = [op for op in all_ops if op[0] != "equal"] if not ops or len(ops) > 2: return None equal_len = sum((i2 - i1) for tag, i1, i2, j1, j2 in all_ops if tag == "equal") if equal_len < 2: return None tag, i1, i2, j1, j2 = ops[0] removed, inserted = a[i1:i2], b[j1:j2] if i1 == 0: pos = "prefix" elif i2 == len(a): pos = "suffix" else: pos = "infix" return (pos, removed, inserted) def find_transformation_families(pairs: list[tuple[str, str]]) -> list[str]: """Clusters GENUINELY PAIRED forms (same row only) sharing an identical clean edit signature. Emits a family only if 2+ separate given pairs share it -- one occurrence is worse than silence.""" groups = defaultdict(list) for a, b in pairs: if not a or not b or a == b: continue sig = edit_signature(a, b) if sig: groups[sig].append((a, b)) families = [] for sig, grp in groups.items(): unique_pairs = list(dict.fromkeys(grp)) if len(unique_pairs) >= 2: pos, removed, inserted = sig removed_disp = removed if removed else "(nothing)" inserted_disp = inserted if inserted else "(nothing)" examples = "; ".join(f"{a}->{b}" for a, b in unique_pairs[:4]) families.append((len(unique_pairs), f"{pos} change: '{removed_disp}' -> '{inserted_disp}' (seen in: {examples})")) families.sort(key=lambda x: -x[0]) return [f for _, f in families] def detect_reduplication(forms: list[str]) -> list[str]: """Flags a word only if it contains an exact adjacent doubled substring (length >= 2). Emits nothing if absent.""" findings = [] for w in forms: n = len(w) found = False for length in range(2, n // 2 + 1): for start in range(0, n - 2 * length + 1): chunk = w[start:start + length] nxt = w[start + length:start + 2 * length] if chunk == nxt: findings.append(f"reduplication in '{w}': '{chunk}' repeated") found = True break if found: break return findings def build_symbolic_evidence(context: str) -> str: """Returns "" if no supported transformation family and no reduplication is found -- augments the prompt only with real, multi-supported evidence. Never replaces context.""" forms = extract_forms_from_context(context) pairs = extract_explicit_pairs(context) families = find_transformation_families(pairs) if pairs else [] redup = detect_reduplication(forms) if forms else [] lines = [] if families: lines.append("Transformation families found (patterns supported by multiple examples):") for f in families[:3]: lines.append(f"- {f}") if redup: lines.append("Reduplication detected:") for r in redup[:2]: lines.append(f"- {r}") if not lines: return "" return ("\n\nSYMBOLIC EVIDENCE (deterministically computed from the examples above; " "may be incomplete -- verify against the examples, do not trust blindly):\n" + "\n".join(lines)) # --------------------------------------------------------------------------- # Closed-answer-space pre-constraint, match_letters only. Deterministic, # read-only, zero generate() calls of its own. Extracts the closed set of # option letters genuinely present in the context (always explicitly # given -- "A. water", "B. child", ...) and states that exact set as a # soft hint in the prompt. FAIL-OPEN: if extraction isn't clean and # unambiguous, the hint is skipped -- baseline behavior for that row is # byte-identical to not having this module at all. # --------------------------------------------------------------------------- def extract_match_letter_options(context: str) -> list[str] | None: """Returns a sorted list of option letters if extraction is CLEAN and UNAMBIGUOUS, else None. Deliberately strict: must never guess.""" found = set() for line in context.splitlines(): for m in re.finditer(r"(?:^|\s)([A-Z])[.\)]\s+\S", line): found.add(m.group(1)) if not found: return None letters = sorted(found) expected = [chr(ord("A") + i) for i in range(len(letters))] if letters != expected: return None if not (2 <= len(letters) <= 26): return None return letters def build_messages(context: str, query: str, task_type: str) -> tuple[list[dict], int | None]: """The proven decomposition-and-verification scaffold, augmented with the symbolic evidence layer and the match_letters closed-option hint. Nothing else about the reasoning instructions has changed since the version that scored 0.083/0.0296/0.2323 on the real leaderboard.""" preamble, items, count_known = parse_items(query) guidance = TASK_GUIDANCE.get(task_type, DEFAULT_GUIDANCE) symbolic_evidence = build_symbolic_evidence(context) system = ( "You solve puzzles about a language you have never seen. Everything you " "need is in the examples below. Use only the examples, not outside " "knowledge of any language. You may meet a task type you have never " "seen -- read the instruction and examples, and answer in the same " "form they use." ) number_note = "" if task_type == "text_to_num": number_note = ( "\n\nAlso add one more line after your answers, exactly like this:\n" "COMPUTE: expr1 | expr2\n" "where each expr is a plain arithmetic expression (digits, +, -, *, " "parentheses only) for that item's value, one per answer, matching " "the rule you found." ) options_note = "" if task_type == "match_letters": options = extract_match_letter_options(context) if options: options_note = ( f"\n\nThe only valid answers are: {', '.join(options)}. " f"Do not use any other letter." ) if count_known: n_items = len(items) slots = "\n\n".join(f"Question {i + 1}: {it}\nAnswer {i + 1}:" for i, it in enumerate(items)) user = ( f"EXAMPLES:\n{context.strip()}" f"{symbolic_evidence}\n\n" f"--- The examples end here. The questions begin below. ---\n\n" f"For each question: find the rule that explains ALL the examples above " f"(not just one). Check it against every example before answering. " f"For this task type, {guidance}\n\n" f"{preamble}\n\n{slots}\n\n" f"After answering all {n_items} questions, finish with exactly one line, " f"all {n_items} answers in order separated by ' | ':\n" f"FINAL ANSWERS: answer1 | answer2" f"{number_note}" f"{options_note}" ) else: n_items = None user = ( f"EXAMPLES:\n{context.strip()}" f"{symbolic_evidence}\n\n" f"--- The examples end here. The question begins below. ---\n\n" f"Find the rule that explains ALL the examples above (not just one). " f"Check it against every example before answering. " f"For this task type, {guidance}\n\n" f"{preamble}\n\n" f"Answer every item asked above, in order, one per answer. Finish " f"with exactly one line, all your answers in order separated by ' | ':\n" f"FINAL ANSWERS: answer1 | answer2" f"{number_note}" f"{options_note}" ) return [{"role": "system", "content": system}, {"role": "user", "content": user}], n_items def build_repair_messages(query: str, n_items: int | None, bad_text: str) -> list[dict]: n_desc = f"exactly {n_items}" if n_items is not None else "one per item asked" system = "You reformat answers. Output nothing except the requested line." user = ( f"Question:\n{query.strip()}\n\n" f"A previous attempt produced:\n{bad_text[:600]}\n\n" f"Extract or restate {n_desc} final answers, in order, as ONE line:\n" f"FINAL ANSWERS: answer1 | answer2" ) return [{"role": "system", "content": system}, {"role": "user", "content": user}] # --------------------------------------------------------------------------- # Safe arithmetic: no exec(), no eval() of arbitrary code. # --------------------------------------------------------------------------- _ALLOWED_BINOPS = (pyast.Add, pyast.Sub, pyast.Mult) def safe_arithmetic(expr: str) -> float | int | None: try: tree = pyast.parse(expr.strip(), mode="eval") except Exception: return None def _eval(node): if isinstance(node, pyast.Expression): return _eval(node.body) if isinstance(node, pyast.Constant) and isinstance(node.value, (int, float)): return node.value if isinstance(node, pyast.BinOp) and isinstance(node.op, _ALLOWED_BINOPS): left, right = _eval(node.left), _eval(node.right) if left is None or right is None: return None if isinstance(node.op, pyast.Add): return left + right if isinstance(node.op, pyast.Sub): return left - right if isinstance(node.op, pyast.Mult): return left * right if isinstance(node, pyast.UnaryOp) and isinstance(node.op, pyast.USub): v = _eval(node.operand) return -v if v is not None else None return None return _eval(tree) def clean_answer(a: str) -> str: """Broadened: strips "Answer N:", "is:", "the answer is:", "final answer:" prefixes, applies NFC Unicode normalization and collapses internal whitespace runs -- both target exact-match killers the organizers' own documented normalization does not cover (Unicode form, internal whitespace).""" a = re.sub(r"(?i)^\s*(the\s+)?(final\s+)?answer\s*\d*\s*(is)?\s*:\s*", "", a).strip() a = re.sub(r"(?i)^\s*is\s*:\s*", "", a).strip() a = a.strip("* ") a = unicodedata.normalize("NFC", a) a = re.sub(r"\s{2,}", " ", a) return a.strip(" .\"'\u201c\u201d\u2018\u2019") def extract(text: str) -> tuple[list[str], int | None]: """Fixed against three real bugs found on real Linguini output: (1) markdown-bold marker with content on the NEXT line, not same line; (2) a following COMPUTE: line bleeding into the answer list; (3) NO marker found + answers dumped on one pipe-separated line -- splits each fallback line further by "|" instead of treating the whole line as one answer.""" m = list(re.finditer(r"final answers?\s*:?\s*\**", text, flags=re.IGNORECASE)) if m: tail = text[m[-1].end():] stop = re.search(r"(?i)compute\s*:", tail) if stop: tail = tail[:stop.start()] tail = tail.replace("**", " ").strip() candidate = " ".join(tail.splitlines()) parts = [clean_answer(p) for p in candidate.split("|") if p.strip()] if parts: return parts, m[-1].start() lines = [ln.strip() for ln in text.splitlines() if ln.strip()] fallback = [] for ln in lines: ln_clean = re.sub(r"^\s*\d+\s*[.\)]\s*", "", ln) if "|" in ln_clean: fallback.extend(clean_answer(p) for p in ln_clean.split("|") if p.strip()) else: fallback.append(clean_answer(ln_clean)) return fallback, None def extract_compute_overrides(text: str, n_answers: int) -> dict[int, str]: m = re.search(r"compute\s*:\s*(.+)", text, flags=re.IGNORECASE) if not m: return {} exprs = [e.strip() for e in m.group(1).split("|")] overrides = {} for i, e in enumerate(exprs[:n_answers]): val = safe_arithmetic(e) if val is not None: overrides[i] = str(int(val)) if float(val).is_integer() else str(val) return overrides # --------------------------------------------------------------------------- # Generation: defensive against both chat-template return shapes (the # sandbox's transformers version may return a bare tensor from # apply_chat_template rather than a dict), and against a constrained # decoding attempt failing for any reason. # --------------------------------------------------------------------------- def generate(tok, model, messages: list[dict], max_new_tokens: int, constraint_fn=None) -> str: """constraint_fn: optional prefix_allowed_tokens_fn, default None means byte-identical behavior to an unconstrained call. If a constrained attempt fails for ANY reason, falls back to a fully UNCONSTRAINED generation (not a retry with the same broken kwarg) -- the two concerns (dict-vs-tensor API shape, constrained-vs-unconstrained) are isolated from each other so a failure in one never masks as the other.""" def _try_generate(gen_kwargs): try: enc = tok.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", return_dict=True, ).to(model.device) input_len = enc["input_ids"].shape[-1] with torch.no_grad(): out = model.generate(**enc, **gen_kwargs) except Exception: ids = tok.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", ).to(model.device) input_len = ids.shape[-1] with torch.no_grad(): out = model.generate(ids, **gen_kwargs) return out, input_len base_kwargs = {"max_new_tokens": max_new_tokens, "do_sample": False} if constraint_fn is not None: try: out, input_len = _try_generate({**base_kwargs, "prefix_allowed_tokens_fn": constraint_fn}) except Exception: out, input_len = _try_generate(base_kwargs) else: out, input_len = _try_generate(base_kwargs) return tok.decode(out[0][input_len:], skip_special_tokens=True).strip() # Adapted from a top-1 public submission's prefix_allowed_tokens_fn # technique -- but scoped correctly for OUR prompt architecture. Their # version applies across an ENTIRE generation because their prompt has no # reasoning phase (plain newline-per-answer output). Ours does have a # reasoning phase (decomposition + "FINAL ANSWERS:" marker); applying a # letter-only constraint there would silently break the model's ability to # reason at all. Scoped here to ONLY the repair call, whose entire # expected output is already a short answer line. Fail-open throughout. _LETTER_CONSTRAINT_CACHE: dict = {} def build_letter_constraint_fn(tok, valid_letters: list[str]): cache_key = (id(tok), tuple(sorted(valid_letters))) if cache_key in _LETTER_CONSTRAINT_CACHE: return _LETTER_CONSTRAINT_CACHE[cache_key] try: allowed_chars = set(valid_letters) | set(" |\n\t\r") eos = tok.eos_token_id pieces = [] for token_id in range(len(tok)): if token_id == eos: continue piece = tok.decode([token_id], skip_special_tokens=False) if piece and all(c in allowed_chars for c in piece): pieces.append(token_id) allowed_ids = ([eos] if eos is not None else []) + pieces def allowed(_batch_id, _input_ids): return allowed_ids if allowed_ids else list(range(len(tok))) _LETTER_CONSTRAINT_CACHE[cache_key] = allowed return allowed except Exception: return None EXPLANATION_SYSTEM = ( "Summarize the following reasoning into a few short bullet points: the " "rule or pattern found in the data and the key evidence for the answer. " "Be concise and structured -- do not repeat the full reasoning." ) EXPLANATION_FALLBACK = "Answer derived from patterns found in the examples above." def dynamic_tokens_cap(n_items: int | None, time_based_cap: int) -> int: """Item-count-aware token budget, evidenced by a top-1 public submission citing truncation on multi-item problems as "a pure unforced loss". Combined with, not replacing, the time-based adaptation via min() -- a multi-item problem gets more room, but never more than time allows.""" if not n_items: return time_based_cap item_based_cap = max(TOKENS_CAP_FLOOR, min(TOKENS_CAP_CEIL, n_items * TOKENS_PER_ITEM + TOKENS_PER_ITEM_BASE)) return min(time_based_cap, item_based_cap) def process_row(tok, model, row: dict, n_rows: int, n_done: int, per_row_budget: float) -> tuple[dict, bool]: """Processes one row. Returns (result_row, ok) -- ok=False means a fallback row was produced after an exception, not a real answer.""" try: remaining = TIME_LIMIT_S - elapsed_seconds() budget_left_rows = max(n_rows - n_done, 1) row_budget = remaining / budget_left_rows time_based_cap = 1280 if row_budget > per_row_budget else 640 task_type = row.get("task_type", "") messages, n_items = build_messages(row["context"], row["query"], task_type) tokens_cap = dynamic_tokens_cap(n_items, time_based_cap) text = generate(tok, model, messages, tokens_cap) answers, marker_pos = extract(text) if task_type == "text_to_num": overrides = extract_compute_overrides(text, len(answers)) for idx, val in overrides.items(): if idx < len(answers): answers[idx] = val # Repair only on TRUE extraction failure (no marker / nothing found) # -- not a mere count difference, since extra answers are harmless # and our own count guess may be the thing that's wrong. if (marker_pos is None or not answers) and remaining > SETUP_BUFFER_S: repair_constraint = None if task_type == "match_letters": repair_options = extract_match_letter_options(row["context"]) if repair_options: repair_constraint = build_letter_constraint_fn(tok, repair_options) repair_text = generate(tok, model, build_repair_messages(row["query"], n_items, text), 128, constraint_fn=repair_constraint) rep, rep_pos = extract(repair_text) if rep: answers, marker_pos = rep, rep_pos if n_items is not None: if len(answers) < n_items: answers = answers + [answers[-1] if answers else ""] * (n_items - len(answers)) elif len(answers) > n_items and marker_pos is None: answers = answers[:n_items] # else: marker found, more answers than our guess -> keep them all if not answers: answers = [""] # Explanation: dedicated call if time is comfortable, else a cheap # truncated fallback -- never blank, never a second full generation # under time pressure. remaining_after = TIME_LIMIT_S - elapsed_seconds() budget_left_after = max(n_rows - n_done - 1, 0) comfortable = remaining_after > (budget_left_after + 1) * per_row_budget * 1.3 if comfortable: try: explanation = generate( tok, model, [{"role": "system", "content": EXPLANATION_SYSTEM}, {"role": "user", "content": text}], 300, ) or EXPLANATION_FALLBACK except Exception: explanation = EXPLANATION_FALLBACK else: snippet = re.sub(r"\s{2,}", " ", text[:300]).strip() explanation = snippet if snippet else EXPLANATION_FALLBACK return {"id": row["id"], "pred": json.dumps(answers, ensure_ascii=False), "explanation": explanation}, True except Exception as exc: try: _, fallback_items, fk = parse_items(row["query"]) n_fallback = len(fallback_items) if fk else 1 except Exception: n_fallback = 1 print(f"ROW ERROR on {row.get('id', '?')}: {exc}", flush=True) return {"id": row["id"], "pred": json.dumps([""] * n_fallback, ensure_ascii=False), "explanation": EXPLANATION_FALLBACK}, False def main() -> None: if not INPUT_CSV.exists(): emergency_submission_csv(f"input CSV not found: {INPUT_CSV}") raise FileNotFoundError(f"Missing input CSV: {INPUT_CSV}") try: # Read test.csv and checkpoint a placeholder submission FIRST, # before the slowest and most failure-prone step (model loading, # ~7-9 min for this checkpoint size) even starts -- matching a # top-1 public submission's proven order. If loading hangs or gets # killed, something valid already exists on disk. df = pd.read_csv(INPUT_CSV, dtype=str).fillna("") placeholder_rows = [{"id": rid, "pred": json.dumps([""]), "explanation": "Placeholder written before model load."} for rid in df["id"].tolist()] write_submission_csv(placeholder_rows) print(f"Pre-load checkpoint written for {len(placeholder_rows)} rows.", flush=True) tok, model = load_model() except Exception as exc: emergency_submission_csv(f"tokenizer/model load or test.csv read failed: {exc}") raise n_rows = len(df) actual_setup_elapsed = elapsed_seconds() per_row_budget = max(20, (TIME_LIMIT_S - actual_setup_elapsed) / max(n_rows, 1)) print(f"Setup took {actual_setup_elapsed:.0f}s (estimated {SETUP_BUFFER_S:.0f}s) | " f"per_row_budget={per_row_budget:.0f}s for {n_rows} rows", flush=True) rows: list[dict] = [] processed_ids: set[str] = set() try: for _, row in df.iterrows(): result_row, _ok = process_row(tok, model, row, n_rows, len(rows), per_row_budget) rows.append(result_row) processed_ids.add(row["id"]) write_submission_csv(rows) print(f"{len(rows)}/{n_rows} elapsed={elapsed_seconds():.0f}s", flush=True) if elapsed_seconds() > TIME_LIMIT_S - EXIT_RESERVE_S: print("Time budget nearly exhausted, stopping early.", flush=True) break # Guarantee one row per test.csv id, even under a timeout. for _, row in df.iterrows(): if row["id"] in processed_ids: continue try: _, fallback_items, fk = parse_items(row["query"]) n_fallback = len(fallback_items) if fk else 1 except Exception: n_fallback = 1 rows.append({"id": row["id"], "pred": json.dumps([""] * n_fallback, ensure_ascii=False), "explanation": EXPLANATION_FALLBACK}) write_submission_csv(rows) print(f"DONE. Wrote {len(rows)} rows in {elapsed_seconds():.0f}s.", flush=True) except Exception as exc: # Final safety net: even if something escapes every inner # try/except above, whatever rows were collected so far still get # written. emergency_submission_csv(f"main loop failed: {exc}", rows_so_far=rows if rows else None) print(f"FATAL, but submission.csv was written with {len(rows)} rows. Error: {exc}", flush=True) if __name__ == "__main__": main()