script.py

#import subprocess, sys
#subprocess.run(
#    [
#        sys.executable,
#        "-m",
#        "pip",
#        "install",
#        "-q",
#        "transformers>=4.43",
#        "accelerate>=0.30",
#        "torch>=2.2",
#        "pandas",
#    ],
#    check=True,
#)

import os
import re
import json

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

from incontext_examples import INCONTEXT_EXAMPLES


# The repository is the working directory at runtime.
# The evaluation sandbox has no internet access.
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

MODEL_ID = "."
TEST_PATH = "/tmp/data/test.csv"
OUTPUT_PATH = "submission.csv"

MAX_NEW_TOKENS = 1536

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token


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


df = pd.read_csv(
    TEST_PATH,
    dtype=str,
).fillna("")


required_columns = {
    "id",
    "context",
    "query",
    "task_type",
}

missing_columns = required_columns - set(df.columns)

if missing_columns:
    raise ValueError(
        "The test CSV is missing required columns: "
        f"{sorted(missing_columns)}"
    )


SYSTEM = (
    "You solve International Linguistics Olympiad problems by reasoning from the "
    "data you are given. You may meet a task type you have never seen: read the "
    "instruction and the examples, and answer in the same form they use. "
    "Common task types and what to give -- "
    "translation: the translated form only, in the language the task asks for; "
    "fill_blanks: only the missing form for each blank; "
    "match_letters: only the option letter, for example A, B, or C; "
    "text_to_num: the number in digits; "
    "num_to_text: the number written out in words, in the language asked; "
    "any other type: give exactly what the instruction asks, nothing else. "
    "Reason step by step first. Then write a line that says exactly FINAL ANSWERS: "
    "and, below it, one answer per line in the order the items are asked. "
    "Each final answer must be bare: no numbering, no quotes, and no extra text."
)


def normalize_task_type(task_type):
    """
    Normalize task-type names so capitalization, spaces, and hyphens
    do not prevent a match.

    Examples:
        "Num To Text" -> "num_to_text"
        "fill-blanks" -> "fill_blanks"
    """
    task_type = str(task_type).strip().lower()
    task_type = re.sub(r"[\s-]+", "_", task_type)
    return task_type


def parse_answers(text):
    """
    Keep only nonempty lines after the final FINAL ANSWERS marker.

    If the model adds numbering such as:
        1. answer
        2) answer

    the numbering is removed.
    """
    markers = list(
        re.finditer(
            r"(?im)^\s*final answers?\s*:?\s*$",
            text,
        )
    )

    if markers:
        text = text[markers[-1].end():]

    answers = []

    for line in text.splitlines():
        line = re.sub(
            r"^\s*\d+[.)]\s*",
            "",
            line,
        ).strip()

        if line:
            answers.append(line)

    return answers


def format_example_assistant_response(example):
    """
    Format the stored reasoning trace and answers as an assistant message.
    """
    reasoning = str(
        example.get("reasoning", "")
    ).strip()

    answers = example.get("answers", [])

    if isinstance(answers, str):
        answers = [
            line.strip()
            for line in answers.splitlines()
            if line.strip()
        ]

    formatted_answers = "\n".join(
        str(answer).strip()
        for answer in answers
    )

    if reasoning:
        return (
            f"{reasoning}\n\n"
            f"FINAL ANSWERS:\n"
            f"{formatted_answers}"
        )

    return (
        f"FINAL ANSWERS:\n"
        f"{formatted_answers}"
    )


def get_incontext_example(task_type):
    """
    Retrieve the demonstration matching the current problem's task type.

    Returns None when no demonstration exists for that task type.
    """
    normalized_type = normalize_task_type(task_type)

    # Normalize the dictionary keys as well, so keys such as
    # "num-to-text" and "num_to_text" are treated equivalently.
    normalized_examples = {
        normalize_task_type(key): value
        for key, value in INCONTEXT_EXAMPLES.items()
    }

    return normalized_examples.get(normalized_type)


def build_messages(row):
    """
    Construct the conversation for one test problem.

    The conversation contains:
      1. the system prompt;
      2. a solved demonstration of the same task type, when available;
      3. the actual test problem.
    """
    task_type = normalize_task_type(
        row["task_type"]
    )

    example = get_incontext_example(
        task_type
    )

    messages = [
        {
            "role": "system",
            "content": SYSTEM,
        }
    ]

    if example is not None:
        example_context = str(
            example.get("context", "")
        ).strip()

        example_query = str(
            example.get("query", "")
        ).strip()

        messages.extend(
            [
                {
                    "role": "user",
                    "content": (
                        f"Task type: {task_type}\n\n"
                        f"{example_context}\n\n"
                        f"{example_query}"
                    ),
                },
                {
                    "role": "assistant",
                    "content": format_example_assistant_response(
                        example
                    ),
                },
            ]
        )
    else:
        print(
            "Warning: no in-context example found for "
            f"task_type={task_type!r}, id={row['id']!r}. "
            "Running the problem without a demonstration."
        )

    problem_context = str(
        row["context"]
    ).strip()

    problem_query = str(
        row["query"]
    ).strip()

    messages.append(
        {
            "role": "user",
            "content": (
                f"Task type: {task_type}\n\n"
                f"{problem_context}\n\n"
                f"{problem_query}"
            ),
        }
    )

    return messages


rows = []


for _, row in df.iterrows():
    messages = build_messages(
        row
    )

    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    ).to(model.device)

    input_length = inputs[
        "input_ids"
    ].shape[-1]

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    generated_text = tokenizer.decode(
        output[0][input_length:],
        skip_special_tokens=True,
    ).strip()

    answers = parse_answers(
        generated_text
    )

    rows.append(
        {
            "id": row["id"],
            "pred": json.dumps(
                answers,
                ensure_ascii=False,
            ),
        }
    )


submission_df = pd.DataFrame(
    rows
)

submission_df.to_csv(
    OUTPUT_PATH,
    index=False,
)

print(
    f"Saved {len(submission_df)} predictions to {OUTPUT_PATH}"
)