script.py

"""IOL-AI 2026 submission — v5 hybrid (solver + budget-managed gemma4:12b).

Everything here was validated on a local dev set (see README):
  pass 0  symbolic numeral solver (numeral_solver.py, shipped in this repo);
          answers text_to_num / num_to_text exactly when the system fits the
          searched grammar family, returns None -> LLM fallback otherwise
  pass 1  fast low-token answer for every remaining problem, submission.csv
          atomically checkpointed after every row (a kill never leaves an
          invalid/partial file)
  pass 2  per-row time-sliced re-solve with reasoning; overwrites baseline
          only if the result parses to the right number of answers; translation
          reasoning is deliberately capped LOW (truncated reasoning + a forced
          short answer scored higher than completed reasoning on gemma4:12b)

Repo layout expected: this file as script.py, numeral_solver.py beside it,
gemma4:12b weights in the repo root (load from "."). fp16 12B does not fit a
16 GB T4, so weights load 4-bit via bitsandbytes (the organizer-reference
recipe).
"""
import os
import subprocess
import sys
import time

T0 = time.monotonic()
TIME_LIMIT = 30 * 60
DEADLINE = T0 + TIME_LIMIT - 3 * 60          # 3-min reserve for writes/exit

# Smoke mode (local CPU shakeout, mirrors the leader's IOL_DUMMY pattern):
#   IOL_SMOKE=1        skip bitsandbytes (no CUDA), load fp32 on CPU
#   IOL_MODEL_ID=...   substitute a tiny model for the shipped weights
#   IOL_INPUT=...      read a local CSV instead of /tmp/data/test.csv
SMOKE = os.environ.get("IOL_SMOKE", "0") == "1"

deps = ["transformers>=4.51", "accelerate>=0.30", "torch>=2.2", "pandas"]
if not SMOKE:
    deps.append("bitsandbytes")
subprocess.run([sys.executable, "-m", "pip", "install", "-q", *deps], check=True)

import json
import re

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

import numeral_solver

MODEL_ID = os.environ.get("IOL_MODEL_ID", ".")
INPUT_CSV = os.environ.get("IOL_INPUT", "/tmp/data/test.csv")
MARKER = "FINAL ANSWERS"
TRANSLATION_CAP = 3000        # tokens; see module docstring
GLOBAL_CAP = 8192

BASE_RULES = (
    f"You solve International Linguistics Olympiad problems. Your reasoning budget "
    f"is limited, so be systematic and compact — do not second-guess a hypothesis "
    f"that fits all the data. Verify against every given example once, then commit. "
    f"End with the line '{MARKER}:' followed by one answer per line, in item order, "
    f"with no numbering and no extra text."
)

NUMERAL_CORE = (
    "1) For EVERY example, write one arithmetic equation showing exactly how its "
    "words produce its value. 2) Determine the base from the single-word values. "
    "3) CRITICAL: find pairs of examples using the same words in different orders "
    "with different values — decide which order means multiplication and which "
    "means addition. 4) Only after your equations reproduce ALL examples, "
)

STRATEGIES = {
    "translation": (
        "Method: 1) Align the given sentence pairs and segment every word into "
        "morphemes by comparing entries that share meaning components. 2) Write a "
        "compact table: each root, prefix, and suffix with its meaning, plus the "
        "morpheme order. 3) Compose each requested item from the table. Mind the "
        "direction of translation asked for."
    ),
    "text_to_num": "Method: " + NUMERAL_CORE + "convert each item to digits.",
    "num_to_text": "Method: " + NUMERAL_CORE + "compose each requested number in "
                   "the puzzle language.",
    "fill_blanks": (
        "Method: deduce the paradigm from the completed cells, state the rule for "
        "each row/column, then fill each blank consistently with it."
    ),
    "match_letters": (
        "Method: find anchor items you can pair with certainty first (word length, "
        "repeated letters), then use elimination. Every item gets exactly one match."
    ),
}

SALVAGE = (
    "Time is up. Based on the partial analysis below, give your best answer for "
    "all {n} items RIGHT NOW: one answer per line, in order, no numbering, no "
    "other text, no further reasoning.\n\nThe items to answer:\n{query}\n\n"
    "Partial analysis:\n{tail}"
)


def remaining():
    return DEADLINE - time.monotonic()


tok = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=not SMOKE)
if SMOKE:
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID).eval()
else:
    from transformers import BitsAndBytesConfig
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16),
        device_map="auto", local_files_only=True,
    ).eval()
print(f"[model] loaded at {time.monotonic()-T0:.0f}s", flush=True)


def generate(messages, max_new, max_time=None):
    """Greedy generation. Returns (text, truncated, tokens_per_sec)."""
    enc = tok.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt",
    )
    if hasattr(enc, "keys"):        # BatchEncoding on newer transformers
        inputs = {k: v.to(model.device) for k, v in enc.items()}
    else:                           # bare tensor on older versions
        inputs = {"input_ids": enc.to(model.device)}
    prompt_len = inputs["input_ids"].shape[-1]
    kw = {"max_new_tokens": max_new, "do_sample": False,
          "pad_token_id": tok.eos_token_id}
    if max_time:
        kw["max_time"] = max_time
    t0 = time.monotonic()
    with torch.no_grad():
        out = model.generate(**inputs, **kw)
    n_new = out.shape[-1] - prompt_len
    tps = n_new / max(time.monotonic() - t0, 1e-6)
    text = tok.decode(out[0][prompt_len:], skip_special_tokens=True).strip()
    return text, n_new >= max_new, tps


def item_count(query):
    return len(re.findall(r"^\s*\d+[.)]", query, re.M)) or 1


def extract_answers(text, n):
    m = re.search(rf"{MARKER}\s*:?", text, re.I)
    block = text[m.end():] if m else text
    lines = [ln.strip() for ln in block.splitlines() if ln.strip()]
    if not m:
        lines = lines[-n:]
    lines = [re.sub(r"^\d+[.)]\s*", "", ln) for ln in lines]
    return (lines + [""] * n)[:n]


def valid(answers, n):
    return len(answers) == n and all(answers)


def checkpoint(rows):
    tmp = ".submission.csv.tmp"
    pd.DataFrame(rows, columns=["id", "pred"]).to_csv(tmp, index=False)
    os.replace(tmp, "submission.csv")


def build_messages(r):
    system = BASE_RULES + "\n\n" + STRATEGIES.get(r["task_type"], "")
    return [
        {"role": "system", "content": system},
        {"role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}"},
    ]


df = pd.read_csv(INPUT_CSV, dtype=str).fillna("")
problems = [r for _, r in df.iterrows()]
counts = [item_count(r["query"]) for r in problems]
rows = [{"id": r["id"], "pred": json.dumps([""] * counts[i])}
        for i, r in enumerate(problems)]
checkpoint(rows)

# ── pass 0: symbolic numeral solver ──────────────────────────────────────────
solved = [False] * len(problems)
for i, r in enumerate(problems):
    if r["task_type"] in ("text_to_num", "num_to_text"):
        try:
            ans = numeral_solver.answer(r["context"], r["query"], r["task_type"])
        except Exception as exc:
            print(f"[solver] {r['id']} error: {exc}", flush=True)
            ans = None
        if ans is not None:
            rows[i]["pred"] = json.dumps(ans, ensure_ascii=False)
            solved[i] = True
            print(f"[solver] {r['id']} solved", flush=True)
checkpoint(rows)

# ── pass 1: fast complete baseline ───────────────────────────────────────────
tps_est = 20.0
for i, r in enumerate(problems):
    if solved[i] or remaining() < 60:
        continue
    text, _, tps_est = generate(build_messages(r), max_new=48 * counts[i] + 128,
                                max_time=min(90.0, remaining() / 4))
    rows[i]["pred"] = json.dumps(extract_answers(text, counts[i]),
                                 ensure_ascii=False)
    checkpoint(rows)
    print(f"[base {i+1}/{len(problems)}] done at {time.monotonic()-T0:.0f}s",
          flush=True)

# ── pass 2: time-sliced reasoning upgrades ───────────────────────────────────
for i, r in enumerate(problems):
    if solved[i]:
        continue
    n = counts[i]
    if remaining() < 30:
        print("[deadline] stopping upgrades", flush=True)
        break
    todo = sum(1 for j in range(i, len(problems)) if not solved[j])
    slice_s = max(remaining() / max(todo, 1), 10.0)
    max_new = int(min(max(slice_s * 0.8 * tps_est, 512), GLOBAL_CAP))
    if r["task_type"] == "translation":
        max_new = min(max_new, TRANSLATION_CAP)

    text, truncated, tps_est = generate(build_messages(r), max_new,
                                        max_time=slice_s * 0.8)
    upgraded = extract_answers(text, n) if text else []

    if not valid(upgraded, n) and remaining() > 20:
        msgs = [
            {"role": "system", "content": "Output only the answers, one per line."},
            {"role": "user", "content": SALVAGE.format(
                n=n, query=r["query"].strip(), tail=text[-3000:])},
        ]
        text, _, _ = generate(msgs, max_new=48 * n + 64,
                              max_time=min(60.0, remaining()))
        upgraded = extract_answers(text, n)

    if valid(upgraded, n):
        rows[i]["pred"] = json.dumps(upgraded, ensure_ascii=False)
        checkpoint(rows)
        print(f"[think {i+1}/{len(problems)}] upgraded", flush=True)
    else:
        print(f"[think {i+1}/{len(problems)}] kept baseline", flush=True)

checkpoint(rows)
print(f"wrote submission.csv at {time.monotonic()-T0:.0f}s", flush=True)