528 lines
17 KiB
Python
528 lines
17 KiB
Python
"""IOL-AI 2026 — M1 iterate on best private~0.083 (temp0.6 + user-instr).
|
||
|
||
Keep: user-prompt instructions, /think, sample think @ TEMPERATURE.
|
||
CSV fixes from that run: strip _GCY/gloss; reject essay+alphabet+resample;
|
||
stronger user prompt; think 2048; greedy answer continuation.
|
||
"""
|
||
|
||
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 = "/think"
|
||
|
||
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 = 2048
|
||
ANSWER_CONTINUATION_TOKENS = 512
|
||
COT_MAX_NEW_TOKENS = 1024
|
||
TEMPERATURE = 0.6
|
||
TOP_P = 0.95
|
||
QUALITY_RESAMPLES = 4
|
||
ANSWER_GREEDY = True # format-sensitive answer phase
|
||
|
||
# Empty system — instructions go on the user turn (decode ablation vs v4)
|
||
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."
|
||
)
|
||
|
||
# --- 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 _strip_gloss_keep_form(line: str) -> str:
|
||
s = (line or "").strip()
|
||
s = re.sub(r"\*\*", "", s)
|
||
s = re.split(r"\s+_?(?:GCY|NS|N/A)_?\b", s, maxsplit=1, flags=re.I)[0].strip()
|
||
s = re.sub(r"\s+_?(?:GCY|NS|N/A)_?\s*$", "", s, flags=re.I).strip()
|
||
m = re.match(
|
||
r"^(.+?)\s+[-–—]\s+((?:to|the|a|an|in|of|for|being|means?|black|white|red|green|yellow)\b.*)$",
|
||
s,
|
||
flags=re.I,
|
||
)
|
||
if m:
|
||
s = m.group(1).strip()
|
||
return s.strip()
|
||
|
||
|
||
def _clean_line(line: str) -> str:
|
||
line = _NUMBERING.sub("", line).strip()
|
||
line = _MD_WRAP.sub("", line).strip()
|
||
line = line.replace("\u202f", " ").replace("\xa0", " ")
|
||
return _strip_gloss_keep_form(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) < 8:
|
||
return False
|
||
seq = 0
|
||
for i, L in enumerate(letters):
|
||
if ord(L) == ord("A") + i:
|
||
seq += 1
|
||
else:
|
||
break
|
||
return seq >= 8
|
||
|
||
|
||
def _looks_like_essay(answers: list[str]) -> bool:
|
||
if any(len(str(a)) > 100 for a in answers):
|
||
return True
|
||
blob = " ".join(map(str, answers))
|
||
return bool(
|
||
re.search(
|
||
r"(?i)\b(colors? are expressed|systematic set of lexical|"
|
||
r"these stems are combined|step by step|as an ai|verification)\b",
|
||
blob,
|
||
)
|
||
)
|
||
|
||
|
||
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 _looks_like_essay(answers):
|
||
return False
|
||
if any(re.search(r"(?i)_GCY\b|_NS\b", str(a)) for a in answers):
|
||
return False
|
||
if task_type != "match_letters" and len(answers) >= 4:
|
||
if all(re.fullmatch(r"[ivxlcdm]+", str(a).strip(), flags=re.I) for a in answers):
|
||
return False
|
||
if n_expected is not None and n_expected > 0 and abs(len(answers) - n_expected) > max(
|
||
2, n_expected // 2
|
||
):
|
||
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 = []
|
||
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",
|
||
reasoning_options={"enabled": thinking},
|
||
)
|
||
except TypeError:
|
||
return tok.apply_chat_template(
|
||
messages, add_generation_prompt=True, return_tensors="pt"
|
||
)
|
||
|
||
|
||
def _sample_kwargs():
|
||
return dict(
|
||
do_sample=True,
|
||
temperature=TEMPERATURE,
|
||
top_p=TOP_P,
|
||
)
|
||
|
||
|
||
@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]
|
||
gen_kw = _sample_kwargs()
|
||
|
||
think_out = model.generate(
|
||
prompt_ids,
|
||
max_new_tokens=THINKING_BUDGET,
|
||
pad_token_id=tok.pad_token_id or tok.eos_token_id,
|
||
**gen_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
|
||
|
||
ans_kw = dict(do_sample=False) if ANSWER_GREEDY else gen_kw
|
||
full = model.generate(
|
||
cont.unsqueeze(0),
|
||
max_new_tokens=ANSWER_CONTINUATION_TOKENS,
|
||
pad_token_id=tok.pad_token_id or tok.eos_token_id,
|
||
**ans_kw,
|
||
)[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,
|
||
pad_token_id=tok.pad_token_id or tok.eos_token_id,
|
||
**_sample_kwargs(),
|
||
)[0]
|
||
text = tok.decode(out[prompt_len:], skip_special_tokens=False)
|
||
return _TURN_NOISE.sub("", text).strip()
|
||
|
||
|
||
def _build_user(
|
||
instructions: str,
|
||
context: str,
|
||
query: str,
|
||
*,
|
||
n_guess: int,
|
||
think_token: str = "",
|
||
) -> str:
|
||
parts = [
|
||
instructions.strip(),
|
||
"",
|
||
context.strip(),
|
||
"",
|
||
query.strip(),
|
||
]
|
||
if n_guess:
|
||
parts.append("")
|
||
parts.append(f"(Emit exactly {n_guess} answer line(s) after FINAL ANSWERS:.)")
|
||
if think_token:
|
||
parts.append(think_token.strip())
|
||
return "\n".join(parts)
|
||
|
||
|
||
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_think = _build_user(
|
||
USER_INSTRUCTIONS,
|
||
r["context"],
|
||
r["query"],
|
||
n_guess=n_guess,
|
||
think_token=USER_THINK_TOKEN,
|
||
)
|
||
user_plain = _build_user(
|
||
USER_INSTRUCTIONS_COT,
|
||
r["context"],
|
||
r["query"],
|
||
n_guess=n_guess,
|
||
think_token="",
|
||
)
|
||
|
||
best_answers: list[str] = []
|
||
used_cot = False
|
||
for attempt in range(QUALITY_RESAMPLES):
|
||
torch.manual_seed(1000 + int(i) * 97 + attempt * 31)
|
||
if torch.cuda.is_available():
|
||
torch.cuda.manual_seed_all(1000 + int(i) * 97 + attempt * 31)
|
||
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)
|
||
answers = [_strip_gloss_keep_form(a) for a in answers if _strip_gloss_keep_form(a)]
|
||
if n_exp and n_exp > 0 and len(answers) > n_exp:
|
||
answers = answers[:n_exp]
|
||
if has_usable_answer(answers, n_expected=n_exp, task_type=task):
|
||
best_answers = answers
|
||
break
|
||
if answers and not best_answers:
|
||
best_answers = answers
|
||
print(
|
||
f" quality resample {attempt + 1}/{QUALITY_RESAMPLES} id={r['id']} n={len(answers)}",
|
||
flush=True,
|
||
)
|
||
answers = best_answers
|
||
|
||
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, user_plain, 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)
|
||
cot_answers = [
|
||
_strip_gloss_keep_form(a) for a in cot_answers if _strip_gloss_keep_form(a)
|
||
]
|
||
if has_usable_answer(cot_answers, n_expected=n_exp, task_type=task) or (
|
||
cot_answers and not answers
|
||
):
|
||
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)
|