script.py

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