228 lines
9.1 KiB
Python
228 lines
9.1 KiB
Python
|
"""IOL-AI 2026 submission script.
|
||
|
|
|
||
|
|
Runs inside the competition sandbox (T4, 16 GB, no internet, 30 min limit). The
|
||
|
|
submission repo is the working directory, so the model weights ship alongside this
|
||
|
|
file and load from ".". The hidden test set is mounted at /tmp/data/test.csv; we
|
||
|
|
write submission.csv with one row per problem id.
|
||
|
|
|
||
|
|
Model: Qwen2.5-14B-Instruct-AWQ (Apache-2.0), 4-bit AWQ so it fits the T4's 16 GB.
|
||
|
|
Single-pass prompt: one generation yields both the answers (FINAL ANSWERS block) and
|
||
|
|
a short human-readable explanation (EXPLANATION block, for the human-eval track).
|
||
|
|
"""
|
||
|
|
|
||
|
|
import os
|
||
|
|
|
||
|
|
# The sandbox has no network and the token is revoked before we run: force offline so
|
||
|
|
# transformers/hf_hub never try to reach the Hub (which would error).
|
||
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
||
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||
|
|
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
||
|
|
|
||
|
|
import csv
|
||
|
|
import json
|
||
|
|
import re
|
||
|
|
import time
|
||
|
|
|
||
|
|
# Heavy deps (pandas / torch / transformers) are imported lazily inside the functions
|
||
|
|
# that use them, so the pure-function helpers below can be unit-tested without a GPU
|
||
|
|
# or those packages installed.
|
||
|
|
|
||
|
|
# ----------------------------------------------------------------------------------
|
||
|
|
# Config / timing
|
||
|
|
# ----------------------------------------------------------------------------------
|
||
|
|
MODEL_ID = "." # weights ship in this repo
|
||
|
|
MAX_NEW_TOKENS = 1024 # room to reason; lower = faster/safer on the time limit
|
||
|
|
TEST_CSV = "/tmp/data/test.csv"
|
||
|
|
OUT_CSV = "submission.csv"
|
||
|
|
|
||
|
|
TIME_LIMIT = 30 * 60 # platform's TIME_LIMIT
|
||
|
|
SAFETY = 150 # reserve time to always write the CSV before we're killed
|
||
|
|
START = time.time()
|
||
|
|
DEADLINE = START + TIME_LIMIT - SAFETY
|
||
|
|
|
||
|
|
# ----------------------------------------------------------------------------------
|
||
|
|
# Prompt
|
||
|
|
# ----------------------------------------------------------------------------------
|
||
|
|
SYSTEM = (
|
||
|
|
"You solve International Linguistics Olympiad problems by reasoning from the data "
|
||
|
|
"you are given. Everything you need is in the problem; no outside knowledge of the "
|
||
|
|
"language is required. You may meet a task type you have never seen: read the "
|
||
|
|
"instruction and the examples, and answer in the same form they use. "
|
||
|
|
"Common task types and what to give -- "
|
||
|
|
"translation: the translated form only, in the language the task asks for; "
|
||
|
|
"fill_blanks: only the missing form for each blank; "
|
||
|
|
"match_letters: only the option letter (for example A, B, C); "
|
||
|
|
"text_to_num: the number in digits; "
|
||
|
|
"num_to_text: the number written out in words, in the language asked; "
|
||
|
|
"any other type: give exactly what the instruction asks, nothing else.\n\n"
|
||
|
|
"Reason step by step first. Then write a line that says exactly:\n"
|
||
|
|
"FINAL ANSWERS:\n"
|
||
|
|
"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. Give your best guess for "
|
||
|
|
"every item; never leave one blank.\n"
|
||
|
|
"Then write a line that says exactly:\n"
|
||
|
|
"EXPLANATION:\n"
|
||
|
|
"and below it a SHORT explanation for a human judge (a few bullet points or a small "
|
||
|
|
"table): the rule or pattern you found and the key evidence for your answers. Be "
|
||
|
|
"concise and structured -- do not repeat the full reasoning."
|
||
|
|
)
|
||
|
|
|
||
|
|
# ----------------------------------------------------------------------------------
|
||
|
|
# Pure-function helpers (unit-tested in test_logic.py, no GPU needed)
|
||
|
|
# ----------------------------------------------------------------------------------
|
||
|
|
_ANSWERS_MARKER = re.compile(r"(?im)^\s*final\s+answers?\s*:?\s*$")
|
||
|
|
_EXPL_MARKER = re.compile(r"(?im)^\s*explanation\s*:?\s*$")
|
||
|
|
# Numbered items at the start of a line: "1.", "2)", "17." (translation, number tasks).
|
||
|
|
_LINE_ITEM = re.compile(r"(?m)^\s*(\d+)\s*[.)]")
|
||
|
|
# Parenthesized blanks anywhere on a line: "... | (1) | ..." (fill_blanks tables).
|
||
|
|
_PAREN_ITEM = re.compile(r"\((\d+)\)")
|
||
|
|
_LEADING_NUM = re.compile(r"^\s*\(?\d+\)?\s*[.)]\s*")
|
||
|
|
|
||
|
|
|
||
|
|
def count_items(query):
|
||
|
|
"""Number of numbered items in the query. Handles both line-start numbering
|
||
|
|
("1." / "2)") and inline parenthesized blanks ("(1)", used by fill_blanks). Returns
|
||
|
|
0 when neither is found (e.g. matching queries whose items live in the context) so
|
||
|
|
the caller falls back to trusting the model's own answer count."""
|
||
|
|
q = query or ""
|
||
|
|
line_items = _LINE_ITEM.findall(q)
|
||
|
|
if line_items:
|
||
|
|
return len(line_items)
|
||
|
|
return len(_PAREN_ITEM.findall(q))
|
||
|
|
|
||
|
|
|
||
|
|
def _clean(line):
|
||
|
|
"""Strip a leading '1. '/'2) '/'(3) ' and surrounding quotes/whitespace."""
|
||
|
|
line = _LEADING_NUM.sub("", line.strip())
|
||
|
|
line = line.strip().strip('"').strip("'").strip()
|
||
|
|
return line
|
||
|
|
|
||
|
|
|
||
|
|
def parse_answers(text):
|
||
|
|
"""Answers = lines after the LAST 'FINAL ANSWERS:' marker, up to 'EXPLANATION:'.
|
||
|
|
Falls back to the non-empty lines of the whole text (minus the explanation) when
|
||
|
|
no marker is present."""
|
||
|
|
ans_markers = list(_ANSWERS_MARKER.finditer(text))
|
||
|
|
expl_markers = list(_EXPL_MARKER.finditer(text))
|
||
|
|
expl_start = expl_markers[-1].start() if expl_markers else len(text)
|
||
|
|
|
||
|
|
if ans_markers:
|
||
|
|
segment = text[ans_markers[-1].end():expl_start]
|
||
|
|
else:
|
||
|
|
# No marker: use everything before the explanation block as a best effort.
|
||
|
|
segment = text[:expl_start]
|
||
|
|
|
||
|
|
out = []
|
||
|
|
for line in segment.splitlines():
|
||
|
|
cleaned = _clean(line)
|
||
|
|
if cleaned:
|
||
|
|
out.append(cleaned)
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def parse_explanation(text):
|
||
|
|
"""Explanation = text after the LAST 'EXPLANATION:' marker ('' if absent)."""
|
||
|
|
expl_markers = list(_EXPL_MARKER.finditer(text))
|
||
|
|
if not expl_markers:
|
||
|
|
return ""
|
||
|
|
return text[expl_markers[-1].end():].strip()
|
||
|
|
|
||
|
|
|
||
|
|
def align(preds, n_items):
|
||
|
|
"""Line up predictions with the reference by position: exactly n_items entries,
|
||
|
|
padding short lists with '' and truncating long ones. When n_items is unknown
|
||
|
|
(0) we keep the model's list as-is."""
|
||
|
|
if n_items <= 0:
|
||
|
|
return preds
|
||
|
|
return (preds + [""] * n_items)[:n_items]
|
||
|
|
|
||
|
|
|
||
|
|
# ----------------------------------------------------------------------------------
|
||
|
|
# Model
|
||
|
|
# ----------------------------------------------------------------------------------
|
||
|
|
def load_model():
|
||
|
|
import torch
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
|
|
||
|
|
tok = AutoTokenizer.from_pretrained(MODEL_ID)
|
||
|
|
model = AutoModelForCausalLM.from_pretrained(
|
||
|
|
MODEL_ID, torch_dtype=torch.float16, device_map="auto",
|
||
|
|
).eval()
|
||
|
|
if tok.pad_token_id is None:
|
||
|
|
tok.pad_token = tok.eos_token
|
||
|
|
return tok, model
|
||
|
|
|
||
|
|
|
||
|
|
def generate(tok, model, context, query):
|
||
|
|
import torch
|
||
|
|
|
||
|
|
messages = [
|
||
|
|
{"role": "system", "content": SYSTEM},
|
||
|
|
{"role": "user", "content": f"{context.strip()}\n\n{query.strip()}"},
|
||
|
|
]
|
||
|
|
# Most-robust form across transformers versions: get a bare tensor and build the
|
||
|
|
# attention mask ourselves (avoids the version-dependent return_dict behaviour).
|
||
|
|
input_ids = tok.apply_chat_template(
|
||
|
|
messages, add_generation_prompt=True, return_tensors="pt",
|
||
|
|
).to(model.device)
|
||
|
|
attn = torch.ones_like(input_ids)
|
||
|
|
with torch.no_grad():
|
||
|
|
out = model.generate(
|
||
|
|
input_ids=input_ids,
|
||
|
|
attention_mask=attn,
|
||
|
|
max_new_tokens=MAX_NEW_TOKENS,
|
||
|
|
do_sample=False, # greedy -> reproducible
|
||
|
|
pad_token_id=tok.pad_token_id,
|
||
|
|
)
|
||
|
|
return tok.decode(out[0][input_ids.shape[-1]:], skip_special_tokens=True).strip()
|
||
|
|
|
||
|
|
|
||
|
|
# ----------------------------------------------------------------------------------
|
||
|
|
# Main
|
||
|
|
# ----------------------------------------------------------------------------------
|
||
|
|
def main():
|
||
|
|
import pandas as pd
|
||
|
|
|
||
|
|
df = pd.read_csv(TEST_CSV, dtype=str).fillna("")
|
||
|
|
print(f"loaded {len(df)} problems", flush=True)
|
||
|
|
|
||
|
|
tok, model = load_model()
|
||
|
|
print("model loaded", flush=True)
|
||
|
|
|
||
|
|
rows = []
|
||
|
|
for i, r in df.iterrows():
|
||
|
|
n_items = count_items(r["query"])
|
||
|
|
expl = ""
|
||
|
|
if time.time() > DEADLINE:
|
||
|
|
# Out of time: still emit a full, positionally-aligned (blank) row so a
|
||
|
|
# slow run never zeroes the whole submission via a hard timeout.
|
||
|
|
preds = []
|
||
|
|
print(f"{i + 1}/{len(df)} id={r['id']} SKIPPED (deadline)", flush=True)
|
||
|
|
else:
|
||
|
|
try:
|
||
|
|
text = generate(tok, model, r["context"], r["query"])
|
||
|
|
preds = parse_answers(text)
|
||
|
|
expl = parse_explanation(text)
|
||
|
|
except Exception as e: # never let one problem sink the whole run
|
||
|
|
preds = []
|
||
|
|
print(f"{i + 1}/{len(df)} id={r['id']} ERROR: {e}", flush=True)
|
||
|
|
else:
|
||
|
|
print(f"{i + 1}/{len(df)} id={r['id']} items={n_items} "
|
||
|
|
f"parsed={len(preds)}", flush=True)
|
||
|
|
|
||
|
|
preds = align(preds, n_items)
|
||
|
|
rows.append({
|
||
|
|
"id": r["id"],
|
||
|
|
"pred": json.dumps(preds, ensure_ascii=False),
|
||
|
|
"explanation": expl,
|
||
|
|
})
|
||
|
|
|
||
|
|
pd.DataFrame(rows, columns=["id", "pred", "explanation"]).to_csv(
|
||
|
|
OUT_CSV, index=False, quoting=csv.QUOTE_MINIMAL,
|
||
|
|
)
|
||
|
|
print(f"wrote {OUT_CSV} ({len(rows)} rows) in {time.time() - START:.0f}s", flush=True)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|