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 random
import shutil
import tempfile
import unicodedata
from collections import Counter

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 = 1200  # was 2000; cut to fit T4 wall clock
TEMPERATURE = 0.8
TOP_P = 0.95
MAX_ATTEMPTS = 1  # no per-lang retry; majority vote absorbs misses
NUM_LANGS = 3  # English + 2 others (fixed; was random 3–5)

# Languages tiny-aya-global handles well for chain-of-thought.
REASONING_LANGUAGES = [
    "English",
    "Spanish",
    "French",
    "German",
    "Portuguese",
    "Italian",
    "Dutch",
    "Russian",
    "Arabic",
    "Simplified Chinese",
    "Japanese",
    "Korean",
    "Turkish",
    "Hindi",
    "Indonesian",
    "Vietnamese",
    "Polish",
    "Swedish",
    "Greek",
    "Hebrew",
    # "Swahili",
    "Ukrainian",
    "Romanian",
    "Czech",
    "Hungarian",
]


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_TEMPLATE = (
    "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"
    "IMPORTANT: Write ALL of your step-by-step reasoning in {language}. "
    "Do not mix languages in the reasoning. "
    "The FINAL ANSWERS section must still use the English marker `FINAL ANSWERS:` "
    "and the answer values themselves must follow the TASK TYPE / QUERY requirements "
    "(do not translate those answers into {language} unless the query asks for that).\n\n"
    "As the first part of your answer, reason step by step in {language} 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."
)


# Prefer a dedicated header line; also allow same-line answers after the colon.
FINAL_ANSWERS_LINE_RE = re.compile(
    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."""
    text_after_marker = extract_raw_final(text)
    if not text_after_marker.strip():
        return []
    return parse_answer_lines(text_after_marker, query, task_type)


def normalize_for_vote(text: str, task_type: str) -> str:
    text = unicodedata.normalize("NFC", text.strip())
    if task_type == "match_letters":
        return text.upper()
    return " ".join(text.split())


def majority_vote(
    lang_rollouts: list[tuple[str, list[str]]],
    expected: int,
    task_type: str,
) -> list[str]:
    """Per-item majority vote; ties break toward the English rollout."""
    if expected <= 0:
        return []

    # Prefer rollouts whose length matches the expected answer count.
    eligible = [
        (lang, rollout)
        for lang, rollout in lang_rollouts
        if len(rollout) == expected and any(a.strip() for a in rollout)
    ]
    if not eligible:
        eligible = [
            (lang, rollout)
            for lang, rollout in lang_rollouts
            if any(a.strip() for a in rollout)
        ]
    if not eligible:
        return []

    final: list[str] = []
    for i in range(expected):
        tagged = [
            (lang, rollout[i])
            for lang, rollout in eligible
            if i < len(rollout) and rollout[i].strip()
        ]
        if not tagged:
            final.append("")
            continue

        pairs = [
            (lang, normalize_for_vote(ans, task_type), ans)
            for lang, ans in tagged
        ]
        counter = Counter(norm for _, norm, _ in pairs)
        top_count = max(counter.values())
        tied_norms = {norm for norm, count in counter.items() if count == top_count}

        english_pair = next(
            ((norm, ans) for lang, norm, ans in pairs if lang == "English"),
            None,
        )
        if english_pair is not None and english_pair[0] in tied_norms:
            winner_norm = english_pair[0]
            # Prefer English's surface form when it matches the winning norm.
            final.append(english_pair[1])
            continue

        winner_norm = sorted(tied_norms)[0]
        originals = [ans for _, norm, ans in pairs if norm == winner_norm]
        final.append(Counter(originals).most_common(1)[0][0])
    return final


def sample_reasoning_languages() -> list[str]:
    """Always include English; sample NUM_LANGS-1 others."""
    others = [lang for lang in REASONING_LANGUAGES if lang != "English"]
    return ["English"] + random.sample(others, NUM_LANGS - 1)


def generate_for_language(tok, model, language: str, context: str, task_type: str, query: str) -> str:
    system = SYSTEM_TEMPLATE.format(language=language)
    messages = [
        {"role": "system", "content": system},
        {
            "role": "user",
            "content": (
                f"CONTEXT:{context.strip()}\n"
                f"TASK TYPE:`{task_type}`\n\n"
                f"QUERY:{query.strip()}\n\n"
                f"Remember: reason entirely in {language}."
            ),
        },
    ]
    ids = tok.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt",
    ).to(model.device)

    text = ""
    for attempt in range(1, MAX_ATTEMPTS + 1):
        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 extract_raw_final(text).strip():
            return text
        print(
            f"  [{language}] retry {attempt}/{MAX_ATTEMPTS}: no FINAL ANSWERS",
            flush=True,
        )
    return text


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

# Write incrementally so a wall-clock kill still leaves a partial submission.csv.
with open("submission.csv", "w", encoding="utf-8", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=["id", "pred"])
    writer.writeheader()
    f.flush()

    for idx, r in enumerate(test_rows, start=1):
        languages = sample_reasoning_languages()
        print(f"{idx}/{len(test_rows)} langs={languages}", flush=True)

        lang_rollouts: list[tuple[str, list[str]]] = []
        for language in languages:
            text = generate_for_language(
                tok, model, language, r["context"], r["task_type"], r["query"],
            )
            answers = postprocess_answer(text, r["query"], r["task_type"])
            lang_rollouts.append((language, answers))
            print(f"  [{language}] parsed={answers!r}", flush=True)

        expected = expected_answer_count(r["query"], r["task_type"])
        voted = majority_vote(lang_rollouts, expected, r["task_type"])
        print(f"  vote -> {voted!r}", flush=True)

        writer.writerow({"id": r["id"], "pred": json.dumps(voted, ensure_ascii=False)})
        f.flush()

print("wrote submission.csv", flush=True)