script.py

#!/usr/bin/env python3
"""IOL-AI 2026 submission script.

Reads /tmp/data/test.csv and writes submission.csv.
Model weights must be shipped in the repo; default local path is ./model.
"""
from __future__ import annotations

import json
import os
import re
from pathlib import Path
from typing import Any


def _ensure_deps() -> None:
    """Fail clearly instead of installing anything during evaluation."""
    try:
        import pandas  # noqa: F401
    except ImportError as exc:
        raise RuntimeError(
            "Missing required dependency: pandas. Install dependencies before evaluation; "
            "script.py will not download packages or call the internet at runtime."
        ) from exc


_ensure_deps()

import pandas as pd  # noqa: E402


INPUT_CSV = Path("/tmp/data/test.csv")
OUTPUT_CSV = Path("submission.csv")
MODEL_DIR = os.environ.get("MODEL_DIR", "./model")
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "768"))
DUMMY_MODE = os.environ.get("IOL_DUMMY", "0") == "1"

# Force local/offline loading for HF libraries.
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

ANSWER_KEYS = ("answers", "answer", "pred", "prediction", "predictions")
FENCED_BLOCK_RE = re.compile(r"```(?:json)?\s*(.*?)```", flags=re.S | re.I)
LIST_ITEM_RE = re.compile(
    r"^\s*(?:(?:[-*+]\s*)?(?:\(?\d{1,3}\)?[\).:]|[A-Za-z][\).:])|[-*+\u2013\u2014\u2022\u2023\u2043\u2219\u25e6])\s+(.*?)\s*$"
)
HEADER_LINE_RE = re.compile(r"^\s*(?:answers?|predictions?|output|final answers?)\s*:?\s*$", flags=re.I)
INTRO_LINE_RE = re.compile(r"^\s*(?:here are|the answers are|my answers are)\b", flags=re.I)
JSON_SCAFFOLD_LINE_RE = re.compile(
    r"^\s*(?:[\{\}\[\],]+|[\"']?(?:answers?|predictions?|pred|prediction)[\"']?\s*:\s*\[?)\s*$", flags=re.I
)
QUERY_ITEM_RE = re.compile(r"(?m)^\s*(?:\(?\d{1,3}\)?[\).:])\s+")

OUTPUT_CONTRACT = (
    "Answer every numbered item. Think silently. Return only valid JSON: "
    "a list of final answer strings in order. Do not include explanations, labels, markdown, or extra text. "
    "Preserve Unicode and diacritics exactly."
)

TASK_PROMPTS: dict[str, dict[str, str]] = {
    "translation": {
        "system": (
            "You solve IOL translation tasks from the provided linguistic data. "
            "Infer word order, morphology, agreement, and lexical correspondences from the examples. "
            f"{OUTPUT_CONTRACT}"
        ),
        "user": (
            "Task type: translation\n"
            "Evaluation type: {eval_type}\n"
            "{count_instruction}\n\n"
            "Use the context examples to translate each numbered query item.\n\n"
            "CONTEXT:\n{context}\n\n"
            "QUERY:\n{query}"
        ),
    },
    "fill_blanks": {
        "system": (
            "You solve IOL fill-in-the-blank tasks from the provided patterns. "
            "Infer the missing forms or words needed to complete each item. "
            f"{OUTPUT_CONTRACT}"
        ),
        "user": (
            "Task type: fill_blanks\n"
            "Evaluation type: {eval_type}\n"
            "{count_instruction}\n\n"
            "Fill each blank in the numbered query items using only the context patterns.\n\n"
            "CONTEXT:\n{context}\n\n"
            "QUERY:\n{query}"
        ),
    },
    "match_letters": {
        "system": (
            "You solve IOL letter-matching tasks from the provided correspondences. "
            "Infer which letters, choices, or labels match each numbered item. "
            f"{OUTPUT_CONTRACT}"
        ),
        "user": (
            "Task type: match_letters\n"
            "Evaluation type: {eval_type}\n"
            "{count_instruction}\n\n"
            "Return the matching letter, choice, or label for each numbered query item.\n\n"
            "CONTEXT:\n{context}\n\n"
            "QUERY:\n{query}"
        ),
    },
    "text_to_num": {
        "system": (
            "You solve IOL number-system tasks that convert written forms into numerals. "
            "Infer the numeral system and output the numeric value for each item. "
            f"{OUTPUT_CONTRACT}"
        ),
        "user": (
            "Task type: text_to_num\n"
            "Evaluation type: {eval_type}\n"
            "{count_instruction}\n\n"
            "Convert each numbered written form into its numeric value.\n\n"
            "CONTEXT:\n{context}\n\n"
            "QUERY:\n{query}"
        ),
    },
    "num_to_text": {
        "system": (
            "You solve IOL number-system tasks that convert numerals into written forms. "
            "Infer the numeral system and output the written form for each item. "
            f"{OUTPUT_CONTRACT}"
        ),
        "user": (
            "Task type: num_to_text\n"
            "Evaluation type: {eval_type}\n"
            "{count_instruction}\n\n"
            "Convert each numbered numeric value into the target written form.\n\n"
            "CONTEXT:\n{context}\n\n"
            "QUERY:\n{query}"
        ),
    },
    "fallback": {
        "system": (
            "You solve IOL pattern-inference tasks using only the provided context and query. "
            "Infer the requested transformation or mapping for each numbered item. "
            f"{OUTPUT_CONTRACT}"
        ),
        "user": (
            "Task type: {task_type}\n"
            "Evaluation type: {eval_type}\n"
            "{count_instruction}\n\n"
            "Solve each numbered query item using only the context.\n\n"
            "CONTEXT:\n{context}\n\n"
            "QUERY:\n{query}"
        ),
    },
}


def count_expected_answers(query: str) -> int | None:
    """Estimate number of numbered items in the query.

    Returns None if no reliable numbering is visible.
    """
    # Common IOL format: each item starts with `17.`, `17)`, `(17)` etc.
    matches = QUERY_ITEM_RE.findall(query)
    if matches:
        return len(matches)
    return None


def expected_item_count(query: str) -> int | None:
    """Backward-compatible alias for count_expected_answers."""
    return count_expected_answers(query)


def clean_answer_text(text: str) -> str:
    """Remove only formatting wrappers around an answer."""
    text = text.strip()
    quote_pairs = {
        '"': '"',
        "'": "'",
        "\u201c": "\u201d",
        "\u2018": "\u2019",
    }
    if len(text) >= 2 and quote_pairs.get(text[0]) == text[-1]:
        text = text[1:-1].strip()
    return text


def fit_answer_count(answers: list[str], n_expected: int | None) -> list[str]:
    if n_expected is None:
        return answers
    if len(answers) > n_expected:
        return answers[:n_expected]
    if len(answers) < n_expected:
        return answers + [""] * (n_expected - len(answers))
    return answers


def answer_value_to_text(value: Any) -> str:
    if value is None:
        return ""
    if isinstance(value, str):
        return clean_answer_text(value)
    if isinstance(value, (dict, list)):
        return json.dumps(value, ensure_ascii=False)
    return clean_answer_text(str(value))


def answers_from_json_value(value: Any) -> list[str] | None:
    if isinstance(value, dict):
        for key in ANSWER_KEYS:
            if key in value:
                return answers_from_json_value(value[key])
        return None
    if isinstance(value, list):
        answers: list[str] = []
        for item in value:
            if isinstance(item, dict):
                item_answers = answers_from_json_value(item)
                answers.append(item_answers[0] if item_answers else answer_value_to_text(item))
            else:
                answers.append(answer_value_to_text(item))
        return answers
    if isinstance(value, (str, int, float, bool)) or value is None:
        return [answer_value_to_text(value)]
    return None


def iter_json_values(source: str, scan_embedded: bool) -> list[Any]:
    decoder = json.JSONDecoder()
    source = source.strip()
    if not source:
        return []

    starts: list[int] = []
    if source[:1] in "[{":
        starts.append(0)
    if scan_embedded:
        starts.extend(i for i, char in enumerate(source) if char in "[{" and i not in starts)

    values: list[Any] = []
    seen: set[str] = set()
    for start in starts:
        try:
            value, _ = decoder.raw_decode(source[start:])
        except json.JSONDecodeError:
            continue
        signature = json.dumps(value, ensure_ascii=False, sort_keys=True)
        if signature not in seen:
            values.append(value)
            seen.add(signature)
    return values


def select_answer_candidate(candidates: list[list[str]], n_expected: int | None) -> list[str] | None:
    if not candidates:
        return None
    if n_expected is not None:
        for answers in candidates:
            if len(answers) == n_expected:
                return answers
    return candidates[0]


def has_expected_count(answers: list[str], n_expected: int | None) -> bool:
    return n_expected is None or len(answers) == n_expected


def extract_json_object(text: str) -> dict[str, Any] | None:
    """Try to recover a JSON object from a model response."""
    sources = [text, *FENCED_BLOCK_RE.findall(text)]
    for source in sources:
        for value in iter_json_values(source, scan_embedded=True):
            if isinstance(value, dict):
                return value
    return None


def parse_json_answers(raw_text: str, n_expected: int | None, scan_embedded: bool) -> list[str] | None:
    sources = [raw_text, *FENCED_BLOCK_RE.findall(raw_text)]
    candidates: list[list[str]] = []
    for source in sources:
        for value in iter_json_values(source, scan_embedded=scan_embedded):
            answers = answers_from_json_value(value)
            if answers is not None:
                candidates.append(answers)
    return select_answer_candidate(candidates, n_expected)


def parse_list_item_answers(raw_text: str) -> list[str]:
    answers: list[str] = []
    for line in raw_text.splitlines():
        match = LIST_ITEM_RE.match(line)
        if match:
            answer = clean_answer_text(match.group(1))
            if answer:
                answers.append(answer)
    return answers


def parse_plain_line_answers(raw_text: str, n_expected: int | None) -> list[str]:
    lines: list[str] = []
    for line in raw_text.splitlines():
        answer = clean_answer_text(line)
        if not answer or answer.startswith("```") or HEADER_LINE_RE.match(answer) or JSON_SCAFFOLD_LINE_RE.match(answer):
            continue
        lines.append(answer)

    if n_expected is not None and len(lines) > n_expected:
        filtered = [line for line in lines if not INTRO_LINE_RE.match(line)]
        if len(filtered) >= n_expected:
            return filtered[:n_expected]
    return lines


def parse_model_output(text: str, expected_n: int | None) -> list[str]:
    """Convert model text into a list of answer strings."""
    json_answers = parse_json_answers(text, expected_n, scan_embedded=False)
    if json_answers is not None and has_expected_count(json_answers, expected_n):
        return fit_answer_count(json_answers, expected_n)

    list_answers = parse_list_item_answers(text)
    if list_answers and (json_answers is None or has_expected_count(list_answers, expected_n)):
        return fit_answer_count(list_answers, expected_n)
    if json_answers is not None:
        return fit_answer_count(json_answers, expected_n)

    embedded_json_answers = parse_json_answers(text, expected_n, scan_embedded=True)
    if embedded_json_answers is not None:
        return fit_answer_count(embedded_json_answers, expected_n)

    plain_answers = parse_plain_line_answers(text, expected_n)
    return fit_answer_count(plain_answers, expected_n)


def normalize_answers(raw_text: str, n_expected: int | None) -> list[str]:
    """Backward-compatible alias for parse_model_output."""
    return parse_model_output(raw_text, n_expected)


def build_prompt(row: pd.Series, n_expected: int | None) -> list[dict[str, str]]:
    task_type = str(row.get("task_type", "")).strip()
    task_key = task_type.lower() or "fallback"
    template = TASK_PROMPTS.get(task_key, TASK_PROMPTS["fallback"])
    eval_type = str(row.get("eval_type", "")).strip()
    context = str(row.get("context", "")).strip()
    query = str(row.get("query", "")).strip()
    count_instruction = (
        f"Return exactly {n_expected} answers." if n_expected is not None else "Return one answer per numbered item."
    )
    system = template["system"]
    user = template["user"].format(
        task_type=task_type or "unknown",
        eval_type=eval_type,
        count_instruction=count_instruction,
        context=context,
        query=query,
    )
    return [{"role": "system", "content": system}, {"role": "user", "content": user}]


def validate_model_dir(model_dir: str) -> Path:
    path = Path(model_dir)
    if not path.exists():
        raise FileNotFoundError(
            f"Missing local model folder: {path}. Ship model weights with the repo, "
            "or set MODEL_DIR to an existing local directory. Runtime downloads are disabled."
        )
    if not path.is_dir():
        raise NotADirectoryError(f"MODEL_DIR must be a local directory, got: {path}")
    if not (path / "config.json").exists():
        raise FileNotFoundError(f"Local model folder is missing config.json: {path}")
    return path


def load_model():
    if DUMMY_MODE:
        return None, None
    model_path = validate_model_dir(MODEL_DIR)

    try:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
    except ImportError as exc:
        raise RuntimeError(
            "Missing required model dependency: torch/transformers. Install dependencies before evaluation; "
            "script.py will not download packages, model weights, or call the internet at runtime."
        ) from exc

    tok = AutoTokenizer.from_pretrained(model_path, local_files_only=True, trust_remote_code=True)
    model_kwargs = {
        "device_map": "auto",
        "local_files_only": True,
        "trust_remote_code": True,
    }
    try:
        model = AutoModelForCausalLM.from_pretrained(model_path, dtype=torch.float16, **model_kwargs).eval()
    except TypeError:
        model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, **model_kwargs).eval()
    return tok, model


def model_input_device(model):
    device = getattr(model, "device", None)
    if device is not None:
        return device
    return next(model.parameters()).device


def build_generation_inputs(tok, model, messages: list[dict[str, str]]):
    try:
        encoded = tok.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt",
            return_dict=True,
        )
    except TypeError:
        encoded = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")

    device = model_input_device(model)
    if hasattr(encoded, "to"):
        encoded = encoded.to(device)

    if hasattr(encoded, "keys") and "input_ids" in encoded.keys():
        input_ids = encoded["input_ids"]
        return {key: encoded[key] for key in encoded.keys()}, input_ids.shape[-1]

    encoded = encoded.to(device)
    return {"input_ids": encoded}, encoded.shape[-1]


def generate_one(tok, model, messages: list[dict[str, str]]) -> str:
    if DUMMY_MODE:
        # Useful for testing CSV shape without downloading weights.
        return json.dumps({"answers": ["DUMMY"]}, ensure_ascii=False)

    import torch

    generation_inputs, prompt_len = build_generation_inputs(tok, model, messages)
    with torch.no_grad():
        out = model.generate(
            **generation_inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            pad_token_id=tok.eos_token_id,
        )
    return tok.decode(out[0][prompt_len:], skip_special_tokens=True).strip()


def main() -> None:
    if not INPUT_CSV.exists():
        raise FileNotFoundError(f"Missing input CSV: {INPUT_CSV}")

    df = pd.read_csv(INPUT_CSV, dtype=str).fillna("")
    tok, model = load_model()
    rows: list[dict[str, str]] = []

    for _, row in df.iterrows():
        n_expected = count_expected_answers(str(row.get("query", "")))
        messages = build_prompt(row, n_expected)
        raw = generate_one(tok, model, messages)
        answers = parse_model_output(raw, n_expected)

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

    pd.DataFrame(rows, columns=["id", "pred"]).to_csv(OUTPUT_CSV, index=False)
    print(f"Wrote {OUTPUT_CSV} with {len(rows)} rows")


if __name__ == "__main__":
    main()