118 lines
4.2 KiB
Python
118 lines
4.2 KiB
Python
"""IOL-AI 2026 submission entrypoint — PLAIN baseline harness.
|
|
|
|
Faithful replica of the official How_to_submit.md baseline (the config the
|
|
organizers posted as "Baseline-Qwen2.5-14B-AWQ", public score 0.1227): a minimal
|
|
system prompt, the raw context+query, greedy decoding, 512 new tokens, and a
|
|
naive newline split of the output with NO forced item count. The only additions
|
|
are timeout-safe incremental writes (an all-blank submission.csv up front + a
|
|
periodic flush), which never change the output of a run that finishes in time —
|
|
they only guarantee a valid partial file if we were ever killed.
|
|
|
|
Runs on the platform's T4 (16 GB), offline, within the 30-minute limit. Reads
|
|
/tmp/data/test.csv, writes submission.csv (id,pred) to the working directory.
|
|
Model weights (Qwen2.5-14B-Instruct-AWQ) ship in this repo and load from ".".
|
|
"""
|
|
|
|
import os
|
|
import time
|
|
|
|
# Must be set BEFORE importing transformers: no network at run time.
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
|
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
|
|
|
import json
|
|
|
|
import pandas as pd
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
MODEL_ID = "." # weights are shipped in this repo (offline)
|
|
TEST_CSV = "/tmp/data/test.csv"
|
|
OUT_CSV = "submission.csv"
|
|
|
|
MAX_NEW_TOKENS = 512 # same as the official baseline
|
|
SAFETY_S = 25 * 60 # stop generating past this; buffer before the 30-min cap
|
|
WRITE_EVERY = 3 # flush submission.csv every N problems (timeout safety)
|
|
_START = time.monotonic()
|
|
|
|
# The official baseline's exact system prompt.
|
|
SYSTEM = (
|
|
"You solve International Linguistics Olympiad problems. Answer every numbered "
|
|
"item. Put each answer on its own line, in order, with no numbering and no "
|
|
"extra text."
|
|
)
|
|
|
|
|
|
def elapsed():
|
|
return time.monotonic() - _START
|
|
|
|
|
|
def load_model():
|
|
tok = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID,
|
|
torch_dtype=torch.float16, # T4 has no bfloat16
|
|
device_map="auto",
|
|
).eval()
|
|
if tok.pad_token_id is None:
|
|
tok.pad_token = tok.eos_token
|
|
return tok, model
|
|
|
|
|
|
def main():
|
|
df = pd.read_csv(TEST_CSV, dtype=str).fillna("")
|
|
total = len(df)
|
|
print(f"loaded {total} problems from {TEST_CSV}", flush=True)
|
|
|
|
# Write a COMPLETE all-blank submission.csv up front so a valid file with
|
|
# every id always exists, even if we are killed during load or generation.
|
|
records = [
|
|
{"id": df.iloc[i]["id"], "pred": json.dumps([], ensure_ascii=False)}
|
|
for i in range(total)
|
|
]
|
|
|
|
def flush():
|
|
pd.DataFrame(records, columns=["id", "pred"]).to_csv(OUT_CSV, index=False)
|
|
|
|
flush()
|
|
print(f"wrote blank {OUT_CSV} ({total} rows) at {elapsed():.0f}s", flush=True)
|
|
|
|
tok, model = load_model()
|
|
print(f"model loaded at {elapsed():.0f}s", flush=True)
|
|
|
|
for i in range(total):
|
|
if elapsed() > SAFETY_S:
|
|
print(f"[warn] time budget hit at problem {i}; leaving the rest blank", flush=True)
|
|
break
|
|
r = df.iloc[i]
|
|
messages = [
|
|
{"role": "system", "content": SYSTEM},
|
|
{"role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}"},
|
|
]
|
|
try:
|
|
ids = tok.apply_chat_template(
|
|
messages, add_generation_prompt=True, return_tensors="pt",
|
|
).to(model.device)
|
|
with torch.no_grad():
|
|
out = model.generate(
|
|
ids, max_new_tokens=MAX_NEW_TOKENS,
|
|
do_sample=False, pad_token_id=tok.pad_token_id,
|
|
)
|
|
text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
|
|
answers = [ln.strip() for ln in text.splitlines() if ln.strip()]
|
|
except Exception as exc:
|
|
print(f"[warn] id={r['id']} failed: {exc!r}", flush=True)
|
|
answers = []
|
|
records[i]["pred"] = json.dumps(answers, ensure_ascii=False)
|
|
if (i + 1) % WRITE_EVERY == 0:
|
|
flush()
|
|
print(f"{i + 1}/{total} done at {elapsed():.0f}s", flush=True)
|
|
|
|
flush()
|
|
print(f"final {OUT_CSV} ({total} rows) at {elapsed():.0f}s", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|