# E16: GSM8K-128 greedy 14B Q4_K_M-imat (+0.5B draft for speed only), token cap 1024 instead of 320, to measure how much of E13/E15 error is truncation.
# Greedy speculative decoding is output-equivalent to greedy target decoding, so E15 greedy correctness / SC samples are reused offline.
import json, time, requests
from datasets import load_dataset
URL="http://127.0.0.1:8093"; OUT="results/quality/e16_gsm8k_cap1024.jsonl"
ds=load_dataset("openai/gsm8k","main",split="test")
with open(OUT,"w",encoding="utf-8") as f:
  for k in range(128):
    ex=ds[k]
    p=("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n"+ex["question"]+
       "\nSolve step by step, then finish with a line 'The answer is N.'<|im_end|>\n<|im_start|>assistant\n")
    t=time.time(); j=requests.post(URL+"/completion",json={"prompt":p,"n_predict":1024,"cache_prompt":False,"temperature":0,"top_k":1,"seed":0},timeout=3600).json()
    tm=j.get("timings",{})
    rec={"item":k,"text":j.get("content",""),"timings":tm,"wall":round(time.time()-t,2)}
    f.write(json.dumps(rec)+"\n"); f.flush()
    print(k, tm.get("draft_n"), tm.get("draft_n_accepted"), tm.get("predicted_per_second"), flush=True)
print("E16PYDONE", flush=True)