422 lines
14 KiB
Python
422 lines
14 KiB
Python
"""IOL-AI 2026 submission — Tiny Aya bakeoff v4.
|
||
|
||
Force-close <|END_THINKING|>, 1536+512, parser v3, prompt v3.
|
||
CoT fallback (1024) when format/parse fails.
|
||
M1: set USER_THINK_TOKEN="/think" + instruction-following addendum.
|
||
A1: leave USER_THINK_TOKEN="".
|
||
"""
|
||
|
||
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 = "."
|
||
|
||
# "" for A1 (reasoning_options only); "/think" for M1 multilingual
|
||
USER_THINK_TOKEN = ""
|
||
|
||
import json
|
||
import re
|
||
|
||
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 # T4 30min: capped CoT, not another 1536+512
|
||
|
||
_SYSTEM_BASE = """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
|
||
- 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."""
|
||
|
||
_M1_ADDENDUM = """
|
||
|
||
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, never option text.
|
||
- Never append glosses like "form – meaning" or "word - gloss"; bare answers only.
|
||
- Emit exactly as many answer lines as items asked — no more, no fewer."""
|
||
|
||
SYSTEM = _SYSTEM_BASE + (_M1_ADDENDUM if USER_THINK_TOKEN == "/think" else "")
|
||
SYSTEM_COT = (
|
||
SYSTEM
|
||
+ "\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."
|
||
)
|
||
|
||
# --- parser (inlined from parse_iol.py) ---
|
||
_MD_PREFIX = r"(?:[#*_=\-\s`>]*)"
|
||
_MARKER = re.compile(
|
||
rf"(?im)^{_MD_PREFIX}final\s+answers?{_MD_PREFIX}:?{_MD_PREFIX}\s*(.*)$"
|
||
)
|
||
_NUMBERING = re.compile(r"^\s*(?:\d+[.)]|[-*•])\s*")
|
||
_TURN_NOISE = re.compile(
|
||
r"<\|/?END_OF_TURN_TOKEN\|>|<\|/?START_OF_TURN_TOKEN\|>|"
|
||
r"<\|CHATBOT_TOKEN\|>|<EOS_TOKEN>|<BOS_TOKEN>"
|
||
)
|
||
_RESPONSE_BLOCK = re.compile(
|
||
r"<\|START_RESPONSE\|>(.*?)<\|END_RESPONSE\|>",
|
||
flags=re.S,
|
||
)
|
||
_MD_WRAP = re.compile(r"^[*_`#\s]+|[*_`#\s]+$")
|
||
_TRAILING_LETTER = re.compile(
|
||
r"(?:[–—\-]|→|->)\s*([A-Za-z])(?:\s*[.)]|)\s*$"
|
||
)
|
||
_LEADING_LETTER_OPT = re.compile(r"^([A-Za-z])\s*[.):\-–—]\s+\S")
|
||
_WORD_THEN_LETTER = re.compile(r"^.+\s([A-Za-z])\s*$")
|
||
_REFUSAL = re.compile(
|
||
r"(?i)\b("
|
||
r"i'?m sorry|i am sorry|i don'?t have|i cannot|i can'?t|"
|
||
r"unable to|not able to|no reliable|cannot supply|can'?t supply|"
|
||
r"as an ai|i apologize"
|
||
r")\b"
|
||
)
|
||
|
||
|
||
def _clean_line(line: str) -> str:
|
||
line = _NUMBERING.sub("", line).strip()
|
||
line = _MD_WRAP.sub("", line).strip()
|
||
line = line.replace("\u202f", " ").replace("\xa0", " ")
|
||
return line.strip()
|
||
|
||
|
||
def after_thinking(text: str) -> str:
|
||
"""Prefer content after the last <|END_THINKING|>; else drop an unclosed think block."""
|
||
if END_THINKING in text:
|
||
text = text.rsplit(END_THINKING, 1)[-1]
|
||
elif START_THINKING in text:
|
||
text = ""
|
||
return _TURN_NOISE.sub("", text)
|
||
|
||
|
||
def _as_option_letter(line: str) -> str | None:
|
||
line = _clean_line(line)
|
||
if not line:
|
||
return None
|
||
if len(line) == 1 and line.isalpha():
|
||
return line.upper()
|
||
m = _LEADING_LETTER_OPT.match(line)
|
||
if m:
|
||
return m.group(1).upper()
|
||
m = _TRAILING_LETTER.search(line)
|
||
if m:
|
||
return m.group(1).upper()
|
||
if len(line) <= 40:
|
||
m = _WORD_THEN_LETTER.match(line)
|
||
if m:
|
||
return m.group(1).upper()
|
||
return None
|
||
|
||
|
||
def _expand_line(line: str) -> list[str]:
|
||
line = _clean_line(line)
|
||
if not line:
|
||
return []
|
||
if len(line) == 1 and line.isalpha():
|
||
return [line]
|
||
if _LEADING_LETTER_OPT.match(line) or _TRAILING_LETTER.search(line):
|
||
letter = _as_option_letter(line)
|
||
if letter:
|
||
return [letter]
|
||
if len(line) <= 40 and _WORD_THEN_LETTER.match(line):
|
||
letter = _as_option_letter(line)
|
||
if letter:
|
||
return [letter]
|
||
if "|" in line:
|
||
parts = [p.strip() for p in line.split("|") if p.strip()]
|
||
if len(parts) >= 2:
|
||
if len(parts) >= 4 and len(parts) % 2 == 0:
|
||
left, right = parts[0::2], parts[1::2]
|
||
if sum(" " in r for r in right) >= max(1, len(right) // 2):
|
||
return [_clean_line(x) for x in left if _clean_line(x)]
|
||
if len(parts) == 2:
|
||
a, b = parts
|
||
if (" " in b and " " not in a) or (
|
||
len(b) > 2 * max(len(a), 1) and " " in b
|
||
):
|
||
return [_clean_line(a)] if _clean_line(a) else []
|
||
return [_clean_line(p) for p in parts if _clean_line(p)]
|
||
return [line]
|
||
|
||
|
||
def _dedupe_runaway(parts: list[str]) -> list[str]:
|
||
if len(parts) < 6:
|
||
return parts
|
||
out: list[str] = []
|
||
run = 0
|
||
prev = None
|
||
for p in parts:
|
||
if p == prev:
|
||
run += 1
|
||
if run >= 4:
|
||
break
|
||
else:
|
||
run = 1
|
||
prev = p
|
||
out.append(p)
|
||
return out
|
||
|
||
|
||
def _lines_from_region(region: str, *, allow_all_lines: bool) -> list[str]:
|
||
markers = list(_MARKER.finditer(region))
|
||
if markers:
|
||
last = markers[-1]
|
||
after_parts: list[str] = []
|
||
same = _clean_line(last.group(1) or "")
|
||
if same:
|
||
after_parts.extend(_expand_line(same))
|
||
for line in region[last.end() :].splitlines():
|
||
after_parts.extend(_expand_line(line))
|
||
if after_parts:
|
||
return _dedupe_runaway(after_parts)
|
||
before_parts: list[str] = []
|
||
for line in region[: last.start()].splitlines():
|
||
before_parts.extend(_expand_line(line))
|
||
if before_parts:
|
||
return _dedupe_runaway(before_parts)
|
||
|
||
parts: list[str] = []
|
||
for line in region.splitlines():
|
||
parts.extend(_expand_line(line))
|
||
if not parts:
|
||
return []
|
||
if allow_all_lines:
|
||
return _dedupe_runaway(parts)
|
||
return [parts[-1]]
|
||
|
||
|
||
def parse_answers(
|
||
raw: str,
|
||
*,
|
||
n_expected: int | None = None,
|
||
task_type: str = "",
|
||
) -> list[str]:
|
||
text = after_thinking(raw)
|
||
closed_blocks = _RESPONSE_BLOCK.findall(text)
|
||
answers: list[str] = []
|
||
if closed_blocks:
|
||
for region in reversed(closed_blocks):
|
||
answers = _lines_from_region(region.strip(), allow_all_lines=True)
|
||
if answers:
|
||
break
|
||
if not answers:
|
||
answers = _lines_from_region(text, allow_all_lines=False)
|
||
|
||
if task_type == "match_letters":
|
||
coerced: list[str] = []
|
||
for a in answers:
|
||
letter = _as_option_letter(a)
|
||
coerced.append(letter if letter else a)
|
||
answers = coerced
|
||
|
||
if n_expected is not None and n_expected > 0 and len(answers) > n_expected:
|
||
answers = answers[:n_expected]
|
||
return answers
|
||
|
||
|
||
def _looks_like_alphabet_dump(answers: list[str]) -> bool:
|
||
letters = [a.strip().upper() for a in answers if len(a.strip()) == 1 and a.strip().isalpha()]
|
||
if len(letters) < 15:
|
||
return False
|
||
# sequential A,B,C… for a long prefix
|
||
seq = 0
|
||
for i, L in enumerate(letters):
|
||
if ord(L) == ord("A") + i:
|
||
seq += 1
|
||
else:
|
||
break
|
||
return seq >= 15
|
||
|
||
|
||
def _looks_like_refusal(answers: list[str]) -> bool:
|
||
blob = " ".join(answers)
|
||
return bool(_REFUSAL.search(blob)) or len(blob) > 400 and "dictionary" in blob.lower()
|
||
|
||
|
||
def has_usable_answer(
|
||
answers: list[str],
|
||
*,
|
||
n_expected: int | None = None,
|
||
task_type: str = "",
|
||
) -> bool:
|
||
if not answers or not any(a.strip() for a in answers):
|
||
return False
|
||
if _looks_like_refusal(answers):
|
||
return False
|
||
if _looks_like_alphabet_dump(answers):
|
||
return False
|
||
if n_expected is not None and n_expected > 0 and len(answers) != n_expected:
|
||
return False
|
||
if task_type == "match_letters":
|
||
letters = [a for a in answers if len(a) == 1 and a.isalpha()]
|
||
if len(letters) < max(1, int(0.8 * len(answers))):
|
||
return False
|
||
return True
|
||
|
||
|
||
def _n_items_guess(query: str) -> int:
|
||
nums = re.findall(r"(?m)^\s*(?:\(?\d+[.)]|\d+\))", query)
|
||
return len(nums) if nums else 0
|
||
|
||
|
||
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"Tokenizer missing end-think token {END_THINKING!r}")
|
||
return int(end_id)
|
||
|
||
|
||
def _build_prompt_ids(tok, system: str, user: str, *, thinking: bool):
|
||
messages = [
|
||
{"role": "system", "content": system},
|
||
{"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_budget(model, tok, prompt_ids, end_id: int):
|
||
device = next(model.parameters()).device
|
||
prompt_ids = prompt_ids.to(device)
|
||
prompt_len = prompt_ids.shape[-1]
|
||
|
||
think_out = model.generate(
|
||
prompt_ids,
|
||
max_new_tokens=THINKING_BUDGET,
|
||
do_sample=False,
|
||
pad_token_id=tok.pad_token_id or tok.eos_token_id,
|
||
)[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_CONTINUATION_TOKENS,
|
||
do_sample=False,
|
||
pad_token_id=tok.pad_token_id or tok.eos_token_id,
|
||
)[0]
|
||
text = tok.decode(full[prompt_len:], skip_special_tokens=False)
|
||
return _TURN_NOISE.sub("", text).strip()
|
||
|
||
|
||
@torch.inference_mode()
|
||
def generate_plain(model, tok, prompt_ids, max_new_tokens: int):
|
||
device = next(model.parameters()).device
|
||
prompt_ids = prompt_ids.to(device)
|
||
prompt_len = prompt_ids.shape[-1]
|
||
out = model.generate(
|
||
prompt_ids,
|
||
max_new_tokens=max_new_tokens,
|
||
do_sample=False,
|
||
pad_token_id=tok.pad_token_id or tok.eos_token_id,
|
||
)[0]
|
||
text = tok.decode(out[prompt_len:], skip_special_tokens=False)
|
||
return _TURN_NOISE.sub("", text).strip()
|
||
|
||
|
||
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():
|
||
n_guess = _n_items_guess(r["query"])
|
||
task = str(r.get("task_type", "") or "")
|
||
n_exp = n_guess or None
|
||
|
||
user = f"{r['context'].strip()}\n\n{r['query'].strip()}"
|
||
if n_guess:
|
||
user += f"\n\n(Emit exactly {n_guess} answer line(s) after FINAL ANSWERS:.)"
|
||
user_think = user
|
||
if USER_THINK_TOKEN:
|
||
user_think += f"\n{USER_THINK_TOKEN}"
|
||
|
||
ids = _build_prompt_ids(tok, SYSTEM, user_think, thinking=True)
|
||
text = generate_with_think_budget(model, tok, ids, end_id)
|
||
answers = parse_answers(text, n_expected=n_exp, task_type=task)
|
||
used_cot = False
|
||
|
||
if not has_usable_answer(answers, n_expected=n_exp, task_type=task):
|
||
used_cot = True
|
||
# CoT: no /think, thinking channel off, capped budget
|
||
cot_ids = _build_prompt_ids(tok, SYSTEM_COT, user, thinking=False)
|
||
cot_text = generate_plain(model, tok, cot_ids, COT_MAX_NEW_TOKENS)
|
||
cot_answers = parse_answers(cot_text, n_expected=n_exp, task_type=task)
|
||
if has_usable_answer(cot_answers, n_expected=n_exp, task_type=task):
|
||
answers = cot_answers
|
||
|
||
rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
|
||
print(
|
||
f"[{i + 1}/{len(df)}] {len(answers)} answers cot={used_cot}",
|
||
flush=True,
|
||
)
|
||
|
||
pd.DataFrame(rows).to_csv("submission.csv", index=False)
|
||
print("wrote submission.csv", flush=True)
|