| """IOL-AI 2026 — M1 (/think) + Offelia-style techniques. |
| |
| Keep Tiny Aya reasoning (/think). Add: |
| - cardinality: count items, tell model exact N, truncate/pad |
| - task-aware + phonetic-bracket detector (Offelia) |
| - parser hygiene: drop essay lines after FINAL ANSWERS |
| - targeted self-consistency only on match_letters / fill_blanks (k=3) |
| - induction → apply (rules sheet then answers) |
| """ |
| import os |
| import subprocess |
| import sys |
|
|
|
|
| def _install_bundled_deps() -> None: |
| wheels_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "wheels") |
| if not os.path.isdir(wheels_dir): |
| return |
| subprocess.run( |
| [ |
| sys.executable, |
| "-m", |
| "pip", |
| "install", |
| "-q", |
| "--no-index", |
| f"--find-links={wheels_dir}", |
| "transformers==4.56.2", |
| ], |
| check=True, |
| ) |
|
|
|
|
| _install_bundled_deps() |
|
|
| os.environ["HF_HUB_OFFLINE"] = "1" |
| os.environ["TRANSFORMERS_OFFLINE"] = "1" |
| MODEL_ID = "." |
| USER_THINK_TOKEN = "/think" |
|
|
| import json |
| import random |
| import re |
| from collections import Counter |
|
|
| import pandas as pd |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| END_THINKING = "<|END_THINKING|>" |
| START_THINKING = "<|START_THINKING|>" |
|
|
| THINKING_BUDGET = 1536 |
| ANSWER_CONTINUATION_TOKENS = 512 |
| COT_MAX_NEW_TOKENS = 1024 |
| INDUCT_MAX_NEW_TOKENS = 512 |
| THINK_TEMPERATURE = 0.6 |
| THINK_TOP_P = 0.95 |
| |
| SC_TASKS = frozenset({"match_letters", "fill_blanks"}) |
| SC_K = 3 |
| SYSTEM = "" |
|
|
| TASK_INSTRUCTIONS = { |
| "translation": ( |
| "This is a TRANSLATION task. Give only the translated form, in the language " |
| "the task asks for. No explanation, no source form, just the translation." |
| ), |
| "fill_blanks": ( |
| "This is a FILL-IN-THE-BLANKS task. Give only the missing form for each blank, " |
| "nothing else." |
| ), |
| "match_letters": ( |
| "This is a MATCHING task. Each numbered item must be answered with a SINGLE " |
| "OPTION LETTER only (for example: C). Do NOT write the word, meaning, or " |
| "translation -- only the letter that matches." |
| ), |
| "text_to_num": ( |
| "This is a TEXT-TO-NUMBER task. Give the number in digits only (for example: 111)." |
| ), |
| "num_to_text": ( |
| "This is a NUMBER-TO-TEXT task. Write the number out in words, in the language " |
| "the task asks for. Give only the written-out form." |
| ), |
| } |
| TASK_DEFAULT = ( |
| "Give exactly what the instruction asks for, in the same form the examples use, " |
| "and nothing else." |
| ) |
| 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." |
| ) |
| _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" |
| r"|write\s+(it\s+)?in\s+the\s+[\w'\u2019-]+\s+orthography" |
| r"|in\s+the\s+regular\s+orthography" |
| ) |
| _ASKS_TRANSCRIPTION = re.compile(r"(?i)\b(transcribe|transcription|phonetic(ally)?)\b") |
| _TURN_NOISE = re.compile( |
| r"<\|/?END_OF_TURN_TOKEN\|>|<\|/?START_OF_TURN_TOKEN\|>|" |
| r"<\|CHATBOT_TOKEN\|>|<EOS_TOKEN>|<BOS_TOKEN>" |
| ) |
| _MARKER = re.compile(r"(?im)^\s*final answers?\s*:?\s*$") |
|
|
|
|
| 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 |
|
|
|
|
| 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|" |
| r"in summary|colors? are expressed|these stems)\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 = after_thinking(text) |
| markers = list(_MARKER.finditer(text)) |
| if markers: |
| text = text[markers[-1].end() :] |
| answers = [] |
| for line in text.splitlines(): |
| line = re.sub(r"^\s*\d+[.)]\s*", "", line).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 after_thinking(text: str) -> str: |
| if END_THINKING in text: |
| text = text.rsplit(END_THINKING, 1)[-1] |
| elif START_THINKING in text: |
| text = "" |
| return _TURN_NOISE.sub("", text) |
|
|
|
|
| def build_instructions(task_type: str, context: str, query: str) -> str: |
| specific = TASK_INSTRUCTIONS.get(str(task_type).strip().lower(), TASK_DEFAULT) |
| parts = [ |
| "You solve International Linguistics Olympiad (IOL) problems from the data you are given.", |
| specific, |
| "Put answers ONLY after a line that says exactly: FINAL ANSWERS:", |
| "Bare answers only: no numbering, no quotes, no commentary, no _GCY/_NS glosses.", |
| "Never dump the alphabet. Never write an essay under FINAL ANSWERS:.", |
| ] |
| if is_phonetic_task(context, query): |
| parts.append(PHONETIC_INSTRUCTION) |
| return "\n\n".join(parts) |
|
|
|
|
| def build_user( |
| instructions: str, |
| context: str, |
| query: str, |
| *, |
| n_items: int, |
| think_token: str = "", |
| rules: str = "", |
| mode: str = "answer", |
| ) -> str: |
| parts = [instructions.strip(), "", context.strip()] |
| if rules.strip(): |
| parts += ["", "RULES:", rules.strip()] |
| parts += ["", query.strip()] |
| if mode == "induct": |
| parts += [ |
| "", |
| "Deduce linguistic RULES from CONTEXT only. Do NOT answer QUERY.", |
| "Write a bullet list under a line that says exactly: RULES:", |
| ] |
| elif n_items > 0: |
| parts += [ |
| "", |
| f"There are exactly {n_items} items to answer. " |
| f"Give exactly {n_items} answers after FINAL ANSWERS:, " |
| "one per line, no more and no fewer.", |
| ] |
| if think_token: |
| parts.append(think_token.strip()) |
| return "\n".join(parts) |
|
|
|
|
| def _end_thinking_id(tok) -> int: |
| end_id = tok.convert_tokens_to_ids(END_THINKING) |
| if end_id is None or end_id == tok.unk_token_id: |
| ids = tok.encode(END_THINKING, add_special_tokens=False) |
| if len(ids) == 1: |
| end_id = ids[0] |
| if end_id is None or end_id == tok.unk_token_id: |
| raise RuntimeError(f"missing {END_THINKING}") |
| return int(end_id) |
|
|
|
|
| def _build_prompt_ids(tok, user: str, *, thinking: bool): |
| messages = [{"role": "user", "content": user}] |
| try: |
| return tok.apply_chat_template( |
| messages, |
| add_generation_prompt=True, |
| return_tensors="pt", |
| reasoning_options={"enabled": thinking}, |
| ) |
| except TypeError: |
| return tok.apply_chat_template( |
| messages, add_generation_prompt=True, return_tensors="pt" |
| ) |
|
|
|
|
| @torch.inference_mode() |
| def generate_with_think( |
| model, |
| tok, |
| prompt_ids, |
| end_id: int, |
| *, |
| sample_think: bool, |
| think_budget: int = THINKING_BUDGET, |
| answer_tokens: int = ANSWER_CONTINUATION_TOKENS, |
| ): |
| device = next(model.parameters()).device |
| prompt_ids = prompt_ids.to(device) |
| prompt_len = prompt_ids.shape[-1] |
| think_kw = ( |
| dict(do_sample=True, temperature=THINK_TEMPERATURE, top_p=THINK_TOP_P) |
| if sample_think |
| else dict(do_sample=False) |
| ) |
| think_out = model.generate( |
| prompt_ids, |
| max_new_tokens=think_budget, |
| pad_token_id=tok.pad_token_id or tok.eos_token_id, |
| **think_kw, |
| )[0] |
| gen_ids = think_out[prompt_len:].tolist() |
| if end_id not in gen_ids: |
| cont = torch.cat( |
| [think_out, torch.tensor([end_id], device=device, dtype=think_out.dtype)] |
| ) |
| else: |
| cont = think_out |
| |
| full = model.generate( |
| cont.unsqueeze(0), |
| max_new_tokens=answer_tokens, |
| do_sample=False, |
| pad_token_id=tok.pad_token_id or tok.eos_token_id, |
| )[0] |
| return _TURN_NOISE.sub("", tok.decode(full[prompt_len:], skip_special_tokens=False)).strip() |
|
|
|
|
| @torch.inference_mode() |
| def generate_plain(model, tok, prompt_ids, max_new: int, *, sample: bool = False): |
| device = next(model.parameters()).device |
| prompt_ids = prompt_ids.to(device) |
| prompt_len = prompt_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( |
| prompt_ids, |
| max_new_tokens=max_new, |
| pad_token_id=tok.pad_token_id or tok.eos_token_id, |
| **kw, |
| )[0] |
| return _TURN_NOISE.sub("", tok.decode(out[prompt_len:], skip_special_tokens=False)).strip() |
|
|
|
|
| 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)] |
|
|
|
|
| def extract_rules(text: str) -> str: |
| text = after_thinking(text) |
| m = list(re.finditer(r"(?im)^\s*rules?\s*:?\s*$", text)) |
| if m: |
| return text[m[-1].end() :].strip()[:2000] |
| return text.strip()[:2000] |
|
|
|
|
| tok = AutoTokenizer.from_pretrained(MODEL_ID) |
| end_id = _end_thinking_id(tok) |
| model = AutoModelForCausalLM.from_pretrained( |
| MODEL_ID, torch_dtype=torch.float16, device_map="auto" |
| ).eval() |
|
|
| df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("") |
| rows = [] |
| for i, r in df.iterrows(): |
| task = str(r.get("task_type", "") or "") |
| n_items = count_items(r["query"]) |
| instr = build_instructions(task, r["context"], r["query"]) |
|
|
| |
| induct_user = build_user( |
| instr, |
| r["context"], |
| r["query"], |
| n_items=0, |
| think_token=USER_THINK_TOKEN, |
| mode="induct", |
| ) |
| induct_ids = _build_prompt_ids(tok, induct_user, thinking=True) |
| |
| induct_text = generate_with_think( |
| model, |
| tok, |
| induct_ids, |
| end_id, |
| sample_think=False, |
| think_budget=INDUCT_MAX_NEW_TOKENS, |
| answer_tokens=256, |
| ) |
| rules = extract_rules(induct_text) |
|
|
| |
| def one_answer(seed: int, sample_think: bool) -> list[str]: |
| torch.manual_seed(seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(seed) |
| user = build_user( |
| instr, |
| r["context"], |
| r["query"], |
| n_items=n_items, |
| think_token=USER_THINK_TOKEN, |
| rules=rules, |
| mode="answer", |
| ) |
| ids = _build_prompt_ids(tok, user, thinking=True) |
| text = generate_with_think( |
| model, tok, ids, end_id, sample_think=sample_think |
| ) |
| return parse_answers(text, n_items=n_items) |
|
|
| if task in SC_TASKS: |
| samples = [ |
| one_answer(1000 + int(i) * 97 + k * 17, sample_think=True) |
| for k in range(SC_K) |
| ] |
| answers = majority_vote(samples, n_items) |
| print(f" targeted SC k={SC_K} task={task}", flush=True) |
| else: |
| answers = one_answer(1000 + int(i) * 97, sample_think=True) |
|
|
| |
| if n_items > 0 and sum(1 for a in answers if a.strip()) < max(1, n_items // 2): |
| answers = one_answer(42 + int(i), sample_think=False) |
|
|
| rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)}) |
| pd.DataFrame(rows).to_csv("submission.csv", index=False) |
| print(f"[{i + 1}/{len(df)}] n={n_items} got={len(answers)} phon={is_phonetic_task(r['context'], r['query'])}", flush=True) |
|
|
| print("wrote submission.csv", flush=True) |
|
|
