# ============================================================ # IOL-AI 2026: ALGEBRAIC INDUCTION SOLVER (SUBMISSION SCRIPT) # ============================================================ import os # The evaluation sandbox has no internet access. # These environment variables force Transformers to use local files only. os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" os.environ["TOKENIZERS_PARALLELISM"] = "false" import gc, re, time, torch, json import pandas as pd from collections import Counter from transformers import AutoTokenizer, AutoModelForCausalLM # ──────────────────────────────────────────────────── # 1. CONFIGURATION AND MODEL LOAD # ──────────────────────────────────────────────────── try: del model, tok except NameError: pass gc.collect() torch.cuda.empty_cache() MODEL_ID = "." # Load weights directly from the repository MAX_TOKEN_BUDGET = 2048 INDUCT_MAX_TOKENS = 800 MAX_ATTEMPTS = 3 GLOBAL_TIME_LIMIT = 1700 # 28.3 minutes (safe margin under 30 min limit) SC_TASKS = frozenset({"match_letters", "fill_blanks"}) SC_K = 3 print("Loading tokenizer and model...", flush=True) tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, dtype=torch.float16, device_map="auto", trust_remote_code=True ).eval() print("Model loaded successfully.", flush=True) tok.padding_side = "left" if tok.pad_token_id is None: tok.pad_token = tok.eos_token or tok.unk_token # ──────────────────────────────────────────────────── # 2. EOS DISCOVERY # ──────────────────────────────────────────────────── def get_eos_ids(tokenizer, model): eos_ids = set() if tokenizer.eos_token_id is not None: eos_ids.add(int(tokenizer.eos_token_id)) meos = getattr(model.generation_config, "eos_token_id", None) if meos: if isinstance(meos, (list, tuple, set)): eos_ids.update(int(x) for x in meos) else: eos_ids.add(int(meos)) return sorted(list(eos_ids)) EOS_IDS = get_eos_ids(tok, model) EOS_SET = set(EOS_IDS) model.generation_config.eos_token_id = EOS_IDS model.generation_config.pad_token_id = tok.pad_token_id # ──────────────────────────────────────────────────── # 3. DYNAMIC CONTEXT READERS # ──────────────────────────────────────────────────── _IPA_HINT = re.compile(r"[\u0250-\u02AF\u02B0-\u02FF\u0300-\u036F\u1D00-\u1D7Føœæðθŋɣʔ]") _ASKS_NON_PHONETIC = re.compile(r"(?i)translate\s+into\s+english|write\s+(it\s+)?in\s+the\s+[\w'\u2019-]+\s+orthography|in\s+the\s+regular\s+orthography") _ASKS_TRANSCRIPTION = re.compile(r"(?i)\b(transcribe|transcription|phonetic(ally)?)\b") def _bracketed_forms(text: str) -> list[str]: out = [] for m in re.finditer(r"\[([^\[\]\n]{1,40})\]", text): inner = m.group(1).strip() if not inner or re.fullmatch(r"[\d\s,.\-]+", inner): continue out.append(inner) return out def is_phonetic_task(context: str, query: str, min_forms: int = 3) -> bool: if _ASKS_NON_PHONETIC.search(query): return False if _bracketed_forms(query) and not _ASKS_TRANSCRIPTION.search(query): return False forms = _bracketed_forms(context) + _bracketed_forms(query) if len(forms) < min_forms: return False phonetic_looking = sum(1 for f in forms if _IPA_HINT.search(f) or ":" in f) return phonetic_looking >= max(2, len(forms) // 4) def count_items(query: str) -> int: n = len(re.findall(r"(?m)^\s*\d+[.)]", query)) if n: return n if "blanks" in query.lower(): m = re.search(r"\((\d+)-(\d+)\)", query) if m: return int(m.group(2)) - int(m.group(1)) + 1 return len(re.findall(r"\(\d+\)", query)) or 0 return 0 # ──────────────────────────────────────────────────── # 4. DEFENSIVE PARSING # ──────────────────────────────────────────────────── _TURN_NOISE = re.compile(r"<\|/?END_OF_TURN_TOKEN\|>|<\|/?START_OF_TURN_TOKEN\|>|<\|CHATBOT_TOKEN\|>|||<\|im_end\|>|<\|im_start\|>") _MARKER = re.compile(r"(?im)^\s*final answers?\s*:?\s*$") def _looks_like_prose(line: str) -> bool: if re.search(r"(?i)^(final answers?|answers?|note|reviewing|summary|explanation|verification)\b.*:$", line): return True if re.search(r"(?i)^(here (are|is)|the (final )?answers? (are|is)|based on|therefore|thus|in summary)\b", line): return True if line.rstrip().endswith(":") and len(line) > 3: return True if len(line) > 120: return True return False def _strip_gloss_keep_form(line: str) -> str: s = re.sub(r"\*\*", "", (line or "").strip()) s = re.split(r"\s+_?(?:GCY|NS|N/A)_?\b", s, maxsplit=1, flags=re.I)[0].strip() m = re.match(r"^(.+?)\s+[-–—]\s+((?:to|the|a|an|in|of|for|being|means?)\b.*)$", s, flags=re.I) if m: s = m.group(1).strip() return s.strip() def parse_answers(text: str, n_items: int = 0) -> list[str]: text = _TURN_NOISE.sub("", text or "") markers = list(_MARKER.finditer(text)) if markers: text = text[markers[-1].end():] else: m = re.search(r'(.*?)', text, re.DOTALL | re.IGNORECASE) if m: text = m.group(1) answers = [] for line in text.splitlines(): line = re.sub(r"^\s*\d+[.)]\s*", "", line).strip().strip("`").strip("*") if not line or _looks_like_prose(line): continue line = _strip_gloss_keep_form(line) if not line: continue if re.fullmatch(r"(?:[A-Za-z]\s+)+[A-Za-z]", line): answers.extend([p.upper() for p in line.split()]) continue answers.append(line) if n_items > 0: answers = answers[:n_items] if len(answers) < n_items: answers += [""] * (n_items - len(answers)) return answers def majority_vote(samples: list[list[str]], n_items: int) -> list[str]: usable = [s for s in samples if any(x.strip() for x in s)] if not usable: return [""] * max(n_items, 0) n = n_items or max(len(s) for s in usable) padded = [(list(s) + [""] * n)[:n] for s in usable] counts = Counter(tuple(p) for p in padded) best, c = counts.most_common(1)[0] if c >= 2: return list(best) return [Counter(p[i] for p in padded).most_common(1)[0][0] for i in range(n)] # ──────────────────────────────────────────────────── # 5. PROMPT BUILDERS (ALGEBRAIC INDUCTION) # ──────────────────────────────────────────────────── SYSTEM_BASE = ( "You are an elite computational linguist solving International Linguistics Olympiad problems. " "This is a closed-world puzzle. DO NOT use your knowledge of real-world languages. " "You may meet a task type you have never seen: read the instruction and the examples, and answer in the same form they use. " "You MUST output your reasoning inside tags first. " "After your reasoning is complete, you MUST 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. After FINAL ANSWERS:, output only the answers, exactly one line per " "numbered item, then stop." ) PHONETIC_INSTRUCTION = ( "IMPORTANT -- this problem uses PHONETIC TRANSCRIPTION. The examples write forms " "inside square brackets, like [bø:va]. Your answers must be phonetic transcriptions " "in exactly that same notation: enclosed in square brackets, using the same phonetic " "symbols. Do NOT give an English meaning or gloss -- give the transcribed FORM." ) def build_system(task_type: str, context: str, query: str) -> str: parts = [SYSTEM_BASE] if is_phonetic_task(context, query): parts.append(PHONETIC_INSTRUCTION) task_type = str(task_type).strip().lower() if task_type == "match_letters": parts.append("This is a MATCHING task. Answer with a SINGLE OPTION LETTER only (e.g., C).") elif task_type == "text_to_num": parts.append("This is a TEXT-TO-NUMBER task. Give the number in digits only (e.g., 111).") return "\n\n".join(parts) def build_user(row, n_items: int, rules: str = "", mode: str = "answer", error_feedback: str = None) -> str: content = f"{str(row['context']).strip()}\n\n{str(row['query']).strip()}" if mode == "induct": content += ( "\n\nDeduce the linguistic system as a strict ALGEBRAIC EQUATION SHEET. DO NOT write prose. DO NOT answer the QUERY yet. " "Inside tags, output ONLY the following mathematical notations based on the CONTEXT:\n\n" "1. ALIGNMENT: Define the abstract structure using variables.\n" " - If concatenative: `Word = A + B + C` (e.g., `anguls = angul + s`)\n" " - If infixing: `Word = A + Infix + B` (e.g., `sumulat = s + um + ulat`)\n" " - If templatic/ablaut: `Word = F(Root)` (e.g., `sang = Past(sing)`, `kataba = CaCaCa(k,t,b)`)\n" " - If reduplication: `Word = A + A` (e.g., `bukubuku = buku + buku`)\n" "2. MORPHOLOGY: Map variables to meanings. (e.g., `A = sing`, `Past = F()`, `s = Plural`)\n" "3. PHONOLOGY: Write exact sound changes using rule notation: /input/ -> [output] / environment. (e.g., `/v/ -> [g] / ø:_a`)\n" "4. MATCHING (if applicable): Map forms to options using matrices. (e.g., `u'u = breast = Option A`)\n" "5. NUMBERS (if applicable): Map bases mathematically. (e.g., `123 = 6 * 20^1 + 3 * 20^0`)\n\n" "Then write a line that says exactly: RULES:" ) return content if rules.strip(): content += ( f"\n\nINDUCED RULES:\n{rules.strip()}\n\n" "CRITICAL: You must solve the algebraic equations from the RULES to construct the answers. " "Do NOT guess. Do NOT blindly copy and paste full words from the context. " "Apply the exact functions, morpheme slots, and sound changes to derive the final forms." ) if n_items > 0: content += f"\n\nThere are exactly {n_items} items to answer. Give exactly {n_items} answers after FINAL ANSWERS:, one per line, no more and no fewer." if error_feedback: content += f"\n\nPREVIOUS ATTEMPT FAILED:\n{error_feedback}\n\nFix your equation solving and output the corrected answers again." return content def extract_rules(text: str) -> str: text = _TURN_NOISE.sub("", text or "") m = list(re.finditer(r"(?im)^\s*rules?\s*:?\s*$", text)) if m: return text[m[-1].end():].strip()[:2000] return text.strip()[:2000] # ──────────────────────────────────────────────────── # 6. INFERENCE ENGINE # ──────────────────────────────────────────────────── @torch.inference_mode() def generate(prompt_text: str, max_new_tokens: int, sample: bool = False, seed: int = 0): torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) enc = tok(prompt_text, return_tensors="pt", add_special_tokens=False).to(model.device) plen = enc["input_ids"].shape[1] kw = dict(do_sample=True, temperature=0.6, top_p=0.95) if sample else dict(do_sample=False) out = model.generate( **enc, max_new_tokens=max_new_tokens, use_cache=True, eos_token_id=EOS_IDS, pad_token_id=tok.pad_token_id, **kw ) gid = out[0, plen:] for pos, tid in enumerate(gid.tolist()): if tid in EOS_SET: rt = tok.decode(gid[:pos+1], skip_special_tokens=False).strip() return rt, pos+1 rt = tok.decode(gid, skip_special_tokens=False).strip() return rt, gid.shape[0] # ──────────────────────────────────────────────────── # 7. LOAD DATA & DYNAMIC INFERENCE LOOP # ──────────────────────────────────────────────────── print("Loading test data...", flush=True) df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("") if "id" not in df.columns: df["id"] = df.index results_by_id = {} run_start = time.time() order = df.index.tolist() order.sort(key=lambda idx: len(str(df.loc[idx,"context"]))+len(str(df.loc[idx,"query"]))) current_token_budget = MAX_TOKEN_BUDGET for i, row_id in enumerate(order): row = df.loc[row_id] task_type = str(row.get("task_type", "general")) n_items = count_items(str(row["query"])) t0 = time.time() messages = [] final_answers = [] raw_output = "" status = "FAIL" gen_len = 0 elapsed = time.time() - run_start remaining_time = GLOBAL_TIME_LIMIT - elapsed problems_left = len(order) - i max_allowed_tokens_for_time = max(256, int((remaining_time - (problems_left * 5)) * 10)) if max_allowed_tokens_for_time < current_token_budget: current_token_budget = max_allowed_tokens_for_time # PASS 1: Induction (Greedy, strict algebraic extraction) system_prompt = build_system(task_type, str(row["context"]), str(row["query"])) induct_user = build_user(row, n_items=0, mode="induct") induct_prompt = tok.apply_chat_template( [{"role":"system","content":system_prompt}, {"role":"user","content":induct_user}], add_generation_prompt=True, tokenize=False ) induct_raw, _ = generate(induct_prompt, INDUCT_MAX_TOKENS, sample=False, seed=1) rules = extract_rules(induct_raw) for attempt in range(MAX_ATTEMPTS): error_feedback = messages[-1] if messages else None user_prompt = build_user(row, n_items=n_items, rules=rules, mode="answer", error_feedback=error_feedback) prompt = tok.apply_chat_template( [{"role":"system","content":system_prompt}, {"role":"user","content":user_prompt}], add_generation_prompt=True, tokenize=False ) if task_type in SC_TASKS and attempt == 0: samples = [] raw_samples = [] for k in range(SC_K): raw_out, gen_len = generate(prompt, current_token_budget, sample=True, seed=1000+k*17) raw_samples.append(raw_out) samples.append(parse_answers(raw_out, n_items=n_items)) final_answers = majority_vote(samples, n_items) raw_output = "\n---\n".join(raw_samples) else: raw_output, gen_len = generate(prompt, current_token_budget, sample=False, seed=42+attempt) final_answers = parse_answers(raw_output, n_items=n_items) if len(final_answers) < n_items or not all(final_answers): msg = f"PARSE ERROR: Expected {n_items} answers, but extracted {len([a for a in final_answers if a])}. Ensure you output exactly {n_items} answers inside FINAL ANSWERS:." messages.append(msg) print(f"id={row_id} Attempt {attempt+1}: DENIED - {msg[:80]}", flush=True) continue status = "OK" print(f"id={row_id} Attempt {attempt+1}: SUCCESS", flush=True) break if status != "OK": final_answers = [""] * n_items if n_items > 0 else [] print(f"id={row_id} Failed after {MAX_ATTEMPTS} attempts.", flush=True) wt = time.time() - t0 results_by_id[row_id] = { "id": row["id"], "pred": final_answers } print(f" tok={gen_len:>4}/{current_token_budget} time={wt:.1f}s total={int(elapsed)}s", flush=True) # ──────────────────────────────────────────────────── # 8. WRITE SUBMISSION # ──────────────────────────────────────────────────── out_rows = [] for res in results_by_id.values(): out_rows.append({ "id": res["id"], "pred": json.dumps(res["pred"], ensure_ascii=False) }) pd.DataFrame(out_rows, columns=["id", "pred"]).to_csv("submission.csv", index=False) print("wrote submission.csv", flush=True)