221 lines
8.5 KiB
Python
221 lines
8.5 KiB
Python
"""IOL-AI 2026 submission — v9-refine (same Qwen weights).
|
||
|
||
Two-pass directed self-critique. Pass 1 is v5's exact greedy run (drafts
|
||
written to submission.csv as a safety net). Pass 2 shows the model its own
|
||
draft answers alongside the problem and asks it to verify each one against
|
||
the data — form, language, one-to-one letter use, arithmetic — and re-emit
|
||
corrected FINAL ANSWERS. The draft is kept whenever the refine pass fails
|
||
to produce a parseable answer list or time runs short.
|
||
"""
|
||
|
||
import os
|
||
|
||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
|
||
|
||
import json
|
||
import re
|
||
import time
|
||
|
||
import pandas as pd
|
||
import torch
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
||
TAG = "v9-refine" # submission tag, printed into the eval log
|
||
MODEL_ID = "."
|
||
T_START = time.monotonic()
|
||
TIME_BUDGET_S = 27 * 60
|
||
MAX_NEW_TOKENS = 1280
|
||
# Minimum seconds that must remain before starting a refine generation.
|
||
MIN_S_PER_GEN = 150
|
||
|
||
SYSTEM_PROMPT = (
|
||
"You solve International Linguistics Olympiad problems by reasoning from "
|
||
"the data you are given. Everything needed is in the problem data. "
|
||
"Reason step by step first: work out the vocabulary and rules from the "
|
||
"examples, and double-check each answer against them. "
|
||
"Each answer must be the same kind of form that appears in that position "
|
||
"in the examples: if the examples show a phonetic transcription in "
|
||
"[brackets], give a phonetic transcription in [brackets], not the English "
|
||
"meaning; if they show a word or sentence in the problem language, answer "
|
||
"in the problem language, copying its exact spelling conventions. For "
|
||
"matching items give just the option letter -- the items and options form "
|
||
"a one-to-one matching, so use each letter exactly once, never repeating "
|
||
"or inventing letters. For number items give digits or the spelled-out "
|
||
"number as the query asks. "
|
||
"Then write a line that says exactly FINAL ANSWERS: and, below it, one "
|
||
"answer per line in the order the items are asked -- the bare answer "
|
||
"only, no numbering, no quotes, no extra text."
|
||
)
|
||
|
||
REFINE_PROMPT = (
|
||
"Below is a linguistics problem and DRAFT answers produced on a first "
|
||
"attempt. Carefully verify every draft answer against the problem data: "
|
||
"re-derive the rules from the examples and check each answer follows "
|
||
"them. Look especially for: answers in the wrong form or language "
|
||
"compared to what fills that position in the examples; matching answers "
|
||
"that repeat an option letter or use a letter that is not offered (the "
|
||
"matching is one-to-one, each letter exactly once); number conversions "
|
||
"that do not satisfy the number system of the data; and morphology that "
|
||
"contradicts a pattern in the examples. Keep every draft answer that "
|
||
"checks out, and correct the ones that do not. "
|
||
"Then write a line that says exactly FINAL ANSWERS: and, below it, one "
|
||
"answer per line in the same order -- the bare answer only, no "
|
||
"numbering, no quotes, no extra text.\n\n"
|
||
"{context}\n\n{query}\n\nDRAFT ANSWERS:\n{draft}"
|
||
)
|
||
|
||
|
||
def count_items(query: str):
|
||
"""Count numbered items in a query; None when no numbering is detectable."""
|
||
nums = set()
|
||
for m in re.finditer(r"\((\d{1,3})\s*[-–]\s*(\d{1,3})\)", query):
|
||
a, b = int(m.group(1)), int(m.group(2))
|
||
if a < b and b - a < 40:
|
||
nums.update(range(a, b + 1))
|
||
for m in re.finditer(r"^\s*(\d{1,3})[.)]\s", query, re.M):
|
||
nums.add(int(m.group(1)))
|
||
for m in re.finditer(r"\((\d{1,3})\)", query):
|
||
nums.add(int(m.group(1)))
|
||
nums = {n for n in nums if 1 <= n <= 999}
|
||
return len(nums) if nums else None
|
||
|
||
|
||
def parse_answers(text: str, n_items):
|
||
"""Lines after the last FINAL ANSWERS: marker; pad to n_items, never truncate."""
|
||
marker = None
|
||
for marker in re.finditer(r"(?i)final answers?\s*:", text):
|
||
pass
|
||
if marker is not None:
|
||
zone = text[marker.end() :]
|
||
else:
|
||
zone = text
|
||
|
||
answers = [ln.strip() for ln in zone.splitlines() if ln.strip()]
|
||
answers = [re.sub(r"^\s*\d{1,3}[.)]\s*", "", x) for x in answers]
|
||
answers = [
|
||
x[1:-1].strip() if len(x) >= 2 and x[0] == x[-1] and x[0] in "'\"" else x
|
||
for x in answers
|
||
]
|
||
|
||
if marker is None and n_items is not None and len(answers) > n_items:
|
||
answers = answers[-n_items:]
|
||
if n_items is not None and len(answers) < n_items:
|
||
answers = answers + [""] * (n_items - len(answers))
|
||
return answers
|
||
|
||
|
||
def generate(tok, model, messages):
|
||
chat = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
||
inputs = tok(chat, return_tensors="pt").to(model.device)
|
||
with torch.no_grad():
|
||
out = model.generate(
|
||
**inputs,
|
||
max_new_tokens=MAX_NEW_TOKENS,
|
||
do_sample=False,
|
||
pad_token_id=tok.pad_token_id or tok.eos_token_id,
|
||
)
|
||
return tok.decode(
|
||
out[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=True
|
||
).strip()
|
||
|
||
|
||
def write_submission(df, preds):
|
||
rows = [
|
||
{"id": r["id"], "pred": json.dumps(preds[i], ensure_ascii=False)}
|
||
for i, (_, r) in enumerate(df.iterrows())
|
||
]
|
||
pd.DataFrame(rows).to_csv("submission.csv", index=False)
|
||
|
||
|
||
def main():
|
||
print(f"submission tag: {TAG}", flush=True)
|
||
print("loading model...", flush=True)
|
||
tok = AutoTokenizer.from_pretrained(MODEL_ID)
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
MODEL_ID,
|
||
torch_dtype=torch.float16, # T4 has no bfloat16
|
||
device_map="auto",
|
||
attn_implementation="sdpa",
|
||
).eval()
|
||
print(f"model loaded in {time.monotonic() - T_START:.0f}s", flush=True)
|
||
|
||
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
|
||
n_items_list = [count_items(r["query"]) for _, r in df.iterrows()]
|
||
|
||
# Pass 1: greedy drafts (identical to the proven v5 run).
|
||
drafts = []
|
||
for i, (_, r) in enumerate(df.iterrows()):
|
||
remaining = TIME_BUDGET_S - (time.monotonic() - T_START)
|
||
if remaining > 30:
|
||
text = generate(
|
||
tok,
|
||
model,
|
||
[
|
||
{"role": "system", "content": SYSTEM_PROMPT},
|
||
{
|
||
"role": "user",
|
||
"content": f"{r['context'].strip()}\n\n{r['query'].strip()}",
|
||
},
|
||
],
|
||
)
|
||
drafts.append(parse_answers(text, n_items_list[i]))
|
||
else:
|
||
drafts.append([""] * (n_items_list[i] or 0))
|
||
print(
|
||
f"draft {i + 1}/{len(df)} done (elapsed={time.monotonic() - T_START:.0f}s)",
|
||
flush=True,
|
||
)
|
||
|
||
# Safety net: a valid submission exists from here on.
|
||
write_submission(df, drafts)
|
||
print("draft submission.csv written", flush=True)
|
||
|
||
# Pass 2: verify-and-correct each draft, time permitting.
|
||
finals = list(drafts)
|
||
try:
|
||
for i, (_, r) in enumerate(df.iterrows()):
|
||
remaining = TIME_BUDGET_S - (time.monotonic() - T_START)
|
||
if remaining < MIN_S_PER_GEN:
|
||
print("time budget reached, keeping remaining drafts", flush=True)
|
||
break
|
||
if not any(drafts[i]):
|
||
continue # nothing to refine
|
||
draft_lines = "\n".join(
|
||
f"{j + 1}. {a}" for j, a in enumerate(drafts[i]) if a
|
||
)
|
||
text = generate(
|
||
tok,
|
||
model,
|
||
[
|
||
{"role": "system", "content": SYSTEM_PROMPT},
|
||
{
|
||
"role": "user",
|
||
"content": REFINE_PROMPT.format(
|
||
context=r["context"].strip(),
|
||
query=r["query"].strip(),
|
||
draft=draft_lines,
|
||
),
|
||
},
|
||
],
|
||
)
|
||
refined = parse_answers(text, n_items_list[i])
|
||
if refined and any(refined):
|
||
finals[i] = refined
|
||
print(
|
||
f"refine {i + 1}/{len(df)} done "
|
||
f"(changed={refined != drafts[i]}, "
|
||
f"elapsed={time.monotonic() - T_START:.0f}s)",
|
||
flush=True,
|
||
)
|
||
except Exception as e: # never let refinement cost us the submission
|
||
print(f"refine pass aborted: {e}", flush=True)
|
||
|
||
write_submission(df, finals)
|
||
print("wrote submission.csv", flush=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|