440 lines
14 KiB
Python
440 lines
14 KiB
Python
"""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
|
||
# Targeted SC only
|
||
SC_TASKS = frozenset({"match_letters", "fill_blanks"})
|
||
SC_K = 3
|
||
SYSTEM = "" # instructions on user turn (best M1 private recipe)
|
||
|
||
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
|
||
# match_letters letter blob
|
||
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
|
||
# greedy answer
|
||
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]
|
||
# prefer full-tuple agreement
|
||
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"])
|
||
|
||
# Pass A: induction (greedy think)
|
||
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)
|
||
# Short greedy think for rules only — keep T4 headroom for answer (+ SC).
|
||
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)
|
||
|
||
# Pass B: answer with rules
|
||
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)
|
||
|
||
# Fallback: no-rules single greedy think if mostly empty
|
||
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)
|