初始化项目,由ModelHub XC社区提供模型
Model: ShayanShamsi/IOL-AI-v2 Source: Original Platform
This commit is contained in:
425
script.py
Normal file
425
script.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user