242 lines
12 KiB
Python
242 lines
12 KiB
Python
#!/usr/bin/env python
|
|
import os, re, json, time, unicodedata
|
|
from collections import Counter, defaultdict
|
|
|
|
T0 = time.time()
|
|
TIME_LIMIT = float(os.environ.get("IOL_TIME_LIMIT", "1800"))
|
|
SAFETY = float(os.environ.get("IOL_SAFETY", "150"))
|
|
DEADLINE = T0 + TIME_LIMIT - SAFETY
|
|
TEST_CSV = os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv")
|
|
OUT_CSV = os.environ.get("IOL_OUT_CSV", "submission.csv")
|
|
MODEL_ID = os.environ.get("IOL_MODEL", ".")
|
|
WANT_EXPL = os.environ.get("IOL_EXPLAIN", "1") == "1"
|
|
TOK_PER_S = float(os.environ.get("IOL_TOKS", "30"))
|
|
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
|
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
|
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
|
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
|
|
|
|
MAX_NEW = int(os.environ.get("IOL_MAXNEW", "600")) # was 900 — faster passes, answers are short
|
|
MAX_SAMPLES= int(os.environ.get("IOL_MAXSAMPLES", "40")) # was 12/24 — more votes
|
|
BATCH_SIZE = int(os.environ.get("IOL_BATCH", "4")) # REVERT from 8 (this caused 0.1485)
|
|
SAMPLE_TEMP= float(os.environ.get("IOL_TEMP", "0.5")) # keep — value that gave 0.2039
|
|
COT = os.environ.get("IOL_COT", "0") == "1" # keep off
|
|
|
|
def log(m): print(f"[{time.time()-T0:7.1f}s] {m}", flush=True)
|
|
def left(): return DEADLINE - time.time()
|
|
|
|
# ---- item count + source fallback ------------------------------------------
|
|
_LINE_NUM=re.compile(r"^[ \t]*(\d{1,3})[.)\]]",re.M); _PAREN_NUM=re.compile(r"\((\d{1,3})\)")
|
|
_RANGE=re.compile(r"\(?(\d{1,3})\s*(?:[-–—]|to)\s*(\d{1,3})\)?"); _LINE_LET=re.compile(r"^[ \t]*([A-Z])[.)\]]\s",re.M)
|
|
def detect_n_items(query, task_type="", context=""):
|
|
q=query or ""
|
|
line=[int(m) for m in _LINE_NUM.findall(q)]; par=[int(m) for m in _PAREN_NUM.findall(q)]
|
|
rng=0
|
|
for a,b in _RANGE.findall(q):
|
|
a,b=int(a),int(b)
|
|
if 0<b-a<60: rng=max(rng,b-a+1)
|
|
cand=max(len(set(line)),len(set(par)))
|
|
if rng and cand and rng!=cand: return cand
|
|
cand=max(cand,len(set(_LINE_LET.findall(q)))); n=max(rng,cand)
|
|
if n>1: return n
|
|
lines=[l.strip() for l in q.splitlines() if l.strip()]
|
|
if len(lines)>1:
|
|
body=lines[1:] if lines[0].endswith((":",".")) else lines
|
|
if body: return len(body)
|
|
if context:
|
|
cn=len(set(int(m) for m in _LINE_NUM.findall(context)))
|
|
if cn>1: return cn
|
|
cl=len(set(_LINE_LET.findall(context)))
|
|
if cl>1: return cl
|
|
return max(n,1)
|
|
def extract_item_sources(query,n):
|
|
q=query or ""; out=[]
|
|
for ln in q.splitlines():
|
|
s=ln.strip()
|
|
if not s: continue
|
|
m=re.match(r"^\(?(\d{1,3})\)?[.):\]]\s*(.+)$",s)
|
|
if m: out.append(m.group(2).strip())
|
|
if not out:
|
|
lines=[l.strip() for l in q.splitlines() if l.strip()]
|
|
if len(lines)>1 and lines[0].endswith((":",".")): out=lines[1:]
|
|
out=[o.split("|")[0].strip() if "|" in o else o for o in out]; out=[o for o in out if o]
|
|
while len(out)<n: out.append(out[-1] if out else "?")
|
|
return out[:n]
|
|
|
|
# ---- parsing ----------------------------------------------------------------
|
|
_STRIP_PREFIX=re.compile(r"^\s*(?:\(?\d{1,3}\)?[.):\]]\s*|[-*•]\s+)"); _FENCE=re.compile(r"^```[a-zA-Z]*\s*$")
|
|
_CHATTY=re.compile(r"^\s*(?:here (?:are|is)\b|answers?\s*:?\s*$|explanation\b|note\b|okay\b|"
|
|
r"solution\b|reasoning\b|analysis\b|translations?\s*:?\s*$|the answers?\b|let me\b|first,|so,|therefore\b|thus\b)",re.I)
|
|
def clean_line(s):
|
|
s=s.strip(); s=_STRIP_PREFIX.sub("",s); s=s.strip().strip("`").strip()
|
|
if len(s)>=2 and s[0]==s[-1] and s[0] in "\"'“”": s=s[1:-1].strip()
|
|
return s.strip()
|
|
def fit_to_n(items,n,fb=None):
|
|
items=[i for i in items if i and i.strip()]
|
|
if len(items)>n: items=items[-n:]
|
|
while len(items)<n:
|
|
if fb and len(items)<len(fb): items.append(fb[len(items)])
|
|
else: items.append(items[-1] if items else "?")
|
|
return items[:n]
|
|
def raw_lines(text,n,fb=None):
|
|
"""Champion base parse: every non-empty line, forced to N, never blank."""
|
|
lines=[ln.strip() for ln in (text or "").splitlines() if ln.strip()]
|
|
return fit_to_n(lines,n,fb)
|
|
def parse_answers(text,n,fb=None):
|
|
"""Robust parse for CoT / samples: slice after ANSWERS:, else salvage tail."""
|
|
if not text: return list(fb[:n]) if fb else ["?"]*n
|
|
m=None
|
|
for m2 in re.finditer(r"(?:^|\n)\s*(?:final\s+)?answers?\s*:\s*\n?",text,re.I): m=m2
|
|
body=text[m.end():] if m else text
|
|
numbered,raw=[],[]
|
|
for ln in body.splitlines():
|
|
if _FENCE.match(ln): continue
|
|
mm=re.match(r"^\s*\(?(\d{1,3})\)?[.):\]]\s*(.+)$",ln.strip())
|
|
if mm:
|
|
v=clean_line(mm.group(2))
|
|
if v and not _CHATTY.match(v): numbered.append((int(mm.group(1)),v))
|
|
c=clean_line(ln)
|
|
if c and not _CHATTY.match(c): raw.append(c)
|
|
if len(numbered)>=n:
|
|
by={}; [by.__setitem__(l,v) for l,v in numbered]
|
|
labs=sorted(by)
|
|
if len(labs)>=n: return [by[l] for l in labs[:n]]
|
|
return fit_to_n(raw,n,fb)
|
|
def norm(s):
|
|
s=unicodedata.normalize("NFC",(s or "").strip().lower()); s=re.sub(r"\s+"," ",s)
|
|
return s.strip(" .!?;:,")
|
|
def vote(cands,anchor=None):
|
|
cands=[c for c in cands if c and c.strip()]
|
|
if anchor is None: anchor=cands[0] if cands else "?"
|
|
if len(cands)<3: return anchor
|
|
groups=defaultdict(list)
|
|
for c in cands: groups[norm(c)].append(c)
|
|
a_sup=len(groups.get(norm(anchor),[])); bk,bn=None,0
|
|
for k,v in groups.items():
|
|
if len(v)>bn: bk,bn=k,len(v)
|
|
if bk is not None and bn>=2 and bn>a_sup: return Counter(groups[bk]).most_common(1)[0][0]
|
|
return anchor
|
|
|
|
# ---- prompts ----------------------------------------------------------------
|
|
SYS_TRIVIAL=("You solve International Linguistics Olympiad problems. Answer every numbered "
|
|
"item. Put each answer on its own line, in order, with no numbering and no extra text.")
|
|
SYS_COT=("You solve International Linguistics Olympiad problems. Everything needed is in the "
|
|
"problem. Briefly analyse: segment the words, find the recurring morphemes and the rules "
|
|
"that order them, and check them against every given example. Keep the analysis short. "
|
|
"Then write a line containing exactly ANSWERS: and, below it, one answer per item in the "
|
|
"order asked — no numbering, no commentary, no blank lines. Never leave an item blank.")
|
|
|
|
def write_submission(path,ids,preds,expl=None):
|
|
import pandas as pd
|
|
rows=[]
|
|
for i in ids:
|
|
rec={"id":i,"pred":json.dumps(preds[i],ensure_ascii=False)}
|
|
if expl is not None: rec["explanation"]=expl.get(i,"")
|
|
rows.append(rec)
|
|
pd.DataFrame(rows).to_csv(path,index=False)
|
|
|
|
def main():
|
|
import pandas as pd, torch
|
|
from transformers import (AutoTokenizer,AutoModelForCausalLM,StoppingCriteria,StoppingCriteriaList)
|
|
torch.backends.cuda.matmul.allow_tf32=True; torch.backends.cudnn.allow_tf32=True
|
|
log(f"MODE: {'CoT' if COT else 'trivial'} | rep=1.0 | max_samples={MAX_SAMPLES} temp={SAMPLE_TEMP}")
|
|
|
|
df=pd.read_csv(TEST_CSV,dtype=str).fillna("")
|
|
ids=[str(x) for x in df["id"].tolist()]
|
|
ns=[detect_n_items(r.get("query",""),r.get("task_type",""),r.get("context","")) for _,r in df.iterrows()]
|
|
log(f"loaded {len(df)} problems, {sum(ns)} items")
|
|
srcs={i:extract_item_sources(r.get("query",""),n) for i,(_,r),n in zip(ids,df.iterrows(),ns)}
|
|
|
|
preds={i:list(srcs[i]) for i in ids}; expl={i:"" for i in ids} if WANT_EXPL else None
|
|
write_submission(OUT_CSV,ids,preds,expl); log(f"placeholder written ({len(ids)} rows)")
|
|
|
|
class Deadline(StoppingCriteria):
|
|
def __init__(self,t): self.t=t
|
|
def __call__(self,i,s,**k): return time.time()>self.t
|
|
log("loading model ...")
|
|
tok=AutoTokenizer.from_pretrained(MODEL_ID,trust_remote_code=True)
|
|
if tok.pad_token is None: tok.pad_token=tok.eos_token
|
|
tok.padding_side="left"
|
|
def _load(dm):
|
|
try: return AutoModelForCausalLM.from_pretrained(MODEL_ID,torch_dtype=torch.float16,device_map=dm,trust_remote_code=True).eval()
|
|
except TypeError: return AutoModelForCausalLM.from_pretrained(MODEL_ID,dtype=torch.float16,device_map=dm,trust_remote_code=True).eval()
|
|
try: model=_load({"":0} if torch.cuda.is_available() else "auto")
|
|
except Exception as e: log(f"pinned load failed ({e}); auto"); model=_load("auto")
|
|
log(f"model ready ({left():.0f}s left)")
|
|
|
|
sysmsg=SYS_COT if COT else SYS_TRIVIAL
|
|
base_parse=parse_answers if COT else raw_lines # <-- the champion distinction
|
|
prompts=[tok.apply_chat_template(
|
|
[{"role":"system","content":sysmsg},
|
|
{"role":"user","content":f"{r['context'].strip()}\n\n{r['query'].strip()}"}],
|
|
tokenize=False,add_generation_prompt=True) for _,r in df.iterrows()]
|
|
|
|
cb=BATCH_SIZE
|
|
def generate(texts,max_new,sample,temp=0.7):
|
|
nonlocal cb
|
|
out=[""]*len(texts); order=sorted(range(len(texts)),key=lambda i:len(texts[i])); i=0
|
|
while i<len(order):
|
|
if left()<25: break
|
|
idx=order[i:i+cb]; chunk=[texts[j] for j in idx]
|
|
try:
|
|
enc=tok(chunk,return_tensors="pt",padding=True,truncation=True,max_length=6144).to(model.device)
|
|
kw=dict(max_new_tokens=max_new,pad_token_id=tok.pad_token_id,repetition_penalty=1.0,
|
|
stopping_criteria=StoppingCriteriaList([Deadline(DEADLINE-10)]))
|
|
kw.update(dict(do_sample=True,temperature=temp,top_p=0.95) if sample else dict(do_sample=False))
|
|
with torch.no_grad(): o=model.generate(**enc,**kw)
|
|
for k,j in enumerate(idx): out[j]=tok.decode(o[k][enc["input_ids"].shape[1]:],skip_special_tokens=True)
|
|
i+=cb
|
|
except torch.cuda.OutOfMemoryError:
|
|
torch.cuda.empty_cache()
|
|
if cb==1: i+=1
|
|
else: cb=max(1,cb//2)
|
|
except Exception: i+=cb
|
|
return out
|
|
|
|
cap=max(MAX_NEW,1200) if COT else MAX_NEW # reasoning needs headroom
|
|
adaptive=int(0.40*max(1.0,left())*TOK_PER_S/max(1,len(df)))
|
|
max_new=max(192,min(cap,adaptive))
|
|
log(f"Pass 1 greedy, {max_new} tok/item")
|
|
t=time.time(); texts=generate(prompts,max_new=max_new,sample=False); c1=time.time()-t
|
|
samples={i:[] for i in ids}
|
|
for i,n,txt in zip(ids,ns,texts):
|
|
a=base_parse(txt,n,srcs[i]); preds[i]=a; samples[i].append(a)
|
|
write_submission(OUT_CSV,ids,preds,expl); log(f"Pass 1 done in {c1:.0f}s")
|
|
|
|
reserve=min(300.0,0.25*c1+60) if WANT_EXPL else 30.0
|
|
ne=0
|
|
while left()-reserve>c1*1.25 and ne<MAX_SAMPLES:
|
|
ne+=1; log(f"sample pass {ne} ({left():.0f}s left)")
|
|
texts=generate(prompts,max_new=max_new,sample=True,temp=SAMPLE_TEMP)
|
|
for i,n,txt in zip(ids,ns,texts):
|
|
if txt: samples[i].append(parse_answers(txt,n,srcs[i]))
|
|
for i,n in zip(ids,ns):
|
|
if len(samples[i])>=3:
|
|
g=samples[i][0]
|
|
preds[i]=fit_to_n([vote([s[k] for s in samples[i] if k<len(s)],
|
|
anchor=g[k] if k<len(g) else None) for k in range(n)],n,srcs[i])
|
|
write_submission(OUT_CSV,ids,preds,expl); log(f"voted over {ne+1} samples")
|
|
log(f"self-consistency: {ne} sample pass(es) completed")
|
|
|
|
if WANT_EXPL and left()>60:
|
|
log(f"explanations ({left():.0f}s left)")
|
|
ex_sys=("You explain International Linguistics Olympiad solutions to a human judge. State "
|
|
"the key rules of the language: morphemes, word order, sound changes. Concise (2-4 sentences).")
|
|
ep=[tok.apply_chat_template([{"role":"system","content":ex_sys},
|
|
{"role":"user","content":f"{r['context'].strip()}\n\n{r['query'].strip()}\n\nAnswers given:\n"
|
|
+"\n".join(f"- {a}" for a in preds[str(r['id'])])+"\n\nBriefly explain the linguistic rules."}],
|
|
tokenize=False,add_generation_prompt=True) for _,r in df.iterrows()]
|
|
for i,e in zip(ids,generate(ep,max_new=200,sample=False)):
|
|
e=re.sub(r"\s+"," ",(e or "").strip())
|
|
if e: expl[i]=e[:1200]
|
|
write_submission(OUT_CSV,ids,preds,expl); log("explanations written")
|
|
|
|
bad=[i for i,n in zip(ids,ns) if len(preds[i])!=n or any(not str(x).strip() for x in preds[i])]
|
|
if bad:
|
|
for i,n in zip(ids,ns): preds[i]=fit_to_n([x for x in preds[i] if str(x).strip()],n,srcs[i])
|
|
write_submission(OUT_CSV,ids,preds,expl)
|
|
log(f"DONE. {len(ids)} rows, {time.time()-T0:.0f}s elapsed.")
|
|
|
|
if __name__=="__main__":
|
|
main()
|