script.py

import os
import subprocess
import sys


def _install_bundled_deps() -> None:
    """Install transformers from bundled wheels (eval sandbox has no PyPI access)."""
    wheels_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "wheels")
    if not os.path.isdir(wheels_dir):
        return
    subprocess.run(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "-q",
            "--no-index",
            f"--find-links={wheels_dir}",
            "transformers==4.56.2",
        ],
        check=True,
    )


_install_bundled_deps()

import re
import csv
import json
import shutil
import tempfile
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# The repo is the working directory at run time, and there is no network.
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
MODEL_ID = "."
MAX_NEW_TOKENS = 2000
TEMPERATURE = 0.8
TOP_P = 0.95
MAX_ATTEMPTS = 5


def load_tokenizer(model_id: str = "."):
    """Load tokenizer, converting tokenizer.json for older tokenizers if needed."""
    tokenizer_path = os.path.join(model_id, "tokenizer.json")
    with open(tokenizer_path, encoding="utf-8") as handle:
        data = json.load(handle)

    merges = data.get("model", {}).get("merges", [])
    if not merges or not isinstance(merges[0], list):
        return AutoTokenizer.from_pretrained(model_id)

    # Older tokenizers expect merge pairs as "a b" strings, not ["a", "b"] lists.
    data["model"]["merges"] = [" ".join(piece) for piece in merges]
    tmpdir = tempfile.mkdtemp()
    for name in ("tokenizer_config.json", "special_tokens_map.json"):
        src = os.path.join(model_id, name)
        if os.path.isfile(src):
            shutil.copy(src, tmpdir)
    with open(os.path.join(tmpdir, "tokenizer.json"), "w", encoding="utf-8") as handle:
        json.dump(data, handle)
    return AutoTokenizer.from_pretrained(tmpdir)


SYSTEM = (
    "You solve International Linguistics Olympiad problems by reasoning from the "
    "data in CONTEXT you are given to solve the problems in QUERY. \n"
    "There are common TASK TYPES that we specify below, but "
    "you may meet a TASK TYPE you have never seen: read the "
    "instruction and the examples, and answer the QUERY in the same form they use.\n\n"
    "Common TASK TYPES and what to return: \n"
    "`translation`: return the translated form only, in the language the task asks for; \n"
    "`fill_blanks`: return only the missing form for each indicated blank "
    "(beware: this could be many different things: a word, a part of a word or a phonetic transcription---pay close attention to what part of the CONTEXT is missing in QUERY); \n"
    "`match_letters`: return only the option letter (for example A, B, C); \n"
    "`text_to_num`: return the number in digits; \n"
    "`num_to_text`: return the number written out in words, in the language asked; \n"
    "any other type: return exactly what the instruction asks for, nothing else. \n\n"
    "As the first part of your answer, reason step by step about (1) the linguistic "
    "rules that can be deduced from the given examples in CONTEXT, and (2) "
    "how to apply them to the given problems in QUERY, and (3) in what format answers need to be returned (words, numbers, phonetic transcriptions, ...). \n"
    "Then write a draft of the final answer. "
    "Subsequently, compare it with the format requirements again, "
    "and verify it's compliant with the deduced rules, and it is complete, i.e. has an answer for each element in QUERY. "
    "If necessary, correct and refine."
    "Finally, write a line that says exactly `FINAL ANSWERS:` "
    "and, below it, write the answers to the items requested in QUERY (not those in CONTEXT),"
    "one answer per line (separated by \n) in the order the items are asked for in the QUERY -- the "
    "bare answer only, no numbering, no quotes, no extra text, according to the given TASK TYPE."
)

SYSTEM_POST_EXPLAIN = (
    "You are a helpful assistant that explains the reasoning behind the answers to the International Linguistics Olympiad problems.\n"
    "You are given the following information:\n"
    "- The context of the problem\n"
    "- The task type\n"
    "- The query\n"
    "- The answer\n"
    "- The reasoning\n"
    "You need to explain the reasoning behind the answer in a way that is easy to understand and concisely focused on the key insights and rules deduced and applied.\n"
    "Do not include any other text, do not include the answer in the explanation, "
    "and do not invent any new information."
)

# Prefer a dedicated header line; also allow same-line answers after the colon.
# Fall back to the last "final answer(s):" substring if no clean header line exists.
FINAL_ANSWERS_LINE_RE = re.compile(
    # Header alone on the line (markdown / punctuation allowed), OR
    # header with first answer on the same line after a colon.
    r"(?im)^[^\w\n]*final answers?[^\w\n]*:?[ \t]*(?=\n|$)|"
    r"(?im)^[^\w\n]*final answers?\s*:\s*"
)
FINAL_ANSWERS_INLINE_RE = re.compile(
    r"(?is)\bfinal answers?\s*:\s*"
)


def extract_raw_final(text: str) -> str:
    """Return text after the last final-answers marker, or '' if none found."""
    line_matches = list(FINAL_ANSWERS_LINE_RE.finditer(text))
    if line_matches:
        return text[line_matches[-1].end() :]

    inline_matches = list(FINAL_ANSWERS_INLINE_RE.finditer(text))
    if inline_matches:
        return text[inline_matches[-1].end() :]

    return ""


def expected_answer_count(query: str, task_type: str) -> int:
    if task_type == "match_letters":
        numbered = re.findall(r"^\s*\d+\.", query, re.MULTILINE)
        return len(numbered) or 1

    if "blanks" in query.lower():
        range_match = re.search(r"\((\d+)-(\d+)\)", query)
        if range_match:
            return int(range_match.group(2)) - int(range_match.group(1)) + 1
        return len(re.findall(r"\(\d+\)", query)) or 1

    numbered = re.findall(r"^\s*\d+[.)]", query, re.MULTILINE)
    return len(numbered) or 1


def split_single_line_answer(text: str, expected: int, task_type: str) -> list[str]:
    text = text.strip()
    if expected <= 1:
        return [text]

    def try_split(pattern: str) -> list[str] | None:
        parts = [part.strip() for part in re.split(pattern, text) if part.strip()]
        return parts if len(parts) == expected else None

    if task_type == "match_letters":
        for pattern in (r"\s+", r",\s*", r";\s*"):
            if result := try_split(pattern):
                return result
        letters = re.findall(r"[A-Za-z]", text)
        if len(letters) == expected:
            return [letter.upper() for letter in letters]
        return [text]

    if task_type in ("text_to_num", "num_to_text"):
        for pattern in (r",\s*", r";\s*", r"\s+"):
            if result := try_split(pattern):
                return result
        return [text]

    for pattern in (r";\s*", r",\s*"):
        if result := try_split(pattern):
            return result
    return [text]


def parse_answer_lines(text_after_marker: str, query: str, task_type: str) -> list[str]:
    """Parse cleaned answer lines from the raw final-answers section."""
    answers = []
    for line in text_after_marker.splitlines():
        stripped_line = line.strip("`").strip()
        if stripped_line == "":
            continue

        match_numbered_prefix = re.match(r"^\s*\d+[.)]\s+(.*)", stripped_line)
        if match_numbered_prefix:
            cleaned_line = match_numbered_prefix.group(1).strip()
        else:
            cleaned_line = stripped_line

        cleaned_line = re.sub(r"\*\*", "", cleaned_line).strip()

        if task_type == "match_letters":
            parts = [
                part.strip("().[]")
                for part in re.split(r"[\s,;]+", cleaned_line)
                if part.strip()
            ]
            if not (
                len(parts) > 1
                and all(re.fullmatch(r"[A-Za-z]", part) for part in parts)
            ):
                match_letter_word = re.match(
                    r"^\s*(?:\(([A-Za-z])\)|\[([A-Za-z])\]|([A-Za-z]))\.?:?\s*(.*)$",
                    cleaned_line,
                )
                if match_letter_word:
                    letter = (
                        match_letter_word.group(1)
                        or match_letter_word.group(2)
                        or match_letter_word.group(3)
                    )
                    cleaned_line = letter.upper()

        if cleaned_line:
            answers.append(cleaned_line)

    expected = expected_answer_count(query, task_type)
    if len(answers) == 1 and expected > 1:
        answers = split_single_line_answer(answers[0], expected, task_type)
    return answers


def postprocess_answer(text, query, task_type):
    """Keep only the content after the last 'FINAL ANSWERS' marker, one answer per line."""
    text_after_marker = extract_raw_final(text)
    if not text_after_marker.strip():
        return []
    return parse_answer_lines(text_after_marker, query, task_type)


tok = load_tokenizer(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.float16, device_map="auto"
).eval()

with open("/tmp/data/test.csv", encoding="utf-8", newline="") as f:
    test_rows = list(csv.DictReader(f))

outputs_queries_types = []
for r in test_rows:

    # Create the prompt.
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content":
        f"CONTEXT:{r['context'].strip()}\nTASK TYPE:`{r['task_type']}`\n\nQUERY:{r['query'].strip()}"},
    ]
    ids = tok.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt",
    ).to(model.device)

    # Generate the answer.
    done = False
    attempts = 0
    text = ""
    while not done and attempts < MAX_ATTEMPTS:
        with torch.no_grad():
            out = model.generate(
                ids,
                max_new_tokens=MAX_NEW_TOKENS,
                do_sample=True,
                temperature=TEMPERATURE,
                top_p=TOP_P,)
        text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
        # Retry until we can extract a non-empty final-answers section.
        attempts += 1
        if extract_raw_final(text).strip():
            done = True
        else:
            print(f'TRYING AGAIN...attempts #{attempts}/{MAX_ATTEMPTS}')

    # Generate the explanation.
    messages_post_explain = [
        {"role": "system", "content": SYSTEM_POST_EXPLAIN},
        {"role": "user", "content":
        f"CONTEXT:{r['context'].strip()}\nTASK TYPE:`{r['task_type']}`\n\nQUERY:{r['query'].strip()}\n\nANSWER:{text.strip()}"},
    ]
    ids_post_explain = tok.apply_chat_template(
        messages_post_explain, add_generation_prompt=True, return_tensors="pt",
    ).to(model.device)
    with torch.no_grad():
        out_post_explain = model.generate(
            ids_post_explain,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False)
    text_post_explain = tok.decode(out_post_explain[0][ids_post_explain.shape[-1]:], skip_special_tokens=True).strip()

    outputs_queries_types.append((text, text_post_explain, r['id'], r['query'], r['task_type']))
    print(f"{len(outputs_queries_types)}/{len(test_rows)} done", flush=True)

rows = []
for answer, explanation, row_id, query, task_type in outputs_queries_types:
    answers = postprocess_answer(answer, query, task_type)
    rows.append({"id": row_id, "pred": json.dumps(answers, ensure_ascii=False), "explanation": explanation})
with open("submission.csv", "w", encoding="utf-8", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=["id", "pred", "explanation"])
    writer.writeheader()
    writer.writerows(rows)
print("wrote submission.csv", flush=True)