# E19: MATH-500 self-consistency baseline (N=4, T=0.7, top_p 0.95, 2048-token cap), Qwen2.5-14B-Instruct Q4_K_M-imat.
# Standard method (Wang et al. 2022), no novelty claim. Needs a llama-server on port 8095 (see run_e19.ps1).
import json, time, requests, sys
from concurrent.futures import ThreadPoolExecutor
from datasets import load_dataset
import math_grade as G
URL="http://127.0.0.1:8095"; OUT="e19_math500_sc4.jsonl"; N=4
ds=load_dataset("HuggingFaceH4/MATH-500",split="test")
def gen(p, s):
    t=time.time(); j=requests.post(URL+"/completion",json={"prompt":p,"n_predict":2048,"cache_prompt":False,"temperature":0.7,"top_p":0.95,"top_k":0,"seed":1000+s},timeout=7200).json()
    tm=j.get("timings",{}); return {"seed":1000+s,"text":j.get("content",""),"predicted_n":tm.get("predicted_n"),"wall":round(time.time()-t,2)}
done=set()
try:
    for l in open(OUT,encoding="utf-8"): done.add(json.loads(l)["item"])
except FileNotFoundError: pass
with open(OUT,"a",encoding="utf-8") as f, ThreadPoolExecutor(N) as ex:
  for k in range(len(ds)):
    if k in done: continue
    e=ds[k]
    p=("<|im_start|>system\nPlease reason step by step, and put your final answer within \\boxed{}.<|im_end|>\n<|im_start|>user\n"+e["problem"]+"<|im_end|>\n<|im_start|>assistant\n")
    t=time.time(); S=list(ex.map(lambda s: gen(p,s), range(N)))
    for s in S: s["boxed"]=G.last_boxed(s["text"]); s["correct"]=bool(G.correct(s["text"], e["answer"]))
    rec={"item":k,"unique_id":e["unique_id"],"level":e["level"],"subject":e["subject"],"gold":e["answer"],"samples":S,"wall":round(time.time()-t,2)}
    f.write(json.dumps(rec)+"\n"); f.flush()
    print(k, sum(s["correct"] for s in S), sum(s["predicted_n"] or 0 for s in S), rec["wall"], flush=True)
print("E19PYDONE", flush=True)