script.py

#!/usr/bin/env python3
"""
IOL-AI 2026 submission: two-phase induction -> application prompting with
self-consistency majority voting.

Robust, transformers-only implementation (NO vLLM, NO runtime install of heavy
libs) — the eval sandbox ships transformers+torch+AWQ support (the organizers'
own Qwen2.5-14B-AWQ baseline runs on it), so we depend only on those.

Guaranteed-output design for the 30-min / 16 GB T4 cap:
  Stage 0 : one fast greedy pass per problem -> write submission.csv immediately
            (a valid, non-zero baseline that survives any later timeout/kill).
  Stage 1 : for each problem, induce the language's rules N times, apply each to
            the query items, majority-vote per item, and OVERWRITE that problem's
            row. Written incrementally, so partial progress is never lost.

Model weights (Qwen3-14B-AWQ) are shipped in this repo and loaded from ".".
"""

import os
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import sys
import time
import json
import re
import csv
import unicodedata
from collections import Counter

START = time.time()

MODEL_ID = os.environ.get("IOL_MODEL_ID", ".")
TEST_CSV = os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv")
OUT_CSV = os.environ.get("IOL_OUT_CSV", "submission.csv")

# Time guards (competition hard cap is 30 min).
STAGE1_DEADLINE_S = float(os.environ.get("IOL_STAGE1_DEADLINE_S", 25 * 60))

N_SAMPLES = int(os.environ.get("IOL_N_SAMPLES", "4"))
IND_MAX_TOKENS = int(os.environ.get("IOL_IND_MAX_TOKENS", "1000"))
APP_MAX_TOKENS = int(os.environ.get("IOL_APP_MAX_TOKENS", "400"))
STAGE0_MAX_TOKENS = int(os.environ.get("IOL_STAGE0_MAX_TOKENS", "400"))
IND_TEMPERATURE = float(os.environ.get("IOL_IND_TEMPERATURE", "0.7"))
BATCH_SEQS = int(os.environ.get("IOL_BATCH_SEQS", "6"))  # max sequences per generate()
ENABLE_THINKING = os.environ.get("IOL_ENABLE_THINKING", "0") == "1"


# ---------------------------------------------------------------------------
# Prompt wording (ported from the paper's chained_prompting_*.py)
# ---------------------------------------------------------------------------

_INDUCTION_TRANSLATION = (
    "Below is a problem sheet from a linguistics exam. Your task is to determine as "
    "much information about the language as possible, purely from the information "
    "provided. You should systematically go through the information provided, and try "
    "to determine the vocabulary meaning of each word (where possible), the syntactic "
    "structure (such as word order), the morphology (including any verb conjugations), "
    "and the meaning of any affixes or subwords. Test every piece of information you "
    "determine against every example provided."
)
_INDUCTION_PATTERN = (
    "Below is a problem sheet from a linguistics exam. Your task is to determine as "
    "much information about the language as possible, purely from the information "
    "provided. You should systematically go through the information provided, and try "
    "to determine the morphological and phonological patterns of the language (such as "
    "noun declension), including the meaning of any subwords or affixes you may see. "
    "Look for systematic patterns in how the language forms syllables, words, and "
    "phrases. Test every piece of information you determine against every example provided."
)
_INDUCTION_NUMBER = (
    "Below is a problem sheet from a linguistics exam. Your task is to determine as "
    "much information about the language and its number system as possible, purely from "
    "the information provided. You should systematically go through the information "
    "provided, and try to determine the vocabulary meaning of each number, the base of "
    "the number system (e.g. decimal, hexadecimal), the syntactic structure (such as "
    "word order), the morphology, and any other patterns you can see in the language's "
    "number system. Test every piece of information you determine against every example "
    "provided."
)
_INDUCTION_MATCHUP = (
    "Below is a problem sheet from a linguistics exam. Your answers to the questions "
    "should rely only on reasoning about the information provided in the sheet. Work out "
    "the correspondences between the items and their meanings, and the vocabulary, "
    "morphology and structure of the language. Test every correspondence you determine "
    "against every example provided."
)
_APPLY_INTRO = {
    "translation": "Based on the information about the language you have determined, solve the following puzzle:",
    "fill_blanks": "Based on the patterns you have identified in the language, solve the following puzzle:",
    "text_to_num": "Based on the information about the language you have determined, solve the following puzzle:",
    "num_to_text": "Based on the information about the language you have determined, solve the following puzzle:",
    "match_letters": "Based on the linguistic patterns and correspondences you have identified, answer the following question:",
}


def induction_text(task_type: str) -> str:
    if task_type == "fill_blanks":
        return _INDUCTION_PATTERN
    if task_type in ("text_to_num", "num_to_text"):
        return _INDUCTION_NUMBER
    if task_type == "match_letters":
        return _INDUCTION_MATCHUP
    return _INDUCTION_TRANSLATION


_ITEM_RE = re.compile(r"(?m)^\s*([0-9]+[.)]|\([0-9]+\)|[0-9]+:)\s*")


def split_query(query: str):
    query = (query or "").strip()
    matches = list(_ITEM_RE.finditer(query))
    if not matches:
        return query, [query] if query else ["?"]
    header = query[: matches[0].start()].strip()
    items = []
    for i, m in enumerate(matches):
        end = matches[i + 1].start() if i + 1 < len(matches) else len(query)
        items.append(query[m.end():end].strip())
    return header, items


def application_body(header, items, task_type):
    intro = _APPLY_INTRO.get(task_type, _APPLY_INTRO["translation"])
    lines = [intro]
    if header:
        lines.append(header)
    for i, it in enumerate(items, 1):
        lines.append(f"{i}. {it}")
    n = len(items)
    note = "For each numbered item give ONLY the letter of its correct match. " if task_type == "match_letters" else ""
    keys = ", ".join(f'"{i}": ""' for i in range(1, n + 1))
    lines.append(
        f"\n{note}Answer every item. Give your answer STRICTLY as a single JSON object "
        f"with one key per item number (as a string), plus a short \"explanation\" key "
        f"summarising the rules you used (2-4 short points, no reasoning trace):\n"
        f"{{{keys}, \"explanation\": \"\"}}"
    )
    return "\n".join(lines)


def stage0_prompt(context, header, items, task_type):
    return f"{context.strip()}\n\n{application_body(header, items, task_type)}"


def induction_prompt(context, task_type):
    return f"{induction_text(task_type)}\n\n{context.strip()}"


def application_prompt(context, task_type, rule, header, items):
    return (
        induction_prompt(context, task_type)
        + "\n\n" + rule.strip()
        + "\n\n" + application_body(header, items, task_type)
    )


# ---------------------------------------------------------------------------
# Parsing + self-consistency vote
# ---------------------------------------------------------------------------

def normalize_answer(answer) -> str:
    if not isinstance(answer, str):
        answer = str(answer)
    answer = unicodedata.normalize("NFC", answer)
    return answer.strip().strip('"').strip("'").rstrip(".").strip().lower()


def _strip_think(text: str) -> str:
    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)


def extract_json(text: str):
    text = _strip_think(text)
    cands = re.findall(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL)
    cands += re.findall(r"\{(?:[^{}]|\{[^{}]*\})*\}", text, re.DOTALL)
    for c in reversed(cands):
        try:
            o = json.loads(c)
            if isinstance(o, dict):
                return o
        except Exception:
            continue
    return {}


def parse_items(text: str, n_items: int):
    obj = extract_json(text)
    explanation = ""
    answers = [""] * n_items
    if obj:
        explanation = str(obj.get("explanation", "") or "")
        for i in range(1, n_items + 1):
            for key in (str(i), i):
                if key in obj and str(obj[key]).strip():
                    answers[i - 1] = str(obj[key]).strip()
                    break
    if not any(answers):
        lines = [ln.strip() for ln in _strip_think(text).splitlines() if ln.strip()]
        lines = [ln for ln in lines if not ln.lower().startswith(("here", "based on", "```"))]
        for i in range(min(n_items, len(lines))):
            answers[i] = re.sub(r"^\s*[0-9]+[.):]\s*", "", lines[i])
    return answers, explanation


def _chrf(a: str, b: str, n: int = 3) -> float:
    a, b = a.lower(), b.lower()
    if not a and not b:
        return 1.0
    if not a or not b:
        return 0.0
    total = 0.0
    for k in range(1, n + 1):
        ag = Counter(a[i:i + k] for i in range(len(a) - k + 1))
        bg = Counter(b[i:i + k] for i in range(len(b) - k + 1))
        if not ag or not bg:
            continue
        inter = sum((ag & bg).values())
        p = inter / max(sum(ag.values()), 1)
        r = inter / max(sum(bg.values()), 1)
        total += 0.0 if (p + r) == 0 else 2 * p * r / (p + r)
    return total / n


def majority_vote(candidates):
    valid = [c for c in candidates if isinstance(c, str) and c.strip()]
    if not valid:
        return ""
    groups = {}
    for c in valid:
        groups.setdefault(normalize_answer(c), []).append(c)
    counts = {k: len(v) for k, v in groups.items()}
    top = max(counts.values())
    winners = [k for k, n in counts.items() if n == top]
    if len(winners) == 1:
        return Counter(groups[winners[0]]).most_common(1)[0][0]
    best, best_score = valid[0], -1.0
    for c in valid:
        s = sum(_chrf(c, o) for o in valid if o is not c)
        if s > best_score:
            best, best_score = c, s
    return best


# ---------------------------------------------------------------------------
# Model + batched generation
# ---------------------------------------------------------------------------

class Model:
    def __init__(self):
        import torch
        from transformers import AutoTokenizer, AutoModelForCausalLM
        self.torch = torch
        self.tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
        if self.tok.pad_token_id is None:
            self.tok.pad_token = self.tok.eos_token
        self.tok.padding_side = "left"
        # Load with .to("cuda") (no `accelerate`/`device_map` dependency) so we
        # rely only on the base transformers+torch stack the sandbox guarantees.
        self.model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID, dtype=torch.float16, trust_remote_code=True,
        ).eval()
        self.model.to("cuda" if torch.cuda.is_available() else "cpu")

    def _render(self, prompt):
        kwargs = {}
        try:
            return self.tok.apply_chat_template(
                [{"role": "user", "content": prompt}],
                tokenize=False, add_generation_prompt=True,
                enable_thinking=ENABLE_THINKING,
            )
        except TypeError:
            return self.tok.apply_chat_template(
                [{"role": "user", "content": prompt}],
                tokenize=False, add_generation_prompt=True,
            )

    def generate(self, prompts, max_new_tokens, do_sample=False, temperature=1.0, n=1):
        """Return list (len=len(prompts)) of lists (len=n) of decoded strings."""
        torch = self.torch
        results = [[] for _ in prompts]
        # Expand: each prompt contributes n sequences; chunk so chunk<=BATCH_SEQS.
        chunk = max(1, BATCH_SEQS // max(n, 1))
        for s in range(0, len(prompts), chunk):
            idxs = list(range(s, min(s + chunk, len(prompts))))
            texts = [self._render(prompts[i]) for i in idxs]
            enc = self.tok(texts, return_tensors="pt", padding=True, truncation=True,
                           max_length=7000).to(self.model.device)
            gen_kwargs = dict(
                max_new_tokens=max_new_tokens,
                num_return_sequences=n,
                pad_token_id=self.tok.pad_token_id,
            )
            if do_sample:
                gen_kwargs.update(do_sample=True, temperature=temperature, top_p=0.9)
            else:
                gen_kwargs.update(do_sample=False)
            with torch.no_grad():
                out = self.model.generate(**enc, **gen_kwargs)
            gen = out[:, enc["input_ids"].shape[1]:]
            dec = self.tok.batch_decode(gen, skip_special_tokens=True)
            # dec is len(idxs)*n, grouped per input.
            for bi, i in enumerate(idxs):
                results[i] = dec[bi * n:(bi + 1) * n]
        return results


# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------

def load_problems():
    rows = []
    with open(TEST_CSV, newline="", encoding="utf-8") as f:
        for r in csv.DictReader(f):
            header, items = split_query(r.get("query", ""))
            rows.append({
                "id": r["id"],
                "task_type": (r.get("task_type") or "").strip(),
                "context": r.get("context", ""),
                "header": header,
                "items": items,
            })
    return rows


def write_out(results, order):
    with open(OUT_CSV, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=["id", "pred", "explanation"])
        w.writeheader()
        for pid in order:
            r = results[pid]
            w.writerow({
                "id": pid,
                "pred": json.dumps(r["pred"], ensure_ascii=False),
                "explanation": r.get("explanation", "") or "",
            })


def run():
    probs = load_problems()
    order = [p["id"] for p in probs]
    print(f"[info] {len(probs)} problems", flush=True)

    # Initialise so a valid file exists even if load crashes mid-way.
    results = {p["id"]: {"pred": [""] * len(p["items"]), "explanation": ""} for p in probs}
    write_out(results, order)

    m = Model()
    print(f"[info] model loaded at {time.time()-START:.0f}s", flush=True)

    # ---- Stage 0: fast greedy baseline for every problem ----
    s0_prompts = [stage0_prompt(p["context"], p["header"], p["items"], p["task_type"]) for p in probs]
    s0 = m.generate(s0_prompts, STAGE0_MAX_TOKENS, do_sample=False, n=1)
    for p, outs in zip(probs, s0):
        ans, expl = parse_items(outs[0], len(p["items"]))
        # never leave blank -> keep best-effort guesses
        results[p["id"]] = {"pred": ans, "explanation": expl}
    write_out(results, order)
    print(f"[info] stage 0 done at {time.time()-START:.0f}s", flush=True)

    # ---- Stage 1: two-phase induction->application + self-consistency ----
    for p in probs:
        if time.time() - START > STAGE1_DEADLINE_S:
            print("[warn] stage1 deadline reached; keeping stage0 for the rest", flush=True)
            break
        try:
            ind = m.generate([induction_prompt(p["context"], p["task_type"])],
                             IND_MAX_TOKENS, do_sample=True,
                             temperature=IND_TEMPERATURE, n=N_SAMPLES)[0]
            app_prompts = [
                application_prompt(p["context"], p["task_type"], _strip_think(rule),
                                   p["header"], p["items"])
                for rule in ind
            ]
            app = m.generate(app_prompts, APP_MAX_TOKENS, do_sample=False, n=1)
            n_items = len(p["items"])
            samples, expl = [], results[p["id"]].get("explanation", "")
            for outs in app:
                a, e = parse_items(outs[0], n_items)
                samples.append(a)
                if e and not expl:
                    expl = e
            final = [majority_vote([s[j] for s in samples if j < len(s) and s[j].strip()])
                     for j in range(n_items)]
            # keep any stage0 answer if voting produced an empty for that item
            s0_ans = results[p["id"]]["pred"]
            final = [f if f.strip() else (s0_ans[j] if j < len(s0_ans) else "")
                     for j, f in enumerate(final)]
            if len(expl) > 600:
                expl = expl[:600].rsplit(" ", 1)[0] + "..."
            results[p["id"]] = {"pred": final, "explanation": expl}
            write_out(results, order)
        except Exception as e:
            print(f"[warn] stage1 failed for {p['id']}: {e}", flush=True)
            continue
    write_out(results, order)
    print(f"[info] done at {time.time()-START:.0f}s", flush=True)


if __name__ == "__main__":
    try:
        run()
    except Exception:
        import traceback
        traceback.print_exc()
        # Ensure a well-formed file exists no matter what.
        try:
            with open(OUT_CSV) as f:
                pass
        except Exception:
            try:
                probs = load_problems()
                with open(OUT_CSV, "w", newline="", encoding="utf-8") as f:
                    w = csv.DictWriter(f, fieldnames=["id", "pred", "explanation"])
                    w.writeheader()
                    for p in probs:
                        w.writerow({"id": p["id"],
                                    "pred": json.dumps([""] * len(p["items"]), ensure_ascii=False),
                                    "explanation": ""})
            except Exception:
                traceback.print_exc()