245 lines
9.8 KiB
Python
245 lines
9.8 KiB
Python
"""IOL-AI 2026 v3: Qwen2.5-7B-Instruct in 4-bit, per-item prompting + reasoning.
|
||
|
||
Runs inside the competition sandbox (no internet). The submission repo is the
|
||
working directory; model weights are shipped in it and loaded from ".".
|
||
Reads /tmp/data/test.csv, writes submission.csv (columns: id, pred).
|
||
|
||
v3 over v2: Qwen2.5-7B-Instruct instead of 1.5B, quantized to 4-bit (NF4) at
|
||
load time with bitsandbytes. (AWQ was a dead end: autoawq pins an exact torch
|
||
version and the sandbox's package index serves no torch wheels at all, so the
|
||
preinstalled torch must be used as-is.)
|
||
|
||
v2 over the baseline:
|
||
- One generation per numbered item (not one per problem), so a misparsed line
|
||
can't shift every later answer.
|
||
- The model reasons first, then emits "FINAL: <answer>"; only that is kept.
|
||
- Task-type-specific answer-format instructions (exact match is unforgiving).
|
||
- A one-shot worked example demonstrating the reasoning + FINAL format.
|
||
- A global time budget that shrinks reasoning space as the deadline nears.
|
||
"""
|
||
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
|
||
START_TIME = time.time()
|
||
|
||
# IOL_SKIP_PIP=1 skips dependency install for local testing (where torch etc.
|
||
# are already set up); the sandbox always installs.
|
||
if os.environ.get("IOL_SKIP_PIP") != "1":
|
||
# Do NOT list torch here: the sandbox's package index has no torch wheels
|
||
# (any resolver attempt to fetch one fails), so the preinstalled torch
|
||
# must satisfy every requirement as-is.
|
||
subprocess.run(
|
||
[sys.executable, "-m", "pip", "install", "-q",
|
||
"transformers>=4.43", "accelerate>=0.30", "pandas"],
|
||
check=True,
|
||
)
|
||
# bitsandbytes installed separately and non-fatally: if the sandbox index
|
||
# lacks it, the script falls back to fp16 loading below.
|
||
subprocess.run(
|
||
[sys.executable, "-m", "pip", "install", "-q", "bitsandbytes"],
|
||
check=False,
|
||
)
|
||
|
||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||
|
||
# IOL_TEST_CSV / IOL_MODEL_ID are only for local testing; in the sandbox the
|
||
# defaults are used.
|
||
TEST_CSV = os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv")
|
||
MODEL_ID = os.environ.get("IOL_MODEL_ID", ".")
|
||
OUT_CSV = os.environ.get("IOL_OUT_CSV", "submission.csv")
|
||
# Stop reasoning generously before the sandbox's 30-minute hard kill; the
|
||
# budget also has to cover pip install and model load (already elapsed).
|
||
TIME_BUDGET_S = int(os.environ.get("IOL_TIME_BUDGET_S", 25 * 60))
|
||
|
||
import json
|
||
import re
|
||
|
||
# Progress prints contain IPA characters; don't let a non-UTF-8 console
|
||
# (Windows local testing) crash the run.
|
||
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
|
||
|
||
import pandas as pd
|
||
import torch
|
||
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
|
||
|
||
SYSTEM_PROMPT = (
|
||
"You solve International Linguistics Olympiad problems: given a few "
|
||
"examples from an unfamiliar language, you deduce its rules and apply "
|
||
"them. Think step by step, briefly. Then give your answer on a new line "
|
||
"starting with 'FINAL:'. The FINAL line must contain only the answer "
|
||
"itself: no explanation, no quotes, no numbering."
|
||
)
|
||
|
||
# One worked example (Hakhun, from the competition page) demonstrating the
|
||
# reasoning style and the FINAL format.
|
||
ONE_SHOT_USER = """DATA
|
||
Here are sentences in Hakhun with their English translations:
|
||
ŋa ka kɤ ne | Do I go?
|
||
nɤ ʒip tuʔ ne | Did you sleep?
|
||
ati ka la | He goes.
|
||
ŋa ʒip ku | I sleep.
|
||
|
||
TASK
|
||
Translate into English.
|
||
Answer format: only the translation, nothing else.
|
||
|
||
ITEM
|
||
nɤ ʒip ku ne"""
|
||
|
||
ONE_SHOT_ASSISTANT = """From the data: ŋa = I, nɤ = you, ati = he. ka = go, ʒip = sleep. \
|
||
Sentence-final ne marks a question. The verb suffix differs by subject and tense: \
|
||
kɤ/ku are present forms (kɤ with ŋa+ka, ku with ŋa+ʒip), tuʔ is past, la is 3rd person present. \
|
||
Here: nɤ (you) + ʒip (sleep) + ku (present) + ne (question) = present-tense question about you sleeping.
|
||
FINAL: Do you sleep?"""
|
||
|
||
# Answer-format hints per task_type; unknown types get the generic hint.
|
||
TASK_HINTS = {
|
||
"translation": "only the translation, nothing else",
|
||
"fill_blanks": "only the word or words that fill the blank",
|
||
"match_letters": "only the letter or letters of the matching choice",
|
||
"text_to_num": "only the number in digits",
|
||
"num_to_text": "only the number written out in the problem language",
|
||
}
|
||
GENERIC_HINT = "only the answer itself, as short as possible"
|
||
|
||
# A numbered item looks like "17. ..." or "17) ..." at the start of a line.
|
||
ITEM_RE = re.compile(r"^\s*\(?(\d+)\s*[.)]\s*", re.MULTILINE)
|
||
# Models drift to variants like "Final Answer:" or "Answer:"; accept them all.
|
||
FINAL_RE = re.compile(r"(?:FINAL(?:\s+ANSWER)?|ANSWER)\s*:\s*(.*)", re.IGNORECASE)
|
||
|
||
|
||
def split_query(query: str):
|
||
"""Split the query into (instruction, [(number, item_text), ...])."""
|
||
matches = list(ITEM_RE.finditer(query))
|
||
if not matches:
|
||
return query.strip(), [("1", query.strip())]
|
||
instruction = query[: matches[0].start()].strip()
|
||
items = []
|
||
for i, m in enumerate(matches):
|
||
end = matches[i + 1].start() if i + 1 < len(matches) else len(query)
|
||
items.append((m.group(1), query[m.end(): end].strip()))
|
||
return instruction, items
|
||
|
||
|
||
QUOTED_RE = re.compile(r"[\"“”'']([^\"“”'']+)[\"“”'']")
|
||
|
||
|
||
def extract_final(text: str) -> str:
|
||
"""Take the last FINAL: line; fall back to the last non-empty line."""
|
||
finals = FINAL_RE.findall(text)
|
||
if finals:
|
||
return finals[-1].strip().strip('"').strip()
|
||
lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
|
||
if not lines:
|
||
return ""
|
||
last = lines[-1]
|
||
# Without a FINAL marker the last line is often prose like
|
||
# 'The sentence "X" translates to "Y".' — take the last quoted span.
|
||
quoted = QUOTED_RE.findall(last)
|
||
if quoted and len(last) > 2 * len(quoted[-1]):
|
||
return quoted[-1].strip()
|
||
return last
|
||
|
||
|
||
def main():
|
||
tok = AutoTokenizer.from_pretrained(MODEL_ID)
|
||
# NF4 4-bit quantization: ~5 GB on the T4, fp16 compute. Well-supported
|
||
# on T4 (no bf16 needed) and avoids the fp16-overflow garbage seen when
|
||
# running full-precision Qwen2.5 in fp16.
|
||
try:
|
||
quant = BitsAndBytesConfig(
|
||
load_in_4bit=True,
|
||
bnb_4bit_quant_type="nf4",
|
||
bnb_4bit_compute_dtype=torch.float16,
|
||
bnb_4bit_use_double_quant=True,
|
||
)
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
MODEL_ID, quantization_config=quant, device_map="auto"
|
||
).eval()
|
||
print("Loaded 4-bit (bitsandbytes NF4).", flush=True)
|
||
except Exception as e:
|
||
# Last resort if bitsandbytes is unavailable: fp16 barely fits the
|
||
# 16 GB T4 for 7B and risks Qwen2.5 fp16 overflow, but a degraded
|
||
# run beats no submission.csv at all.
|
||
print(f"4-bit load failed ({e!r}); falling back to fp16.", flush=True)
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
MODEL_ID, torch_dtype=torch.float16, device_map="auto"
|
||
).eval()
|
||
|
||
df = pd.read_csv(TEST_CSV, dtype=str).fillna("")
|
||
|
||
# Count items up front so the time budget can be split evenly.
|
||
parsed = []
|
||
total_items = 0
|
||
for _, r in df.iterrows():
|
||
instruction, items = split_query(r["query"])
|
||
parsed.append((r, instruction, items))
|
||
total_items += len(items)
|
||
|
||
def generate(messages, max_new_tokens):
|
||
prompt = tok.apply_chat_template(
|
||
messages, add_generation_prompt=True, tokenize=False
|
||
)
|
||
enc = tok(prompt, return_tensors="pt").to(model.device)
|
||
with torch.no_grad():
|
||
out = model.generate(
|
||
**enc,
|
||
max_new_tokens=max_new_tokens,
|
||
do_sample=False,
|
||
temperature=None,
|
||
top_p=None,
|
||
top_k=None,
|
||
pad_token_id=tok.eos_token_id,
|
||
)
|
||
return tok.decode(
|
||
out[0][enc["input_ids"].shape[-1]:], skip_special_tokens=True
|
||
).strip()
|
||
|
||
rows = []
|
||
done_items = 0
|
||
for r, instruction, items in parsed:
|
||
hint = TASK_HINTS.get(r["task_type"], GENERIC_HINT)
|
||
answers = []
|
||
for num, item_text in items:
|
||
remaining_s = TIME_BUDGET_S - (time.time() - START_TIME)
|
||
remaining_items = total_items - done_items
|
||
# Even share of the remaining budget, assuming ~10 tokens/s for
|
||
# NF4 7B on a T4; clamp to a sane range. If time has essentially
|
||
# run out, emit a 32-token answer-only guess.
|
||
share_s = max(remaining_s, 1) / max(remaining_items, 1)
|
||
max_new = int(min(384, max(32, share_s * 10)))
|
||
|
||
user_msg = (
|
||
f"DATA\n{r['context'].strip()}\n\n"
|
||
f"TASK\n{instruction or r['query'].strip()}\n"
|
||
f"Answer format: {hint}.\n\n"
|
||
f"ITEM\n{item_text}"
|
||
)
|
||
messages = [
|
||
{"role": "system", "content": SYSTEM_PROMPT},
|
||
{"role": "user", "content": ONE_SHOT_USER},
|
||
{"role": "assistant", "content": ONE_SHOT_ASSISTANT},
|
||
{"role": "user", "content": user_msg},
|
||
]
|
||
text = generate(messages, max_new)
|
||
answer = extract_final(text)
|
||
answers.append(answer)
|
||
done_items += 1
|
||
elapsed = time.time() - START_TIME
|
||
print(f"[{r['id']}] item {num}: {answer!r} "
|
||
f"(max_new={max_new}, t={elapsed:.0f}s)", flush=True)
|
||
|
||
rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
|
||
|
||
pd.DataFrame(rows).to_csv(OUT_CSV, index=False)
|
||
print(f"Wrote {len(rows)} rows to {OUT_CSV} "
|
||
f"in {time.time() - START_TIME:.0f}s")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|