393 lines
12 KiB
Python
393 lines
12 KiB
Python
|
"""IOL-AI 2026 submission — Qwen3-8B v4.
|
||
|
|
|
||
|
|
Force-close </think>, think 1024 + answer 384, parser v3, CoT 768 on format fail.
|
||
|
|
4-bit BitsAndBytes for T4 16GB / 30min.
|
||
|
|
Thinking tokens: <think> … </think>
|
||
|
|
"""
|
||
|
|
|
||
|
|
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 = "."
|
||
|
|
|
||
|
|
THINKING_BUDGET = 1024
|
||
|
|
ANSWER_CONTINUATION_TOKENS = 384
|
||
|
|
COT_MAX_NEW_TOKENS = 768
|
||
|
|
|
||
|
|
START_THINKING = "<think>"
|
||
|
|
END_THINKING = "</think>"
|
||
|
|
|
||
|
|
import json
|
||
|
|
import re
|
||
|
|
|
||
|
|
import pandas as pd
|
||
|
|
import torch
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
||
|
|
|
||
|
|
bnb = BitsAndBytesConfig(
|
||
|
|
load_in_4bit=True,
|
||
|
|
bnb_4bit_quant_type="nf4",
|
||
|
|
bnb_4bit_compute_dtype=torch.float16,
|
||
|
|
)
|
||
|
|
|
||
|
|
SYSTEM = """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."""
|
||
|
|
|
||
|
|
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."
|
||
|
|
)
|
||
|
|
|
||
|
|
_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*")
|
||
|
|
_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"
|
||
|
|
)
|
||
|
|
_THINK_TAGS = re.compile(
|
||
|
|
re.escape(START_THINKING) + r"|" + re.escape(END_THINKING)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
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:
|
||
|
|
if END_THINKING in text:
|
||
|
|
text = text.rsplit(END_THINKING, 1)[-1]
|
||
|
|
elif START_THINKING in text:
|
||
|
|
text = ""
|
||
|
|
return _THINK_TAGS.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)
|
||
|
|
answers = _lines_from_region(text, allow_all_lines=False)
|
||
|
|
if task_type == "match_letters":
|
||
|
|
coerced = []
|
||
|
|
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
|
||
|
|
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))
|
||
|
|
|
||
|
|
|
||
|
|
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) or _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]
|
||
|
|
elif len(ids) > 1:
|
||
|
|
# multi-token marker: use last id as force-append (best-effort)
|
||
|
|
end_id = ids[-1]
|
||
|
|
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",
|
||
|
|
enable_thinking=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()
|
||
|
|
# Also check decoded text in case end token is multi-piece
|
||
|
|
think_text = tok.decode(think_out[prompt_len:], skip_special_tokens=False)
|
||
|
|
if end_id not in gen_ids and END_THINKING not in think_text:
|
||
|
|
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]
|
||
|
|
return tok.decode(full[prompt_len:], skip_special_tokens=False).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]
|
||
|
|
return tok.decode(out[prompt_len:], skip_special_tokens=False).strip()
|
||
|
|
|
||
|
|
|
||
|
|
tok = AutoTokenizer.from_pretrained(MODEL_ID)
|
||
|
|
end_id = _end_thinking_id(tok)
|
||
|
|
model = AutoModelForCausalLM.from_pretrained(
|
||
|
|
MODEL_ID, quantization_config=bnb, 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:.)"
|
||
|
|
|
||
|
|
ids = _build_prompt_ids(tok, SYSTEM, user, 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_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)
|