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.\n"
    "Structure your response in exactly two sections with these headers:\n\n"
    "**Reasoning:**\n"
    "(your step-by-step analysis of linguistic rules, how to apply them, and answer format)\n\n"
    "**Final Answers:**\n"
    "(one answer per line, in query order -- bare answers only, no numbering, no quotes, "
    "no extra text, according to the 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
    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()
        # IF NO FINAL ANSWER keyword is used, try again.
        attempts += 1
        if "final answer" not in text.lower() or text.lower().split('final answer')[1].split('\n')==0:
            print(f'TRYING AGAIN...attempts #{attempts+1}/{MAX_ATTEMPTS}')
        else:
            done = True

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

# Postprocess and store the answers.
def _normalize_header(line: str) -> str:
    header = line.strip().lower()
    header = re.sub(r"^#{1,3}\s*", "", header)
    header = re.sub(r"\*+", "", header)
    return header.rstrip(":").strip()


def _is_reasoning_header(line: str) -> bool:
    header = _normalize_header(line)
    return header == "reasoning" or header.endswith(" reasoning")


def _is_final_answers_header(line: str) -> bool:
    header = _normalize_header(line)
    return header in ("final answer", "final answers")


def split_response(text: str) -> tuple[str, str]:
    """Return (explanation, raw final-answers section text)."""
    lines = text.splitlines(keepends=True)
    offset = 0
    reasoning_content_start = None
    explanation_end = None
    last_final_content_start = None

    reasoning_header_seen = False
    final_header_indices: list[int] = []

    for line in lines:
        if _is_reasoning_header(line) and not reasoning_header_seen:
            reasoning_header_seen = True
            reasoning_content_start = offset + len(line)
        elif _is_final_answers_header(line):
            final_header_indices.append(offset)
            if reasoning_content_start is not None and explanation_end is None:
                explanation_end = offset
            last_final_content_start = offset + len(line)
        offset += len(line)

    raw_final = text[last_final_content_start:].strip() if last_final_content_start is not None else ""

    explanation = ""
    if reasoning_content_start is not None and explanation_end is not None:
        explanation = text[reasoning_content_start:explanation_end].strip()
    if not explanation:
        explanation = raw_final
    return explanation, raw_final


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):
    """Extract explanation and parsed final answers from model output."""
    explanation, raw_final = split_response(text)
    if not raw_final:
        return [], explanation
    return parse_answer_lines(raw_final, query, task_type), explanation

rows = []
for answer, row_id, query, task_type in outputs_queries_types:
    answers, explanation = 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)