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()
|