script.py

"""Offline IOL-AI 2026 direct-pass submission for Qwen3.5-9B + our LoRA."""

from __future__ import annotations

import csv
import gc
import os
from pathlib import Path
import random
import time
from typing import Any

from iol_contract import (
    expected_answer_count,
    parse_answer_lines,
    serialized_prediction,
    validate_id_sequence,
)
from runtime_bootstrap import (
    assert_runtime_versions,
    bootstrap_local_runtime,
    configure_offline_environment,
    patch_torch24_for_single_gpu,
)


ROOT = Path(__file__).resolve().parent
INPUT_CSV = Path(os.environ.get("IOL_INPUT_CSV", "/tmp/data/test.csv"))
OUTPUT_CSV = Path(os.environ.get("IOL_OUTPUT_CSV", str(ROOT / "submission.csv")))
ADAPTER_DIR = ROOT / "adapter"
SYSTEM_PROMPT_PATH = ROOT / "system_prompt.txt"
SEED = 3407
MAX_NEW_TOKENS = int(os.environ.get("IOL_MAX_NEW_TOKENS", "512"))
HARD_STOP_SECONDS = float(os.environ.get("IOL_HARD_STOP_SECONDS", str(27.5 * 60)))


def set_reproducible_seed(torch: Any) -> None:
    random.seed(SEED)
    torch.manual_seed(SEED)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(SEED)


def read_input() -> list[dict[str, str]]:
    if INPUT_CSV.resolve().parent != Path("/tmp/data") and "IOL_INPUT_CSV" not in os.environ:
        raise RuntimeError("competition input must be /tmp/data/test.csv")
    with INPUT_CSV.open("r", encoding="utf-8-sig", newline="") as handle:
        reader = csv.DictReader(handle)
        required = {"id", "context", "query"}
        missing = required.difference(reader.fieldnames or [])
        if missing:
            raise ValueError(f"test.csv is missing columns: {sorted(missing)}")
        rows = [{key: value or "" for key, value in row.items()} for row in reader]
    if not rows:
        raise ValueError("test.csv is empty")
    ids = [row["id"] for row in rows]
    validate_id_sequence(ids, ids)
    return rows


def build_direct_prompt(row: dict[str, str], expected_n: int) -> str:
    # This matches the non-thinking DIRECT interface used during our SFT.
    return (
        "TASK_MODE=DIRECT\n"
        f"Return exactly {expected_n} answer line(s), in order, with no explanation.\n\n"
        "CONTEXT\n"
        f"{row['context'].strip()}\n\n"
        "QUERY\n"
        f"{row['query'].strip()}"
    )


def typed_text_messages(prompt: str) -> list[dict[str, Any]]:
    if not SYSTEM_PROMPT_PATH.is_file():
        raise RuntimeError("system_prompt.txt is missing")
    system_prompt = SYSTEM_PROMPT_PATH.read_text(encoding="utf-8").strip()
    if not system_prompt:
        raise RuntimeError("system_prompt.txt is empty")
    return [
        {
            "role": "system",
            "content": [{"type": "text", "text": system_prompt}],
        },
        {
            "role": "user",
            "content": [{"type": "text", "text": prompt}],
        },
    ]


def load_model_and_processor(torch: Any):
    from peft import PeftModel
    from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig

    if not ADAPTER_DIR.joinpath("adapter_config.json").is_file():
        raise RuntimeError("adapter/adapter_config.json is missing")
    processor = AutoProcessor.from_pretrained(
        ROOT,
        local_files_only=True,
        trust_remote_code=False,
    )
    quantization = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.float16,
    )
    base = AutoModelForMultimodalLM.from_pretrained(
        ROOT,
        local_files_only=True,
        trust_remote_code=False,
        dtype=torch.float16,
        quantization_config=quantization,
        device_map="auto",
        low_cpu_mem_usage=True,
    )
    model = PeftModel.from_pretrained(base, ADAPTER_DIR, is_trainable=False)
    model.eval()
    return model, processor


def generate_direct(model: Any, processor: Any, row: dict[str, str], expected_n: int, torch: Any) -> str:
    prompt = build_direct_prompt(row, expected_n)
    rendered_prompt = processor.apply_chat_template(
        typed_text_messages(prompt),
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    encoded = processor(
        text=rendered_prompt,
        return_tensors="pt",
    )
    device = next(model.parameters()).device
    encoded = {name: tensor.to(device) for name, tensor in encoded.items()}
    input_length = encoded["input_ids"].shape[-1]
    with torch.inference_mode():
        output = model.generate(
            **encoded,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            use_cache=True,
            pad_token_id=getattr(processor.tokenizer, "pad_token_id", None),
            eos_token_id=getattr(processor.tokenizer, "eos_token_id", None),
        )
    generated = output[0, input_length:]
    text = processor.decode(generated, skip_special_tokens=True)
    del encoded, output, generated
    return text


def write_submission(rows: list[dict[str, str]], predictions: dict[str, list[str]]) -> None:
    temp = OUTPUT_CSV.with_suffix(OUTPUT_CSV.suffix + ".tmp")
    output_ids: list[str] = []
    with temp.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=["id", "pred"])
        writer.writeheader()
        for row in rows:
            row_id = row["id"]
            answers = predictions[row_id]
            expected_n = expected_answer_count(row)
            if len(answers) != expected_n or not all(isinstance(item, str) for item in answers):
                raise ValueError(f"invalid prediction shape for id={row_id}")
            writer.writerow({"id": row_id, "pred": serialized_prediction(answers)})
            output_ids.append(row_id)
        handle.flush()
        os.fsync(handle.fileno())
    validate_id_sequence([row["id"] for row in rows], output_ids)
    os.replace(temp, OUTPUT_CSV)


def main() -> None:
    configure_offline_environment()
    started = time.monotonic()
    rows = read_input()
    expected = {row["id"]: expected_answer_count(row) for row in rows}
    predictions = {row["id"]: [""] * expected[row["id"]] for row in rows}

    vendor = bootstrap_local_runtime(ROOT)
    versions = assert_runtime_versions(vendor)
    import torch

    if not torch.cuda.is_available():
        raise RuntimeError("a CUDA GPU is required")
    patches = patch_torch24_for_single_gpu()
    set_reproducible_seed(torch)
    print(
        f"runtime ready: rows={len(rows)} torch={torch.__version__} "
        f"transformers={versions['transformers']} patches={','.join(patches) or 'none'}",
        flush=True,
    )
    model, processor = load_model_and_processor(torch)
    print(
        f"model loaded: gpu={torch.cuda.get_device_name(0)} "
        f"vram_gib={torch.cuda.memory_allocated() / 2**30:.2f}",
        flush=True,
    )

    for index, row in enumerate(rows, 1):
        elapsed = time.monotonic() - started
        if elapsed >= HARD_STOP_SECONDS:
            print(f"hard stop reached after {index - 1}/{len(rows)} rows", flush=True)
            break
        row_id = row["id"]
        row_started = time.monotonic()
        try:
            raw = generate_direct(model, processor, row, expected[row_id], torch)
            predictions[row_id] = parse_answer_lines(raw, expected[row_id])
        except torch.cuda.OutOfMemoryError:
            gc.collect()
            torch.cuda.empty_cache()
            print(f"row={row_id} failed: OutOfMemoryError", flush=True)
        except Exception as exc:
            print(f"row={row_id} failed: {type(exc).__name__}", flush=True)
        write_submission(rows, predictions)
        print(
            f"row={row_id} done {index}/{len(rows)} "
            f"answers={len(predictions[row_id])} seconds={time.monotonic() - row_started:.1f}",
            flush=True,
        )

    write_submission(rows, predictions)
    print(f"wrote {OUTPUT_CSV.name} in {time.monotonic() - started:.1f}s", flush=True)


if __name__ == "__main__":
    main()