118 lines
3.8 KiB
Python
118 lines
3.8 KiB
Python
"""IOL-AI 2026 submission entrypoint.
|
|
|
|
Runs on the platform's T4 (16 GB), offline, within a 30-minute limit. Reads the
|
|
hidden test set from /tmp/data/test.csv and writes submission.csv (id,pred) to
|
|
the working directory. Model weights ship inside this repo and load from ".".
|
|
|
|
Baseline model: Qwen/Qwen2.5-1.5B-Instruct. The score lever here is the harness
|
|
(task-aware prompts + robust parsing) in iol_harness.py, not the model size.
|
|
"""
|
|
|
|
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
|
|
|
|
import iol_harness as H
|
|
|
|
MODEL_ID = "." # weights are shipped in this repo (offline)
|
|
TEST_CSV = "/tmp/data/test.csv"
|
|
OUT_CSV = "submission.csv"
|
|
|
|
TIME_LIMIT_S = 30 * 60
|
|
SAFETY_S = 26 * 60 # stop generating past this; still write what we have
|
|
_START = time.monotonic()
|
|
|
|
# Optional per-repo harness config (reasoning-model repos ship one to turn on CoT).
|
|
CFG = {}
|
|
if os.path.exists("harness_cfg.json"):
|
|
try:
|
|
CFG = json.load(open("harness_cfg.json"))
|
|
except Exception:
|
|
CFG = {}
|
|
COT = bool(CFG.get("cot", False))
|
|
MAX_CAP = int(CFG.get("max_new_cap", 1024))
|
|
|
|
|
|
def elapsed():
|
|
return time.monotonic() - _START
|
|
|
|
|
|
def load_model():
|
|
tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID,
|
|
torch_dtype=torch.float16, # T4 has no bfloat16
|
|
device_map="auto",
|
|
trust_remote_code=True,
|
|
).eval()
|
|
if tok.pad_token_id is None:
|
|
tok.pad_token = tok.eos_token
|
|
return tok, model
|
|
|
|
|
|
def generate(tok, model, messages, max_new_tokens):
|
|
inputs = tok.apply_chat_template(
|
|
messages,
|
|
add_generation_prompt=True,
|
|
return_tensors="pt",
|
|
).to(model.device)
|
|
with torch.no_grad():
|
|
out = model.generate(
|
|
inputs,
|
|
max_new_tokens=max_new_tokens,
|
|
do_sample=False, # greedy => reproducible
|
|
num_beams=1,
|
|
pad_token_id=tok.pad_token_id,
|
|
)
|
|
return tok.decode(out[0][inputs.shape[-1] :], skip_special_tokens=True).strip()
|
|
|
|
|
|
def main():
|
|
df = pd.read_csv(TEST_CSV, dtype=str).fillna("")
|
|
total = len(df)
|
|
print(f"loaded {total} problems from {TEST_CSV}", flush=True)
|
|
|
|
tok, model = load_model()
|
|
print(f"model loaded at {elapsed():.0f}s", flush=True)
|
|
|
|
records = []
|
|
for i, row in df.iterrows():
|
|
rid = row["id"]
|
|
context = row.get("context", "")
|
|
query = row.get("query", "")
|
|
task_type = row.get("task_type", "")
|
|
n = H.count_items(context, query, task_type)
|
|
|
|
preds = [""] * n # safe default so every item has a slot
|
|
if elapsed() < SAFETY_S:
|
|
try:
|
|
messages = H.build_messages(context, query, task_type, n, cot=COT)
|
|
budget = H.max_new_tokens_for(task_type, n, cot=COT, cap=MAX_CAP)
|
|
text = generate(tok, model, messages, budget)
|
|
preds = H.parse_answers(text, n)
|
|
except Exception as exc: # never let one row sink the whole run
|
|
print(f"[warn] id={rid} failed: {exc!r}", flush=True)
|
|
else:
|
|
print(f"[warn] time budget hit; blanking id={rid}", flush=True)
|
|
|
|
records.append({"id": rid, "pred": json.dumps(preds, ensure_ascii=False)})
|
|
if (i + 1) % 5 == 0 or (i + 1) == total:
|
|
print(f"{i + 1}/{total} done ({elapsed():.0f}s)", flush=True)
|
|
|
|
pd.DataFrame(records, columns=["id", "pred"]).to_csv(OUT_CSV, index=False)
|
|
print(f"wrote {OUT_CSV} ({len(records)} rows) at {elapsed():.0f}s", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|