script.py

"""IOL-AI 2026 submission script.

Runs fully offline against /tmp/data/test.csv and writes submission.csv with
columns: id, pred (JSON list of answer strings), explanation (omitted).
"""
import os

os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

import time

T_START = time.monotonic()

import json
import re
from collections import Counter

import pandas as pd
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# --- Configuration -----------------------------------------------------

# Production always loads from "." (weights shipped alongside this script in
# the HF model repo). This constant only documents which Hub model those
# weights came from, for keeping the Colab test notebook in sync.
MODEL_ID_FOR_LOCAL_TEST = "Qwen/Qwen2.5-7B-Instruct-AWQ"
# Fallback if Colab timing shows 7B has slack to spare: swap MODEL_ID_FOR_LOCAL_TEST
# to "Qwen/Qwen2.5-14B-Instruct-AWQ" and re-run the notebook before re-uploading.

MODEL_PATH = "."
DATA_PATH = "/tmp/data/test.csv"
OUTPUT_PATH = "submission.csv"

TIME_BUDGET_SEC = 27 * 60  # 3-min safety margin under the 30-min hard limit
BATCH_SIZE = 4
MIN_NEW_TOKENS = 128
MAX_NEW_TOKENS = 1024

# Self-consistency: sample k >= 1 completions per row and majority-vote per
# answer position when the time budget has slack. k=1 is plain greedy decode
# (current default, deterministic). See adaptive_k().
MIN_K = 1
MAX_K = 1
SAMPLE_TEMPERATURE = 0.7

# Which SYSTEM prompt variant to use — "full_cot" or "plain_io". Published
# evidence on whether heavy chain-of-thought helps or hurts a 7B-class model
# on this task type is mixed (modeLing, arXiv:2406.17038), so both variants
# are kept and should be A/B tested on the Colab harness rather than assumed.
PROMPT_VARIANT = "plain_io"

# Worked example shared by both prompt variants: a made-up toy language (not
# a real one) demonstrating (a) the answer must be a word-form in the target
# language, never an English gloss, and (b) subject/object roles must be
# re-derived per query, not copied from a similar-looking example. Directly
# targets the two failure modes observed in real Colab runs (language
# confusion — arXiv:2406.20052 shows few-shot examples largely eliminate it;
# and swapped argument roles).
_WORKED_EXAMPLE = """Worked example (a made-up toy language, unrelated to your actual task — the
words and rules below apply ONLY to this example; never reuse them):

CONTEXT:
mota = "dog"        mota-ta = "dog" (as object)
suno = "cat"         suno-ta = "cat" (as object)
kire = "sees"
mota suno-ta kire = "the dog sees the cat"
suno mota-ta kire = "the cat sees the dog"

QUERY:
Translate into the made-up language: "the cat sees the dog"

IDENTIFY PATTERNS: word order is SUBJECT OBJECT-ta VERB; subject is unmarked,
object takes suffix "-ta", verb is always last.
RULE TABLE:
- subject noun -> bare noun, placed first
- object noun -> noun + "-ta", placed second
- verb -> placed last, unchanged
TEST: "mota suno-ta kire" = dog(subj) cat-ta(obj) sees = "the dog sees the
cat" — matches. "suno mota-ta kire" = cat(subj) dog-ta(obj) sees = "the cat
sees the dog" — matches.
APPLY: query is "the cat sees the dog" -> subject=cat=suno,
object=dog=mota-ta, verb=kire.
VERIFY: answer is a word-form in the made-up language, not an English gloss.
Roles checked: cat is subject (first, unmarked), dog is object (second,
"-ta" suffix) — matches "the cat sees the dog", not reversed.

FINAL ANSWERS:
suno mota-ta kire

--- end worked example. Now solve the actual task below using ONLY the data
in its own CONTEXT and QUERY — never reuse the words or rules above. ---"""

_ANSWER_FORMAT_RULES = """Answer formatting rules by task_type:
- translation: give the full translated phrase or sentence for each numbered item,
  written as a natural fluent sentence exactly like the answer column in the given
  examples — never as a morpheme-by-morpheme gloss with parentheses like
  "you(shuddered)" or "we(spat(on him))".
- fill_blanks: give only the missing word(s)/form for each numbered blank.
- match_letters: give the matching letter or number for each item.
- text_to_num: give only the numeral for each item.
- num_to_text: give only the word(s) for each item.

Output a line that says exactly:
FINAL ANSWERS:
followed by one answer per line, in the same order and count as the numbered
items in the query, with no numbering, quotes, or extra commentary — just the
bare answer text for each line."""

SYSTEM_FULL_COT = f"""You are an expert linguist solving International Linguistics Olympiad problems.
You are given a self-contained set of data in a language you have never seen before,
plus a query asking you to apply what that data teaches you. You must reason only from
the data given — never from memorized knowledge of real-world languages.

Work in stages, showing your work:
1. IDENTIFY PATTERNS: look for recurring forms, affixes, word order, sound
   correspondences, or structural regularities in the given data.
2. RULE TABLE: write your hypotheses as short "TRIGGER -> TRANSFORMATION"
   bullet lines (e.g. "subject noun -> bare noun, placed first"), not prose —
   this keeps rules mechanical and easy to re-check, rather than vague
   descriptions that drift away from the actual data.
3. TEST HYPOTHESES: check each rule-table line against every example in the
   given data. Discard or refine any rule that doesn't hold up on every
   example.
4. APPLY: use only the rules that survived testing to answer the query.
5. VERIFY: before writing FINAL ANSWERS, re-check every draft answer against
   two common mistakes:
   - LANGUAGE: look at the answer column in the given examples, not the
     instructions, to see what language/form your answer must be in. If the
     examples' answers are word-forms in the unfamiliar language, your answer
     must also be a word-form in that language — never substitute an English
     gloss or description of the meaning, even if you're unsure of the exact
     form; give your best-guess constructed form instead.
   - ROLES: for anything involving "who did what to whom" (subject, object,
     possessor, giver/receiver), re-read the query and re-derive which
     argument fills which role from scratch. Do not assume the same word
     order or role assignment as a similar-looking training example — verify
     it against the actual affixes/markers in that example.

{_WORKED_EXAMPLE}

{_ANSWER_FORMAT_RULES}"""

SYSTEM_PLAIN_IO = f"""You are an expert linguist solving International Linguistics Olympiad problems.
You are given a self-contained set of data in a language you have never seen before,
plus a query asking you to apply what that data teaches you. Reason only from the
data given — never from memorized knowledge of real-world languages. Give your
best-guess answer in the same language/form as the examples' answers — never
substitute an English gloss. Double-check which argument is subject vs object
before answering.

{_WORKED_EXAMPLE}

{_ANSWER_FORMAT_RULES}"""

_SYSTEM_VARIANTS = {"full_cot": SYSTEM_FULL_COT, "plain_io": SYSTEM_PLAIN_IO}
SYSTEM = _SYSTEM_VARIANTS[PROMPT_VARIANT]

USER_TEMPLATE = """Task type: {task_type}
Eval type: {eval_type}
Working language -> Task language: {work_lang} -> {task_lang}

CONTEXT:
{context}

QUERY:
{query}"""


# --- Query parsing helpers ---------------------------------------------

_RANGE_RE = re.compile(r"\((\d+)\s*[-–]\s*(\d+)\)")
_LEADING_NUM_RE = re.compile(r"^\s*(\d+)[.)]", re.MULTILINE)
_INLINE_PAREN_NUM_RE = re.compile(r"\((\d+)\)")


def _is_contiguous(nums: list[int]) -> bool:
    if not nums:
        return False
    uniq = sorted(set(nums))
    return uniq == list(range(uniq[0], uniq[-1] + 1))


def count_expected_items(query: str) -> int:
    """Count how many answers the query expects, robust to numbering style."""
    range_match = _RANGE_RE.search(query)
    if range_match:
        start, end = int(range_match.group(1)), int(range_match.group(2))
        if end >= start:
            return end - start + 1

    leading_nums = [int(n) for n in _LEADING_NUM_RE.findall(query)]
    if _is_contiguous(leading_nums):
        return len(set(leading_nums))

    inline_nums = [int(n) for n in _INLINE_PAREN_NUM_RE.findall(query)]
    if _is_contiguous(inline_nums):
        return len(set(inline_nums))

    non_empty_lines = [ln for ln in query.splitlines() if ln.strip()]
    return max(len(non_empty_lines) - 1, 1)


# --- Answer parsing helpers ----------------------------------------------

_MARKER_RE = re.compile(r"FINAL ANSWERS:\s*", re.IGNORECASE)
_LINE_NUM_PREFIX_RE = re.compile(r"^\s*\(?\d+\)?[.):]?\s*")
_QUOTE_STRIP_RE = re.compile(r'^["\'](.*)["\']$')


def _strip_line(line: str) -> str:
    line = line.strip()
    line = _LINE_NUM_PREFIX_RE.sub("", line, count=1)
    m = _QUOTE_STRIP_RE.match(line)
    if m:
        line = m.group(1)
    return line.strip()


def parse_answers(text: str, expected_n: int) -> list[str]:
    """Extract exactly expected_n answer strings from raw model output."""
    parts = _MARKER_RE.split(text, maxsplit=1)
    tail = parts[1] if len(parts) > 1 else text

    lines = [_strip_line(ln) for ln in tail.splitlines()]
    answers = [ln for ln in lines if ln]

    # Model sometimes emits all answers on one "|"-delimited line instead of
    # one per line; only fall back to splitting on "|" when the plain
    # line-split came up short, so genuine single-line answers containing "|"
    # aren't mangled.
    if len(answers) < expected_n and any("|" in a for a in answers):
        expanded = []
        for a in answers:
            if "|" in a:
                expanded.extend(_strip_line(p) for p in a.split("|"))
            else:
                expanded.append(a)
        expanded = [a for a in expanded if a]
        if len(expanded) > len(answers):
            answers = expanded

    if len(answers) < expected_n:
        answers = answers + [""] * (expected_n - len(answers))
    elif len(answers) > expected_n:
        answers = answers[:expected_n]

    return answers


# --- Time-budget-aware batched generation --------------------------------

def elapsed() -> float:
    return time.monotonic() - T_START


def remaining_budget() -> float:
    return TIME_BUDGET_SEC - elapsed()


def build_prompt(tokenizer, row) -> str:
    user_msg = USER_TEMPLATE.format(
        task_type=row["task_type"],
        eval_type=row["eval_type"],
        work_lang=row.get("work_lang", ""),
        task_lang=row.get("task_lang", ""),
        context=row["context"],
        query=row["query"],
    )
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": user_msg},
    ]
    return tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )


def adaptive_max_new_tokens(rows_left: int) -> int:
    if rows_left <= 0:
        return MIN_NEW_TOKENS
    per_row_budget = remaining_budget() / rows_left
    # Rough T4/AWQ-7B throughput assumption used only to size generation length;
    # errs conservative so we degrade gracefully rather than time out.
    tokens_per_sec_estimate = 20.0
    budget_tokens = int(per_row_budget * tokens_per_sec_estimate)
    return max(MIN_NEW_TOKENS, min(MAX_NEW_TOKENS, budget_tokens))


def adaptive_k(rows_left: int, max_new_tokens: int) -> int:
    """How many sampled completions per row the remaining budget affords.

    Only returns >1 once max_new_tokens has already saturated at
    MAX_NEW_TOKENS (i.e. there's more budget than a single generation pass
    needs) — never trades away reasoning length for extra samples.
    """
    if rows_left <= 0 or max_new_tokens <= 0:
        return MIN_K
    per_row_budget = remaining_budget() / rows_left
    tokens_per_sec_estimate = 20.0
    budget_tokens = int(per_row_budget * tokens_per_sec_estimate)
    k = budget_tokens // max_new_tokens
    return max(MIN_K, min(MAX_K, k))


def majority_vote_answers(sampled_answer_lists: list[list[str]]) -> list[str]:
    """Collapse k parsed answer lists (same length, one per sample) into one
    by majority vote per position. Prefers non-empty answers over blanks when
    both appear, since a wrong guess still earns partial chrF credit and a
    blank never does. Ties break toward the first sample's value.
    """
    if not sampled_answer_lists:
        return []
    n = len(sampled_answer_lists[0])
    result = []
    for i in range(n):
        votes = [sample[i] for sample in sampled_answer_lists if i < len(sample)]
        pool = [v for v in votes if v] or votes
        if not pool:
            result.append("")
            continue
        counts = Counter(pool)
        max_count = max(counts.values())
        winner = next(v for v in pool if counts[v] == max_count)
        result.append(winner)
    return result


def pad_row(expected_n: int) -> str:
    return json.dumps([""] * expected_n)


def main() -> None:
    df = pd.read_csv(DATA_PATH, dtype=str).fillna("")

    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH, torch_dtype=torch.float16, device_map="auto"
    )
    model.eval()

    results: dict[str, str] = {}
    rows = df.to_dict("records")

    i = 0
    while i < len(rows):
        rows_left = len(rows) - i

        if remaining_budget() < 30:  # not enough time left to safely attempt a batch
            for row in rows[i:]:
                expected_n = count_expected_items(row["query"])
                results[row["id"]] = pad_row(expected_n)
            break

        batch = rows[i : i + BATCH_SIZE]
        expected_counts = [count_expected_items(r["query"]) for r in batch]
        prompts = [build_prompt(tokenizer, r) for r in batch]

        max_new_tokens = adaptive_max_new_tokens(rows_left)
        k = adaptive_k(rows_left, max_new_tokens)

        inputs = tokenizer(
            prompts, return_tensors="pt", padding=True, truncation=True
        ).to(model.device)

        sampled_decoded: list[list[str]] = []
        for sample_idx in range(k):
            gen_kwargs = dict(
                max_new_tokens=max_new_tokens, pad_token_id=tokenizer.pad_token_id
            )
            if k > 1:
                gen_kwargs.update(do_sample=True, temperature=SAMPLE_TEMPERATURE, top_p=0.9)
            else:
                gen_kwargs.update(do_sample=False)
            try:
                with torch.no_grad():
                    output_ids = model.generate(**inputs, **gen_kwargs)
                input_len = inputs["input_ids"].shape[1]
                decoded = tokenizer.batch_decode(
                    output_ids[:, input_len:], skip_special_tokens=True
                )
            except Exception:
                decoded = [""] * len(batch)
            sampled_decoded.append(decoded)

        for row_idx, (row, expected_n) in enumerate(zip(batch, expected_counts)):
            sampled_answers = [
                parse_answers(sampled_decoded[s][row_idx], expected_n) for s in range(k)
            ]
            answers = majority_vote_answers(sampled_answers)
            results[row["id"]] = json.dumps(answers)

        i += len(batch)

    out_df = pd.DataFrame(
        {"id": df["id"], "pred": [results[i] for i in df["id"]]}
    )
    out_df.to_csv(OUTPUT_PATH, index=False)


if __name__ == "__main__":
    main()