script.py

"""IOL-AI 2026 submission — DeepSeek-R1-Distill v4.

Force-close </think>, think 1024 + answer 384, parser v3, CoT 768 on format fail.
4-bit BitsAndBytes for T4 16GB / 30min.
Thinking tokens: <think> … </think>
Serves both R1-Qwen-7B and R1-Llama-8B Hub repos.
"""

import os

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

THINKING_BUDGET = 1024
ANSWER_CONTINUATION_TOKENS = 384
COT_MAX_NEW_TOKENS = 768

START_THINKING = "<think>"
END_THINKING = "</think>"

import json
import re

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

bnb = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
)

SYSTEM = """You solve International Linguistics Olympiad (IOL) problems from the data you are given.
You may see a task type you have never seen: follow the instruction and examples, and answer in the same form they use.

What to return by task type:
- translation: only the required form in the language the query asks for — do not add extra glosses or "form | meaning" unless asked
- fill_blanks: only the missing form for each blank — no extra glosses
- match_letters: ONLY the option letter (A, B, C, …), one letter per line — never copy option text, never arrows, never "A. word"
- text_to_num: the number in digits only
- num_to_text: the number written out in words, in the language asked
- any other type: exactly what the instruction asks for, nothing else

Answer in the language and form the query asks for. Do not add glosses, translations, or explanations unless the instruction requires them.

Output rules:
- Put answers ONLY after a line that says exactly: FINAL ANSWERS:
- Never put answers before that marker.
- One answer per line; exactly as many lines as items asked in the query.
- Bare answers only: no numbering, no quotes, no commentary, no repeating the question."""

SYSTEM_COT = (
    SYSTEM
    + "\n\nThink step by step about the rules in the examples and how they apply to the query, "
    "then write FINAL ANSWERS: and the answer lines."
)

_MD_PREFIX = r"(?:[#*_=\-\s`>]*)"
_MARKER = re.compile(
    rf"(?im)^{_MD_PREFIX}final\s+answers?{_MD_PREFIX}:?{_MD_PREFIX}\s*(.*)$"
)
_NUMBERING = re.compile(r"^\s*(?:\d+[.)]|[-*•])\s*")
_MD_WRAP = re.compile(r"^[*_`#\s]+|[*_`#\s]+$")
_TRAILING_LETTER = re.compile(
    r"(?:[–—\-]|→|->)\s*([A-Za-z])(?:\s*[.)]|)\s*$"
)
_LEADING_LETTER_OPT = re.compile(r"^([A-Za-z])\s*[.):\-–—]\s+\S")
_WORD_THEN_LETTER = re.compile(r"^.+\s([A-Za-z])\s*$")
_REFUSAL = re.compile(
    r"(?i)\b("
    r"i'?m sorry|i am sorry|i don'?t have|i cannot|i can'?t|"
    r"unable to|not able to|no reliable|cannot supply|can'?t supply|"
    r"as an ai|i apologize"
    r")\b"
)
_THINK_TAGS = re.compile(
    re.escape(START_THINKING) + r"|" + re.escape(END_THINKING)
)


def _clean_line(line: str) -> str:
    line = _NUMBERING.sub("", line).strip()
    line = _MD_WRAP.sub("", line).strip()
    line = line.replace("\u202f", " ").replace("\xa0", " ")
    return line.strip()


def after_thinking(text: str) -> str:
    if END_THINKING in text:
        text = text.rsplit(END_THINKING, 1)[-1]
    elif START_THINKING in text:
        text = ""
    return _THINK_TAGS.sub("", text)


def _as_option_letter(line: str) -> str | None:
    line = _clean_line(line)
    if not line:
        return None
    if len(line) == 1 and line.isalpha():
        return line.upper()
    m = _LEADING_LETTER_OPT.match(line)
    if m:
        return m.group(1).upper()
    m = _TRAILING_LETTER.search(line)
    if m:
        return m.group(1).upper()
    if len(line) <= 40:
        m = _WORD_THEN_LETTER.match(line)
        if m:
            return m.group(1).upper()
    return None


def _expand_line(line: str) -> list[str]:
    line = _clean_line(line)
    if not line:
        return []
    if len(line) == 1 and line.isalpha():
        return [line]
    if _LEADING_LETTER_OPT.match(line) or _TRAILING_LETTER.search(line):
        letter = _as_option_letter(line)
        if letter:
            return [letter]
    if len(line) <= 40 and _WORD_THEN_LETTER.match(line):
        letter = _as_option_letter(line)
        if letter:
            return [letter]
    if "|" in line:
        parts = [p.strip() for p in line.split("|") if p.strip()]
        if len(parts) >= 2:
            if len(parts) >= 4 and len(parts) % 2 == 0:
                left, right = parts[0::2], parts[1::2]
                if sum(" " in r for r in right) >= max(1, len(right) // 2):
                    return [_clean_line(x) for x in left if _clean_line(x)]
            if len(parts) == 2:
                a, b = parts
                if (" " in b and " " not in a) or (
                    len(b) > 2 * max(len(a), 1) and " " in b
                ):
                    return [_clean_line(a)] if _clean_line(a) else []
            return [_clean_line(p) for p in parts if _clean_line(p)]
    return [line]


def _dedupe_runaway(parts: list[str]) -> list[str]:
    if len(parts) < 6:
        return parts
    out: list[str] = []
    run = 0
    prev = None
    for p in parts:
        if p == prev:
            run += 1
            if run >= 4:
                break
        else:
            run = 1
            prev = p
        out.append(p)
    return out


def _lines_from_region(region: str, *, allow_all_lines: bool) -> list[str]:
    markers = list(_MARKER.finditer(region))
    if markers:
        last = markers[-1]
        after_parts: list[str] = []
        same = _clean_line(last.group(1) or "")
        if same:
            after_parts.extend(_expand_line(same))
        for line in region[last.end() :].splitlines():
            after_parts.extend(_expand_line(line))
        if after_parts:
            return _dedupe_runaway(after_parts)
        before_parts: list[str] = []
        for line in region[: last.start()].splitlines():
            before_parts.extend(_expand_line(line))
        if before_parts:
            return _dedupe_runaway(before_parts)
    parts: list[str] = []
    for line in region.splitlines():
        parts.extend(_expand_line(line))
    if not parts:
        return []
    if allow_all_lines:
        return _dedupe_runaway(parts)
    return [parts[-1]]


def parse_answers(
    raw: str,
    *,
    n_expected: int | None = None,
    task_type: str = "",
) -> list[str]:
    text = after_thinking(raw)
    answers = _lines_from_region(text, allow_all_lines=False)
    if task_type == "match_letters":
        coerced = []
        for a in answers:
            letter = _as_option_letter(a)
            coerced.append(letter if letter else a)
        answers = coerced
    if n_expected is not None and n_expected > 0 and len(answers) > n_expected:
        answers = answers[:n_expected]
    return answers


def _looks_like_alphabet_dump(answers: list[str]) -> bool:
    letters = [
        a.strip().upper()
        for a in answers
        if len(a.strip()) == 1 and a.strip().isalpha()
    ]
    if len(letters) < 15:
        return False
    seq = 0
    for i, L in enumerate(letters):
        if ord(L) == ord("A") + i:
            seq += 1
        else:
            break
    return seq >= 15


def _looks_like_refusal(answers: list[str]) -> bool:
    return bool(_REFUSAL.search(" ".join(answers)))


def has_usable_answer(
    answers: list[str],
    *,
    n_expected: int | None = None,
    task_type: str = "",
) -> bool:
    if not answers or not any(a.strip() for a in answers):
        return False
    if _looks_like_refusal(answers) or _looks_like_alphabet_dump(answers):
        return False
    if n_expected is not None and n_expected > 0 and len(answers) != n_expected:
        return False
    if task_type == "match_letters":
        letters = [a for a in answers if len(a) == 1 and a.isalpha()]
        if len(letters) < max(1, int(0.8 * len(answers))):
            return False
    return True


def _n_items_guess(query: str) -> int:
    nums = re.findall(r"(?m)^\s*(?:\(?\d+[.)]|\d+\))", query)
    return len(nums) if nums else 0


def _end_thinking_id(tok) -> int:
    end_id = tok.convert_tokens_to_ids(END_THINKING)
    if end_id is None or end_id == tok.unk_token_id:
        ids = tok.encode(END_THINKING, add_special_tokens=False)
        if len(ids) == 1:
            end_id = ids[0]
        elif len(ids) > 1:
            end_id = ids[-1]
    if end_id is None or end_id == tok.unk_token_id:
        raise RuntimeError(f"Tokenizer missing end-think token {END_THINKING!r}")
    return int(end_id)


def _build_prompt_ids(tok, system: str, user: str):
    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": user},
    ]
    return tok.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt"
    )


def _close_open_think(tok, prompt_ids, end_id: int):
    """R1 template opens <think>; for CoT, close it in the prompt so answers start immediately."""
    device = prompt_ids.device
    text = tok.decode(prompt_ids[0], skip_special_tokens=False)
    if text.rstrip().endswith(START_THINKING) or (
        START_THINKING in text and not text.rstrip().endswith(END_THINKING)
    ):
        # If the prompt ends inside an open think block, append end token
        if END_THINKING not in text.rsplit(START_THINKING, 1)[-1]:
            extra = torch.tensor([[end_id]], device=device, dtype=prompt_ids.dtype)
            return torch.cat([prompt_ids, extra], dim=1)
    return prompt_ids


@torch.inference_mode()
def generate_with_think_budget(model, tok, prompt_ids, end_id: int):
    device = next(model.parameters()).device
    prompt_ids = prompt_ids.to(device)
    prompt_len = prompt_ids.shape[-1]

    think_out = model.generate(
        prompt_ids,
        max_new_tokens=THINKING_BUDGET,
        do_sample=False,
        pad_token_id=tok.pad_token_id or tok.eos_token_id,
    )[0]
    gen_ids = think_out[prompt_len:].tolist()
    think_text = tok.decode(think_out[prompt_len:], skip_special_tokens=False)
    if end_id not in gen_ids and END_THINKING not in think_text:
        cont = torch.cat(
            [think_out, torch.tensor([end_id], device=device, dtype=think_out.dtype)]
        )
    else:
        cont = think_out

    full = model.generate(
        cont.unsqueeze(0),
        max_new_tokens=ANSWER_CONTINUATION_TOKENS,
        do_sample=False,
        pad_token_id=tok.pad_token_id or tok.eos_token_id,
    )[0]
    return tok.decode(full[prompt_len:], skip_special_tokens=False).strip()


@torch.inference_mode()
def generate_plain(model, tok, prompt_ids, max_new_tokens: int):
    device = next(model.parameters()).device
    prompt_ids = prompt_ids.to(device)
    prompt_len = prompt_ids.shape[-1]
    out = model.generate(
        prompt_ids,
        max_new_tokens=max_new_tokens,
        do_sample=False,
        pad_token_id=tok.pad_token_id or tok.eos_token_id,
    )[0]
    return tok.decode(out[prompt_len:], skip_special_tokens=False).strip()


tok = AutoTokenizer.from_pretrained(MODEL_ID)
end_id = _end_thinking_id(tok)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, quantization_config=bnb, device_map="auto"
).eval()

df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")

rows = []
for i, r in df.iterrows():
    n_guess = _n_items_guess(r["query"])
    task = str(r.get("task_type", "") or "")
    n_exp = n_guess or None
    user = f"{r['context'].strip()}\n\n{r['query'].strip()}"
    if n_guess:
        user += f"\n\n(Emit exactly {n_guess} answer line(s) after FINAL ANSWERS:.)"

    ids = _build_prompt_ids(tok, SYSTEM, user)
    text = generate_with_think_budget(model, tok, ids, end_id)
    answers = parse_answers(text, n_expected=n_exp, task_type=task)
    used_cot = False

    if not has_usable_answer(answers, n_expected=n_exp, task_type=task):
        used_cot = True
        cot_ids = _build_prompt_ids(tok, SYSTEM_COT, user)
        cot_ids = _close_open_think(tok, cot_ids.to(next(model.parameters()).device), end_id)
        cot_text = generate_plain(model, tok, cot_ids, COT_MAX_NEW_TOKENS)
        cot_answers = parse_answers(cot_text, n_expected=n_exp, task_type=task)
        if has_usable_answer(cot_answers, n_expected=n_exp, task_type=task):
            answers = cot_answers

    rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
    print(f"[{i + 1}/{len(df)}] {len(answers)} answers cot={used_cot}", flush=True)

pd.DataFrame(rows).to_csv("submission.csv", index=False)
print("wrote submission.csv", flush=True)