| """IOL-AI 2026 — notebook decode (prompt + force-close + T=0.6). |
| |
| Matches linguini-test-iolai.ipynb cells 4–5: |
| - USER_INSTRUCTIONS (+ COT + /think) |
| - think 2048 @ T=0.6; force <|END_THINKING|><|START_RESPONSE|> |
| - answer 512 @ T=0.6 |
| - marker-gated FINAL ANSWERS: parser |
| """ |
|
|
| from __future__ import annotations |
|
|
| 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 = "." |
|
|
| import json |
| import re |
|
|
| import pandas as pd |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| SYSTEM = "" |
|
|
| USER_INSTRUCTIONS = """You solve International Linguistics Olympiad (IOL) problems from the data you are given. |
| You may see a task type you have never seen: follow the instruction and examples, and answer in the same form they use. |
| What to return by task type: |
| - translation: only the required form in the language the query asks for — do not add extra glosses or "form | meaning" unless asked |
| - fill_blanks: only the missing form for each blank — no extra glosses |
| - match_letters: ONLY the option letter (A, B, C, …), one letter per line — never copy option text, never arrows, never "A. word" |
| - text_to_num: the number in digits only |
| - num_to_text: the number written out in words, in the language asked |
| - kinship / sentence matching: the full required sentence or form — NOT roman numerals (i, ii, iii) and NOT an alphabet dump |
| - any other type: exactly what the instruction asks for, nothing else |
| Answer in the language and form the query asks for. Do not add glosses, translations, or explanations unless the instruction requires them. |
| Output rules: |
| - Put answers ONLY after a line that says exactly: FINAL ANSWERS: |
| - Never put answers before that marker. |
| - One answer per line; exactly as many lines as items asked in the query. |
| - Bare answers only: no numbering, no quotes, no commentary, no repeating the question. |
| Extra hard rules: |
| - Never refuse or apologize; always output FINAL ANSWERS: with your best guess. |
| - For match_letters: only bare letters (A, B, C, …) — never dump the alphabet (A B C D E F…), never option text. |
| - Never append tags or glosses: no "_GCY", "_NS", "form – meaning", "word - gloss", or markdown bold. |
| - Never write an essay or explanation of how the language works under FINAL ANSWERS: — only the answer strings. |
| - If the query asks for colour/color forms, output those forms only (one per line), not a linguistics write-up. |
| - Emit exactly as many answer lines as items asked — no more, no fewer.""" |
|
|
| USER_INSTRUCTIONS_COT = ( |
| USER_INSTRUCTIONS |
| + "\n\nThink step by step about the rules in the examples and how they apply to the query, " |
| "then write FINAL ANSWERS: and the answer lines." |
| ) |
|
|
| APPEND_THINK = True |
|
|
| THINK_BUDGET = 2048 |
| ANSWER_BUDGET = 512 |
| TEMPERATURE = 0.6 |
|
|
|
|
| def build_user(context, query): |
| instructions = USER_INSTRUCTIONS_COT if APPEND_THINK else USER_INSTRUCTIONS |
| parts = [instructions.strip(), "", context.strip(), "", query.strip()] |
| if APPEND_THINK: |
| parts.append("/think") |
| return "\n".join(parts) |
|
|
|
|
| _SPECIAL = re.compile( |
| r"<\|START_RESPONSE\|>|<\|END_RESPONSE\|>|" |
| r"<\|/?END_OF_TURN_TOKEN\|>|<\|/?START_OF_TURN_TOKEN\|>|" |
| r"<\|CHATBOT_TOKEN\|>|<\|/?START_THINKING\|>|<\|/?END_THINKING\|>|" |
| r"<EOS_TOKEN>|<BOS_TOKEN>|<PAD>" |
| ) |
| _MARKER = re.compile(r"(?im)\bfinal\s+answers?\b\s*:?\s*") |
| _TURN_NOISE = re.compile( |
| r"<\|/?END_OF_TURN_TOKEN\|>|<\|/?START_OF_TURN_TOKEN\|>|" |
| r"<\|CHATBOT_TOKEN\|>|<EOS_TOKEN>|<BOS_TOKEN>" |
| ) |
|
|
|
|
| def _lines_after(region: str) -> list[str]: |
| answers = [] |
| for line in region.splitlines(): |
| line = _SPECIAL.sub("", line) |
| line = re.sub(r"^\s*(?:\(?\d+[.)]|[-*•])\s*", "", line).strip() |
| if not line: |
| if answers: |
| break |
| continue |
| if re.fullmatch(r"(?i)final\s+answers?\s*:?", line): |
| continue |
| if answers and ( |
| len(line) > 120 |
| or line.lower().startswith(("however", "but the", "to solve", "### ", "**")) |
| ): |
| break |
| answers.append(line) |
| return answers |
|
|
|
|
| def parse_answers(text: str) -> list[str]: |
| raw = text |
|
|
| def _from_region(region: str) -> list[str] | None: |
| markers = list(_MARKER.finditer(region)) |
| if not markers: |
| return None |
| for m in reversed(markers): |
| answers = _lines_after(region[m.end() :]) |
| if answers: |
| return answers |
| return [] |
|
|
| region = raw |
| if "<|END_THINKING|>" in region: |
| region = region.rsplit("<|END_THINKING|>", 1)[-1] |
| blocks = re.findall( |
| r"<\|START_RESPONSE\|>(.*?)<\|END_RESPONSE\|>", region, flags=re.S |
| ) |
| if blocks: |
| region = blocks[-1] |
|
|
| found = _from_region(region) |
| if found is not None: |
| return found |
|
|
| found = _from_region(raw) |
| if found is not None: |
| return found |
|
|
| return [] |
|
|
|
|
| def _build_prompt(tok, system: str, user: str): |
| messages = [] |
| if system.strip(): |
| messages.append({"role": "system", "content": system}) |
| messages.append({"role": "user", "content": user}) |
| try: |
| return tok.apply_chat_template( |
| messages, |
| add_generation_prompt=True, |
| return_tensors="pt", |
| return_dict=True, |
| reasoning_options={"enabled": True}, |
| ) |
| except TypeError: |
| ids = tok.apply_chat_template( |
| messages, add_generation_prompt=True, return_tensors="pt" |
| ) |
| return {"input_ids": ids, "attention_mask": torch.ones_like(ids)} |
|
|
|
|
| @torch.inference_mode() |
| def generate_with_budgets(model, tok, enc): |
| device = next(model.parameters()).device |
| enc = {k: v.to(device) if hasattr(v, "to") else v for k, v in enc.items()} |
| prompt_len = enc["input_ids"].shape[-1] |
| pad_id = tok.pad_token_id or tok.eos_token_id |
| end_ids = tok.encode("<|END_THINKING|>", add_special_tokens=False) |
| start_ids = tok.encode("<|START_RESPONSE|>", add_special_tokens=False) |
|
|
| out = model.generate( |
| **enc, |
| max_new_tokens=THINK_BUDGET, |
| do_sample=True, |
| temperature=TEMPERATURE, |
| eos_token_id=tok.eos_token_id, |
| pad_token_id=pad_id, |
| ) |
| gen = out[0] |
| new_ids = gen[prompt_len:].tolist() |
| forced = False |
| finished = bool(new_ids) and new_ids[-1] == tok.eos_token_id |
|
|
| if end_ids[0] not in new_ids: |
| forced = True |
| cont = torch.tensor( |
| end_ids + start_ids, device=gen.device, dtype=gen.dtype |
| ) |
| gen = torch.cat([gen, cont], dim=0) |
| finished = False |
|
|
| if not finished: |
| ids = gen.unsqueeze(0) |
| out2 = model.generate( |
| input_ids=ids, |
| attention_mask=torch.ones_like(ids), |
| max_new_tokens=ANSWER_BUDGET, |
| do_sample=True, |
| temperature=TEMPERATURE, |
| eos_token_id=tok.eos_token_id, |
| pad_token_id=pad_id, |
| ) |
| gen = out2[0] |
|
|
| text = tok.decode(gen[prompt_len:], skip_special_tokens=False).strip() |
| return _TURN_NOISE.sub("", text).strip(), forced |
|
|
|
|
| tok = AutoTokenizer.from_pretrained(MODEL_ID) |
| 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(): |
| user = build_user(r["context"], r["query"]) |
| enc = _build_prompt(tok, SYSTEM, user) |
| text, forced = generate_with_budgets(model, tok, enc) |
| answers = parse_answers(text) |
| 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={len(answers)} forced={forced}", flush=True) |
|
|
| print("wrote submission.csv", flush=True) |
|
|
