426 lines
17 KiB
Python
426 lines
17 KiB
Python
|
#!/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()
|