Upload script.py with huggingface_hub

script.py

@@ -0,0 +1,425 @@
+#!/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()