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