Files
iol-solver-qwen3/script.py

449 lines
22 KiB
Python
Raw Normal View History

import os, sys, subprocess, importlib, importlib.metadata, unicodedata
from pathlib import Path
os.environ.setdefault("HF_HUB_OFFLINE","1"); os.environ.setdefault("TRANSFORMERS_OFFLINE","1")
SCRIPT_DIR = Path(__file__).resolve().parent if "__file__" in globals() else Path.cwd()
WHEELHOUSE = SCRIPT_DIR / "wheelhouse"
if not WHEELHOUSE.is_dir(): WHEELHOUSE = Path("wheelhouse")
RUNTIME_DIR = Path("/tmp/qwen3deps")
def emergency(reason):
try:
import pandas as pd, json as j
try: ids = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")["id"].tolist()
except Exception: ids = []
pd.DataFrame([{"id":i,"pred":j.dumps([""]),"explanation":str(reason)[:100]} for i in ids],
columns=["id","pred","explanation"]).to_csv("submission.csv", index=False)
except Exception:
try: open("submission.csv","w").write("id,pred,explanation\n")
except Exception: pass
try:
wheels = [str(WHEELHOUSE / w) for w in os.listdir(WHEELHOUSE) if w.endswith(".whl")]
if not wheels: raise FileNotFoundError(f"no wheels in {WHEELHOUSE}")
RUNTIME_DIR.mkdir(parents=True, exist_ok=True)
subprocess.run([sys.executable,"-m","pip","install","--no-index","--no-deps","--upgrade",
"--target",str(RUNTIME_DIR)] + wheels, check=True, timeout=300)
sys.path.insert(0, str(RUNTIME_DIR)); importlib.invalidate_caches()
try: print("transformers:", importlib.metadata.version("transformers"), flush=True)
except Exception: pass
except Exception as e:
emergency(f"wheel install failed: {e}"); raise
import re, json, time
import pandas as pd, torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import ast as _ast, hashlib as _hash
from fractions import Fraction as _Frac
from collections import OrderedDict as _OD
MODEL_ID="."; TIME_LIMIT=30*60; start=time.time()
def write_csv(rows):
import csv
with open("submission.csv.tmp","w",newline="",encoding="utf-8") as f:
w=csv.DictWriter(f,fieldnames=["id","pred","explanation"]); w.writeheader()
for r in rows: w.writerow(r)
os.replace("submission.csv.tmp","submission.csv")
try:
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
write_csv([{"id":i,"pred":json.dumps([""]),"explanation":"placeholder"} for i in df["id"]])
tok = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float16,
device_map="auto", local_files_only=True).eval()
print("loaded, quantized:", getattr(model.config,"quantization_config",None) is not None, flush=True)
except Exception as e:
emergency(f"load failed: {e}"); raise
SYS=("You solve International Linguistics Olympiad problems about a language you have never seen. "
"Everything you need is in the examples. Answer every numbered item, in order. "
"Put each answer on its own line, with no numbering and no extra text.")
def n_expected(query):
items=re.findall(r"(?m)^\s*(\d+)\s*[.\)]", query)
if items: return len(items)
rng=re.search(r"\(\s*(\d+)\s*[-–—]\s*(\d+)\s*\)", query)
if rng:
lo,hi=int(rng.group(1)),int(rng.group(2))
if 0<hi-lo<100: return hi-lo+1
return None
def norm_generic(s):
s = unicodedata.normalize("NFC", s).strip()
s = re.sub(r"^\s*\(?\d+\)?\s*[.):\-]\s+", "", s)
s = re.sub(r"^\s*[-*•·]\s+", "", s)
s = re.sub(r"(?i)^\s*(answer|translation|output)\s*\d*\s*[:.\-]\s+", "", s)
s = s.strip("* ")
s = re.sub(r"\s{2,}", " ", s)
return s.strip()
def norm_number(s):
g = norm_generic(s)
m = re.search(r"-?\d[\d,\. ]*\d|\d", g)
if not m: return g
digits = re.sub(r"[^\d-]", "", m.group(0))
return digits if digits else g
def normalize(task_type, s):
if task_type == "text_to_num":
return norm_number(s)
return norm_generic(s)
def gen(context, query):
msgs=[{"role":"system","content":SYS},
{"role":"user","content":f"{context.strip()}\n\n{query.strip()}"}]
try:
enc=tok.apply_chat_template(msgs,add_generation_prompt=True,enable_thinking=False,
return_tensors="pt",return_dict=True).to(model.device)
ilen=enc["input_ids"].shape[-1]
with torch.no_grad(): out=model.generate(**enc,max_new_tokens=512,do_sample=False)
except Exception:
ids=tok.apply_chat_template(msgs,add_generation_prompt=True,enable_thinking=False,
return_tensors="pt").to(model.device)
ilen=ids.shape[-1]
with torch.no_grad(): out=model.generate(ids,max_new_tokens=512,do_sample=False)
return tok.decode(out[0][ilen:],skip_special_tokens=True).strip()
rows=[]; done=set()
try:
for _,r in df.iterrows():
try:
task_type = r.get("task_type","")
text=gen(r["context"],r["query"])
ans=[normalize(task_type, ln) for ln in text.splitlines() if ln.strip()]
n=n_expected(r["query"])
if n:
if len(ans)<n: ans=ans+[ans[-1] if ans else ""]*(n-len(ans))
elif len(ans)>n: ans=ans[:n]
if not ans: ans=[""]
expl=re.sub(r"\s+"," ",text[:300]).strip() or "derived from the examples"
rows.append({"id":r["id"],"pred":json.dumps(ans,ensure_ascii=False),"explanation":expl})
except Exception as e:
n=n_expected(r["query"]) or 1
rows.append({"id":r["id"],"pred":json.dumps([""]*n,ensure_ascii=False),"explanation":"fallback"})
print("row error",r["id"],e,flush=True)
done.add(r["id"]); write_csv(rows)
print(f"{len(rows)}/{len(df)} t={time.time()-start:.0f}s",flush=True)
if time.time()-start>TIME_LIMIT-60:
print("time up, stopping",flush=True); break
for _,r in df.iterrows():
if r["id"] in done: continue
n=n_expected(r["query"]) or 1
rows.append({"id":r["id"],"pred":json.dumps([""]*n,ensure_ascii=False),"explanation":"fallback"})
# ==========================================================================
# PASS 2: Grammar induction consensus.
# Appended after the proven 0.121 baseline completes. The submission.csv
# already has valid answers at this point. Pass 2 only IMPROVES rows where
# two independent grammar inductions agree exactly -- never empties them.
# ==========================================================================
_P2_INDUCTION_SYS = (
"You study an International Linguistics Olympiad problem. "
"From the examples only, write a RULE SHEET: with 3-8 bullet points "
"covering the grammar: word meanings, word order, morphology, numeral "
"composition, and exact-form constraints. "
"Do not answer the queries. List only rules verifiable from examples."
)
_P2_APPLICATION_SYS = (
"Apply the rule sheet to the IOL queries. "
"Use only the rule sheet and examples. "
"Follow the exact output format. No alternatives, no explanations."
)
_P2_SOFT_DEADLINE = 1620
_P2_RULE_CAP = 90
_P2_APPLY_CAP = 110
_P2_MIN_T = 25
_P2_SAMPLES = 2
_THINK_RE2 = re.compile(r"</?think\b", re.I)
_FINAL_RE2 = re.compile(r"^FINAL\s+ANSWERS\s*:\s*$", re.I)
_CHKRE2 = re.compile(r"^ARITHMETIC\s+CHECKS\s*:\s*$", re.I)
_BEGIN_RE2 = re.compile(r"^BEGIN\s+(ROW_[1-9]\d*)\s*$", re.I)
_END_RE2 = re.compile(r"^END\s+(ROW_[1-9]\d*)\s*$", re.I)
_JUNK2 = re.compile(r"^(?:note|explanation|reason(?:ing)?|answers?|here\s+(?:are|is))\s*:", re.I)
_FMT2 = re.compile(r"^(?:(?:\d{1,3}[.)]|\(\d{1,3}\))(?:\s+|(?=[^\d]))|[-*•]\s+)")
_SAFE_NUM2 = re.compile(r"^[0-9\s.,;+\-*/^=()\[\]{}×÷·\u2212]+$")
_OPT_RE2 = re.compile(r"(?m)^\s*([A-Za-z])[.)]\s")
def _p2_elapsed():
return time.time() - start
def _p2_field(r, k):
v = r.get(k, ""); return "" if v is None else str(v).strip()
def _p2_classify(r):
d = _p2_field(r,"task_type").lower().replace("-","_")
if d in {"translation","fill_blanks","match_letters","text_to_num","num_to_text"}: return d
q = _p2_field(r,"query").lower()
if "fill" in q and "blank" in q: return "fill_blanks"
if "correspondence" in q or ("match" in q and "letter" in q): return "match_letters"
if re.search(r"\bwrite\s+(?:in|as)\s+digits?\b", q): return "text_to_num"
if re.search(r"\bwrite\s+out\b", q): return "num_to_text"
return "translation"
def _p2_n(r):
return n_expected(_p2_field(r,"query"))
def _p2_opt_labels(r):
if _p2_classify(r) != "match_letters": return None
for src in (_p2_field(r,"query"), _p2_field(r,"context")):
labels = _OPT_RE2.findall(src)
numbered = len(re.findall(r"(?m)^\s*\d+[.)]\s+", src))
if len(labels)>=2 and len(set(labels))==len(labels) and numbered==len(labels):
return tuple(labels)
return None
def _p2_safe_eval(expr):
expr = (expr.replace("×","*").replace("÷","/").replace("·","*")
.replace("\u2212","-").replace("^","**").strip())
if not expr or len(expr)>120: return None
try: tree = _ast.parse(expr, mode="eval")
except (SyntaxError,ValueError): return None
def ev(n):
if isinstance(n,_ast.Expression): return ev(n.body)
if isinstance(n,_ast.Constant) and isinstance(n.value,(int,float)) and not isinstance(n.value,bool):
return _Frac(str(n.value))
if isinstance(n,_ast.UnaryOp) and isinstance(n.op,(_ast.UAdd,_ast.USub)):
v=ev(n.operand); return -v if isinstance(n.op,_ast.USub) else v
if isinstance(n,_ast.BinOp):
l,r=ev(n.left),ev(n.right); op=n.op
if isinstance(op,_ast.Add): return l+r
if isinstance(op,_ast.Sub): return l-r
if isinstance(op,_ast.Mult): return l*r
if isinstance(op,_ast.Div):
if r==0: raise ValueError
return l/r
if isinstance(op,_ast.Pow):
if r.denominator!=1 or not 0<=r.numerator<=10: raise ValueError
return l**r.numerator
raise ValueError
try:
v=ev(tree)
return None if abs(v.numerator)>10**15 else v
except Exception: return None
def _p2_verify_arith(answers, checks):
if len(answers)!=len(checks): return False
for ans,chk in zip(answers,checks):
if chk.count("=")!=1: return False
l,r=chk.split("=",1)
lv,rv=_p2_safe_eval(l),_p2_safe_eval(r)
if lv is None or rv is None or lv!=rv: return False
av=_p2_safe_eval(ans)
if av is None or av!=rv: return False
return True
def _p2_invalid(s):
return bool(s.startswith("```") or _THINK_RE2.search(s) or _FINAL_RE2.match(s)
or _CHKRE2.match(s) or _FMT2.match(s) or _JUNK2.match(s)
or (s.startswith("<") and s.endswith(">"))
or s.casefold() in {"n/a","unknown","?","-"})
def _p2_parse_block(text, n):
if not text or n<=0 or _THINK_RE2.search(text) or "```" in text: return None
lines=[l.strip() for l in text.splitlines()]
marks=[i for i,l in enumerate(lines) if _FINAL_RE2.match(l)]
if len(marks)!=1: return None
m=marks[0]
if any(l for l in lines[:m]): return None
answers=[l for l in lines[m+1:] if l]
if len(answers)!=n or any(_p2_invalid(a) for a in answers): return None
chk_marks=[i for i,l in enumerate(lines[:m]) if _CHKRE2.match(l)]
checks=None
if chk_marks:
if len(chk_marks)!=1 or any(lines[:chk_marks[0]]): return None
checks=[l for l in lines[chk_marks[0]+1:m] if l]
if len(checks)!=n: return None
return (answers, checks)
def _p2_parse_group(text, counts):
if not text: return {}
lines=[l.strip() for l in text.splitlines()]
results={}; seen=set(); dups=set()
for si,line in enumerate(lines):
bm=_BEGIN_RE2.match(line)
if not bm: continue
ordinal=int(bm.group(1).split("_",1)[1])-1
if not 0<=ordinal<len(counts): continue
if ordinal in seen: dups.add(ordinal)
seen.add(ordinal)
handle=f"ROW_{ordinal+1}".casefold()
end_idx=None
for j in range(si+1,len(lines)):
if _BEGIN_RE2.match(lines[j]): break
em=_END_RE2.match(lines[j])
if em and em.group(1).casefold()==handle: end_idx=j; break
if end_idx is None: continue
parsed=_p2_parse_block("\n".join(lines[si+1:end_idx]),counts[ordinal])
if ordinal not in dups and parsed is not None: results[ordinal]=parsed
for d in dups: results.pop(d,None)
return results
def _p2_validate(grow, parsed, counts):
valid={}
for ordinal,(answers,checks) in parsed.items():
if len(answers)!=counts[ordinal]: continue
r=grow[ordinal]; fam=_p2_classify(r)
if fam=="text_to_num":
if not all(any(c.isdigit() for c in a) for a in answers): continue
if not all(_SAFE_NUM2.fullmatch(a) for a in answers): continue
if checks is not None and not _p2_verify_arith(answers,checks): continue
if fam=="match_letters":
labels=_p2_opt_labels(r)
if labels is None: continue
if len(set(answers))!=len(answers): continue
if any(a not in labels for a in answers): continue
valid[ordinal]=answers
mords=[o for o,r in enumerate(grow) if _p2_classify(r)=="match_letters"]
if mords:
contracts=[_p2_opt_labels(grow[o]) for o in mords]
contracts=[c for c in contracts if c is not None]
if contracts and all(c==contracts[0] for c in contracts):
labels=contracts[0]
if sum(counts[o] for o in mords)==len(labels):
if any(o not in valid for o in mords):
for o in mords: valid.pop(o,None)
else:
flat=[a for o in mords for a in valid[o]]
if len(set(flat))!=len(flat) or set(flat)!=set(labels):
for o in mords: valid.pop(o,None)
return valid
def _p2_consensus(row, baseline, n, candidates):
valid=[c for c in candidates if c is not None and len(c)==n]
if len(valid)<2: return baseline
a,b=valid[0],valid[1]
if _p2_classify(row)=="match_letters" or len(baseline)!=n:
return a if a==b else baseline
return [a[i] if a[i]==b[i] else baseline[i] for i in range(n)]
def _p2_greedy(msgs, max_tok, time_limit=None):
try:
enc=tok.apply_chat_template(msgs,add_generation_prompt=True,enable_thinking=False,
return_tensors="pt",return_dict=True).to(model.device)
ilen=enc["input_ids"].shape[-1]
kw={"max_new_tokens":min(max_tok,40960-ilen-1),"do_sample":False}
if time_limit: kw["max_time"]=time_limit
with torch.no_grad(): out=model.generate(**enc,**kw)
except Exception:
ids=tok.apply_chat_template(msgs,add_generation_prompt=True,enable_thinking=False,
return_tensors="pt").to(model.device)
ilen=ids.shape[-1]
kw={"max_new_tokens":min(max_tok,40960-ilen-1),"do_sample":False}
if time_limit: kw["max_time"]=time_limit
with torch.no_grad(): out=model.generate(ids,**kw)
decoded=tok.decode(out[0][ilen:],skip_special_tokens=True).strip()
return decoded if decoded else None
def _p2_sampled(msgs, max_tok, time_limit, seed):
torch.manual_seed(seed)
if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
try:
enc=tok.apply_chat_template(msgs,add_generation_prompt=True,enable_thinking=False,
return_tensors="pt",return_dict=True).to(model.device)
ilen=enc["input_ids"].shape[-1]
with torch.no_grad(): out=model.generate(**enc,
max_new_tokens=min(max_tok,40960-ilen-1),max_time=time_limit,
do_sample=True,temperature=0.7,top_p=0.8,top_k=20)
except Exception:
ids=tok.apply_chat_template(msgs,add_generation_prompt=True,enable_thinking=False,
return_tensors="pt").to(model.device)
ilen=ids.shape[-1]
with torch.no_grad(): out=model.generate(ids,
max_new_tokens=min(max_tok,40960-ilen-1),max_time=time_limit,
do_sample=True,temperature=0.7,top_p=0.8,top_k=20)
decoded=tok.decode(out[0][ilen:],skip_special_tokens=True).strip()
return decoded if decoded else None
def _p2_budget(remaining, cap):
if remaining<=0: return 0.0
avail=_P2_SOFT_DEADLINE-_p2_elapsed()
return max(0.0,min(cap,avail/remaining)) if avail>0 else 0.0
def _p2_seed(grow, idx):
payload=json.dumps([{"c":_p2_field(r,"context")[:200]} for r in grow[:1]],sort_keys=True,separators=(",",":"))
h=_hash.sha256(f"{idx}:{payload}".encode()).digest()
return int.from_bytes(h[:8],"big")%(2**31)
def _p2_guidance(grow):
fams={_p2_classify(r) for r in grow}
parts=[]
if "match_letters" in fams: parts.append("Matching: complete one-to-one correspondence.")
if "text_to_num" in fams: parts.append("Numbers: composition rules, verify arithmetic.")
if "translation" in fams: parts.append("Translation: vocabulary, word order, morphology.")
if "fill_blanks" in fams: parts.append("Blanks: exact morphological transformation.")
if "num_to_text" in fams: parts.append("Numeral construction: base, order, word forms.")
return " ".join(parts)
def _p2_induction_msgs(grow):
ctx=_p2_field(grow[0],"context")
hints=[f"[{i}; {_p2_classify(r)}]\n{_p2_field(r,'query')}" for i,r in enumerate(grow,1)]
return [{"role":"system","content":_P2_INDUCTION_SYS},
{"role":"user","content":(f"Focus: {_p2_guidance(grow)}\n\nEXAMPLES:\n{ctx}\n\n"
f"QUERIES (do not answer):\n"+"\n\n".join(hints))}]
def _p2_output_spec(r):
fam=_p2_classify(r)
if fam=="match_letters": return "One option label per line. Complete bijection."
if fam=="text_to_num": return "Digits only. Optional ARITHMETIC CHECKS: block before FINAL ANSWERS:."
if fam=="fill_blanks": return "One filled form per blank."
if fam=="num_to_text": return "One written numeral per item."
return "One translation per item. Exact surface form."
def _p2_app_msgs(grow, rules, counts):
ctx=_p2_field(grow[0],"context")
blocks=[f"ROW_{i} ({_p2_classify(r)}, {n} answers):\n{_p2_field(r,'query')}\nFormat: {_p2_output_spec(r)}"
for i,(r,n) in enumerate(zip(grow,counts),1)]
shapes=[f"BEGIN ROW_{i}\nFINAL ANSWERS:\n<{n} lines>\nEND ROW_{i}"
for i,n in enumerate(counts,1)]
return [{"role":"system","content":_P2_APPLICATION_SYS},
{"role":"user","content":(f"EXAMPLES:\n{ctx}\n\nRULE SHEET:\n{rules}\n\n"
f"QUERIES:\n"+"\n\n".join(blocks)+"\n\nExact structure:\n"+"\n\n".join(shapes))}]
def _p2_clean_rules(text):
text=(text or "").strip()
if not text or _THINK_RE2.search(text) or "```" in text: return None
if re.search(r"FINAL\s+ANSWERS\s*:",text,re.I): return None
lines=[l.rstrip() for l in text.splitlines()]
if lines and re.match(r"^RULE\s+SHEET\s*:?$",lines[0].strip(),re.I): lines=lines[1:]
text="\n".join(lines).strip()
return text if len(text)>=15 else None
# Group rows by shared context, sort by priority
_p2_grps=_OD()
for _p2_ri,(_,_p2_r) in enumerate(df.iterrows()):
_p2_grps.setdefault(_p2_field(_p2_r,"context"),[]).append(_p2_ri)
def _p2_prio(g):
fams={_p2_classify(df.iloc[i]) for i in g}
if fams&{"match_letters","text_to_num"}: fp=0
elif fams&{"fill_blanks","num_to_text"}: fp=1
else: fp=2
return fp,-len(g),g[0]
_p2_all=sorted(_p2_grps.values(),key=_p2_prio)
_p2_rem=len(_p2_all)*_P2_SAMPLES*2
_p2_avail=_P2_SOFT_DEADLINE-_p2_elapsed()
_p2_max=int(_p2_avail/(_P2_SAMPLES*2*_P2_MIN_T)) if _p2_avail>0 else 0
_p2_planned=_p2_all[:_p2_max]
print(f"pass2: {len(_p2_planned)}/{len(_p2_all)} groups t={_p2_elapsed():.0f}s",flush=True)
for _p2_gn,_p2_group in enumerate(_p2_planned,1):
if _p2_elapsed()>_P2_SOFT_DEADLINE: break
_p2_grow=[df.iloc[i] for i in _p2_group]
_p2_counts=[_p2_n(r) or len(json.loads(rows[i]["pred"])) for i,r in zip(_p2_group,_p2_grow)]
if any(c<=0 for c in _p2_counts): _p2_rem-=_P2_SAMPLES*2; continue
_p2_cands=[[] for _ in _p2_group]
for _p2_s in range(_P2_SAMPLES):
_p2_rt=_p2_budget(_p2_rem,_P2_RULE_CAP); _p2_rem-=1
if _p2_rt<_P2_MIN_T: _p2_rem=0; break
_p2_rules=None
try:
_p2_rules=_p2_clean_rules(_p2_sampled(_p2_induction_msgs(_p2_grow),640,_p2_rt,_p2_seed(_p2_grow,_p2_s)) or "")
except Exception as _e: print(f"rule g{_p2_gn} s{_p2_s}: {_e}",flush=True)
_p2_at=_p2_budget(_p2_rem,_P2_APPLY_CAP); _p2_rem-=1
_p2_val={}
if _p2_rules and _p2_at>=_P2_MIN_T:
try:
_p2_raw=_p2_greedy(_p2_app_msgs(_p2_grow,_p2_rules,_p2_counts),1536,_p2_at)
_p2_val=_p2_validate(_p2_grow,_p2_parse_group(_p2_raw or "",_p2_counts),_p2_counts)
except Exception as _e: print(f"apply g{_p2_gn} s{_p2_s}: {_e}",flush=True)
for _p2_j in range(len(_p2_group)): _p2_cands[_p2_j].append(_p2_val.get(_p2_j))
_p2_chg=0
for _p2_j,_p2_idx in enumerate(_p2_group):
_p2_bl=json.loads(rows[_p2_idx]["pred"])
_p2_fin=_p2_consensus(_p2_grow[_p2_j],_p2_bl,_p2_counts[_p2_j],_p2_cands[_p2_j])
if _p2_fin!=_p2_bl: _p2_chg+=1; rows[_p2_idx]["pred"]=json.dumps(_p2_fin,ensure_ascii=False)
write_csv(rows)
print(f"pass2 g{_p2_gn}/{len(_p2_planned)} chg={_p2_chg} t={_p2_elapsed():.0f}s",flush=True)
if _p2_rem<=0: break
write_csv(rows); print("DONE",flush=True)
except Exception as e:
emergency(f"main loop: {e}"); print("FATAL",e,flush=True)