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

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 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 postprocess_answer(text, query, task_type):
    """Keep only the lines after the last 'FINAL ANSWERS:' marker, one answer per line,
    stopping at the first empty line. If eval_type is multiple and only one line as answer, split at whitespace."""
    # Updated regex to be more flexible with surrounding characters
    marker_match = list(re.finditer(r"(?im)^[^\w\n]*final answers?[^\w\n]*:?\s*$", text))
    if marker_match:
        text_after_marker = text[marker_match[-1].end():]
        #print('FOUND FINAL ANSWER', text_after_marker)
    else:
        #print("No 'FINAL ANSWERS:' marker found")
        return []

    answers = []
    answer_lines = text_after_marker.splitlines()
    for i, line in enumerate(answer_lines):
        stripped_line = line.strip('`').strip()

        # Stop processing if an empty line is encountered (not as first line)
        if stripped_line=='':
          continue

        # Use a more precise regex to only remove numbering if it's a prefix to other text
        # This ensures that lines which are just numbers (e.g., '1') are not stripped.
        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

        # Remove any bold markdown '**'
        cleaned_line = re.sub(r"\*\*", "", cleaned_line).strip()

        # Specific handling for 'match_letters' task type to strip extra words
        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()

        # Append the cleaned, non-empty line
        if cleaned_line:
            answers.append(cleaned_line)

    #print('PARSED ANSWERS', answers)

    # Compare against QUERY length: sometimes model forgets newlines
    #print('QUERY', query)
    expected = expected_answer_count(query, task_type)
    query_len = len(query.splitlines()) - 2
    #print(query_len)
    if len(answers) == 1 and expected > 1:
        answers = split_single_line_answer(answers[0], expected, task_type)
    return answers

rows = []
for answer, 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)})
with open("submission.csv", "w", encoding="utf-8", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=["id", "pred"])
    writer.writeheader()
    writer.writerows(rows)
print("wrote submission.csv", flush=True)